update EPSM_path

src/base/camera.cpp

@@ -221,6 +221,7 @@ Camera::Sample Camera::Instance::generate_ray(Expr<uint2> pixel_coord, Expr<floa
     weight *= filter_weight;
     auto c2w = camera_to_world();
     auto o = make_float3(c2w * make_float4(ray->origin(), 1.f));
+    auto pixel_world = make_float3(c2w * make_float4(ray->direction() / ray->direction().z, 1.f));
     auto d = normalize(make_float3x3(c2w) * ray->direction());
     ray->set_origin(o);
     ray->set_direction(d);

src/base/differentiation.cpp

@@ -405,7 +405,8 @@ void Differentiation::add_geom_gradients(Float3 grad_v, Float3 grad_n, Float3 we
     _grad_buffer.value()->atomic(real_grad_offset + triangle.i1 * 8 + 3).fetch_add(grad_n[0] * weight[1]);
     _grad_buffer.value()->atomic(real_grad_offset + triangle.i1 * 8 + 4).fetch_add(grad_n[1] * weight[1]);
     _grad_buffer.value()->atomic(real_grad_offset + triangle.i1 * 8 + 5).fetch_add(grad_n[2] * weight[1]);
-
+
+    //device_log("weight2 {}, grad_v {}, grad_offset {} triangle i2 {}",weight[2], grad_v, real_grad_offset, triangle.i2);
     _grad_buffer.value()->atomic(real_grad_offset + triangle.i2 * 8 + 0).fetch_add(grad_v[0] * weight[2]);
     _grad_buffer.value()->atomic(real_grad_offset + triangle.i2 * 8 + 1).fetch_add(grad_v[1] * weight[2]);
     _grad_buffer.value()->atomic(real_grad_offset + triangle.i2 * 8 + 2).fetch_add(grad_v[2] * weight[2]);
@@ -463,25 +464,40 @@ void Differentiation::register_optimizer(Optimizer::Instance *optimizer) noexcep
 
 void Differentiation::register_geometry_parameter(CommandBuffer &command_buffer, Geometry::MeshData& mesh, Accel& accel, uint instance_id) noexcept {
 
-    LUISA_INFO("register geometry parameter start");
+    //LUISA_INFO("register geometry parameter start");
     auto param_offset = _param_buffer_size;
     auto counter_offset = _counter_size;
     auto grad_offset = _gradient_buffer_size;
     auto param_index = static_cast<uint>(_geometry_params.size());
     //auto channels = 3u;
     //uint buffer_id = mesh.geometry_buffer_id_base;
     auto buffer_view = mesh.vertices;
+
+    auto tri_buffer_view = mesh.triangles;
+    auto triangle_buffer_new = _pipeline.create<Buffer<int>>(tri_buffer_view.size() * 3);
+    auto triangle_buffer_new_view = triangle_buffer_new->view();
+
+    Kernel1D write_tri = [&]() noexcept {
+        auto x = dispatch_x();
+        auto tri = tri_buffer_view->read(x);
+        triangle_buffer_new_view->write(x * 3, (Int)(tri.i0));
+        triangle_buffer_new_view->write(x * 3 + 1, (Int)(tri.i1));
+        triangle_buffer_new_view->write(x * 3 + 2, (Int)(tri.i2));
+    };
+    auto write_tri_shader = pipeline().device().compile(write_tri);
+    command_buffer << write_tri_shader().dispatch(tri_buffer_view.size()) << synchronize();
+
     auto length = buffer_view.size() * 8;
     _counter_size = (_counter_size + length + 8u) & ~0b11u;
     _param_buffer_size = (_param_buffer_size + length + 8u) & ~0b11u;
     _gradient_buffer_size = (_gradient_buffer_size + length + 8u) & ~0b11u;
-    _geometry_params.emplace_back(param_index, instance_id, grad_offset, param_offset, counter_offset, buffer_view, length, 0u, mesh.resource);
+    _geometry_params.emplace_back(param_index, instance_id, grad_offset, param_offset, counter_offset, buffer_view, length, 0u, mesh.resource, triangle_buffer_new_view);
     Kernel1D write_i2o = [&](UInt instance_id, UInt grad_offset) noexcept {
         instance2offset.value()->write(instance_id, grad_offset+1u);
     };
     auto write_i2o_shader = pipeline().device().compile(write_i2o);
     command_buffer << write_i2o_shader(instance_id, grad_offset).dispatch(1u) << synchronize();
-    LUISA_INFO("register geometry parameter finished");
+    //LUISA_INFO("register geometry parameter finished");
 }
 
 void Differentiation::update_parameter_from_external(Stream &stream, luisa::vector<uint> &constants_id, luisa::vector<float4> &constants, luisa::vector<uint> &textures_id, 
@@ -495,33 +511,39 @@ void Differentiation::update_parameter_from_external(Stream &stream, luisa::vect
         stream << image.copy_from(textures[textures_id[i]]);
     }
 
-    LUISA_INFO("working on update paramaters {}",geoms_id.size());
+    //LUISA_INFO("working on update paramaters {}",geoms_id.size());
     // apply geometry parameters
     for (auto i=0;i<geoms_id.size();i++) {
         auto &&p = _geometry_params[geoms_id[i]];
         auto param_offset = p.param_offset();
         auto buffer_view = p.buffer();
         //auto [buffer_view, bindlessbuffer_id] = _pipeline.bindless_arena_buffer<Vertex>(buffer_id);
         auto length = buffer_view.size();
-        LUISA_INFO("here length is {}, size is {}, as<vertex> size is {}",length,geoms[geoms_id[i]].view().size(),geoms[geoms_id[i]].view().as<Vertex>().size());
+        //LUISA_INFO("here length is {}, size is {}, as<vertex> size is {}",length,geoms[geoms_id[i]].view().size(),geoms[geoms_id[i]].view().as<Vertex>().size());
         stream << buffer_view.copy_from(geoms[geoms_id[i]].view().as<Vertex>()) << synchronize() << p.mesh()->build() << synchronize();
         _is_dirty = true;
     }
 }
 
-std::tuple<luisa::vector<void *>, luisa::vector<uint>> Differentiation::get_parameter_from_external(Stream &stream, luisa::vector<uint> &constants_id, luisa::vector<uint> &textures_id, luisa::vector<uint> &geoms_id) noexcept {
+std::tuple<luisa::vector<void *>, luisa::vector<uint>, luisa::vector<void *>, luisa::vector<uint>> Differentiation::get_parameter_from_external(Stream &stream, luisa::vector<uint> &constants_id, luisa::vector<uint> &textures_id, luisa::vector<uint> &geoms_id) noexcept {
     luisa::vector<void*> geom_param{};
     luisa::vector<uint> geom_size{};
+    luisa::vector<void*> tri_param{};
+    luisa::vector<uint> tri_size{};
     // apply geometry parameters
     for (auto i: geoms_id) {
         auto p = _geometry_params[i];
         auto param_offset = p.param_offset();
         auto buffer_view = p.buffer();
+        auto tri_buffer_view = p.tri_buffer();
         auto length = buffer_view.size();
+        auto tri_length = tri_buffer_view.size();
         geom_param.push_back(buffer_view.as<float>().native_handle());
         geom_size.push_back(length*8u);
+        tri_param.push_back(tri_buffer_view.native_handle());
+        tri_size.push_back(tri_length);
     }
-    return std::make_tuple(geom_param, geom_size);
+    return std::make_tuple(geom_param, geom_size, tri_param, tri_size);
 }
 
 std::tuple<luisa::vector<void *>, luisa::vector<void *>> Differentiation::get_gradients(Stream &stream) {
@@ -552,8 +574,8 @@ std::tuple<luisa::vector<void *>, luisa::vector<void *>> Differentiation::get_gr
         auto grad_offset = p.gradient_buffer_offset();
         auto buffer_view = p.buffer();
         auto length = buffer_view.size();
-        LUISA_INFO("here length is {}",length);
-        auto geom_grad_buf_view = _grad_buffer->subview(grad_offset, length);
+        //LUISA_INFO("here length is {}",length);
+        auto geom_grad_buf_view = _grad_buffer->subview(grad_offset, length*8);
         geom_res.push_back(reinterpret_cast<void*>(reinterpret_cast<uint64_t>(geom_grad_buf_view.native_handle())+geom_grad_buf_view.offset_bytes()));
     }
     return std::make_tuple(texture_res, geom_res);

src/base/differentiation.h

@@ -66,16 +66,18 @@ class Differentiation {
         uint _counter_offset;
         float2 _range;
         BufferView<Vertex> _buffer_view;
+        BufferView<int> _tri_buffer_view;
         uint _length;
         uint _buffer_id;
         Mesh *_mesh;
     public:
         GeometryParameter(uint index, uint instance_id, uint grad_offset, uint param_offset,
-                          uint counter_offset, BufferView<Vertex> buffer_view, uint length, uint buffer_id, Mesh *mesh) noexcept
+                          uint counter_offset, BufferView<Vertex> buffer_view, uint length, uint buffer_id, Mesh *mesh, BufferView<int> tri_buffer_view) noexcept
             : _index(index), _instance_id{instance_id}, _grad_offset{grad_offset}, _param_offset{param_offset},
-              _counter_offset{counter_offset}, _buffer_view{buffer_view}, _length(length), _buffer_id(buffer_id), _mesh(mesh) {}
+              _counter_offset{counter_offset}, _buffer_view{buffer_view}, _length(length), _buffer_id(buffer_id), _mesh(mesh), _tri_buffer_view(tri_buffer_view) {}
         [[nodiscard]] auto index() const noexcept { return _index; }
         [[nodiscard]] auto buffer() const noexcept { return _buffer_view; }
+        [[nodiscard]] auto tri_buffer() const noexcept { return _tri_buffer_view; }
         [[nodiscard]] auto buffer_id() const noexcept { return _buffer_id; }
         [[nodiscard]] auto instance_id() const noexcept { return _instance_id; }
         [[nodiscard]] auto gradient_buffer_offset() const noexcept { return _grad_offset; }
@@ -176,7 +178,7 @@ class Differentiation {
     void update_parameter_from_external(Stream &stream, luisa::vector<uint> &constants_id, luisa::vector<float4> &constants, luisa::vector<uint> &textures_id, 
     luisa::vector<Buffer<float4>> &textures, luisa::vector<uint> &geoms_id, luisa::vector<Buffer<float>> &geoms) noexcept;
 
-    std::tuple<luisa::vector<void *>, luisa::vector<uint>> get_parameter_from_external
+    std::tuple<luisa::vector<void *>, luisa::vector<uint>,luisa::vector<void *>, luisa::vector<uint>> get_parameter_from_external
     (Stream &stream, luisa::vector<uint> &constants_id, luisa::vector<uint> &textures_id, luisa::vector<uint> &geoms_id) noexcept;
 
 

src/base/film.h

@@ -43,6 +43,7 @@ class Film : public SceneNode {
         virtual void download(CommandBuffer &command_buffer, float4 *framebuffer) const noexcept = 0;
         virtual bool show(CommandBuffer &command_buffer) const noexcept { return false; }
         virtual void *export_image(CommandBuffer &command_buffer) { return nullptr; }
+        virtual void *export_image_origin(CommandBuffer &command_buffer) { return nullptr; }
         virtual void release() noexcept = 0;
     };
 

src/base/geometry.cpp

@@ -91,7 +91,7 @@ void Geometry::_process_shape(
                 command_buffer << alias_table_buffer_view.copy_from(alias_table.data())
                                << pdf_buffer_view.copy_from(pdf.data());
                 command_buffer << compute::commit();
-                auto geom = MeshGeometry{mesh, vertex_buffer_id, vertex_buffer->view()};
+                auto geom = MeshGeometry{mesh, vertex_buffer_id, vertex_buffer->view(),triangle_buffer->view()};
                 _mesh_cache.emplace(hash, geom);
                 return geom;
             }();
@@ -103,6 +103,7 @@ void Geometry::_process_shape(
             MeshData mesh_data{
                 .resource = mesh_geom.resource,
                 .vertices = mesh_geom.vertices,
+                .triangles = mesh_geom.triangles,
                 .shadow_term = encode_fixed_point(shape->has_vertex_normal() ? shape->shadow_terminator_factor() : 0.f),
                 .intersection_offset = encode_fixed_point(shape->intersection_offset_factor()),
                 .geometry_buffer_id_base = mesh_geom.buffer_id_base,
@@ -197,9 +198,9 @@ bool Geometry::update(CommandBuffer &command_buffer, float time) noexcept {
             //    _accel.set_prim_handle(t.instance_id(), (uint64_t)t.buffer().native_handle());
             //}
             //_pipeline.differentiation()->clear_dirty();
-            LUISA_INFO("start build accel");
+            //LUISA_INFO("start build accel");
             command_buffer << _accel.build() << synchronize();
-            LUISA_INFO("end build accel");
+            //LUISA_INFO("end build accel");
         }
     }
     return updated;

src/base/geometry.h

@@ -47,11 +47,13 @@ class Geometry {
         Mesh *resource;
         uint buffer_id_base;
         BufferView<Vertex> vertices;
+        BufferView<Triangle> triangles;
     };
 
     struct MeshData {
         Mesh *resource;
         BufferView<Vertex> vertices;
+        BufferView<Triangle> triangles;
         uint16_t shadow_term;
         uint16_t intersection_offset;
         uint geometry_buffer_id_base : 22;

src/base/integrator.cpp

@@ -169,17 +169,20 @@ DifferentiableIntegrator::Instance::Instance(
 
 DifferentiableIntegrator::Instance::~Instance() noexcept = default;
 
-void DifferentiableIntegrator::Instance::render_backward(Stream &stream, luisa::vector<Buffer<float>> &grad_in) noexcept {
+luisa::vector<void *> DifferentiableIntegrator::Instance::render_backward(Stream &stream, luisa::vector<Buffer<float>> &grad_in) noexcept {
     CommandBuffer command_buffer{&stream};
     pipeline().differentiation()->clear_gradients(command_buffer);
     LUISA_INFO("Gradients cleared.");
     assert(grad_in.size() == pipeline().camera_count());
+    luisa::vector<void *> result;
     for (auto i = 0u; i < pipeline().camera_count(); i++) {
         auto camera = pipeline().camera(i);
-        //auto pixel_count = resolution.x * resolution.y;
         camera->film()->prepare(command_buffer);
         _render_one_camera_backward(command_buffer, 0,  camera, grad_in[i]);
+        command_buffer << compute::synchronize();
+        result.push_back(camera->film()->export_image(command_buffer));
     }
+    return result;
 }
 
 void DifferentiableIntegrator::Instance::_render_one_camera_backward(

src/base/integrator.h

@@ -44,8 +44,9 @@ class Integrator : public SceneNode {
         [[nodiscard]] auto light_sampler() noexcept { return _light_sampler.get(); }
         [[nodiscard]] auto light_sampler() const noexcept { return _light_sampler.get(); }
         virtual void render(Stream &stream) noexcept = 0;
-        virtual void render_backward(Stream &stream, luisa::vector<Buffer<float>> &grad_in) {
+        virtual luisa::vector<void *> render_backward(Stream &stream, luisa::vector<Buffer<float>> &grad_in) {
             LUISA_INFO("Not implemented!");
+            return luisa::vector<void *>{};
         };
         virtual luisa::vector<void*> render_with_return(Stream &stream) {
             LUISA_INFO("Not implemented!");
@@ -121,7 +122,7 @@ class DifferentiableIntegrator : public ProgressiveIntegrator {
         virtual void _render_one_camera_backward(CommandBuffer &command_buffer, uint iteration, Camera::Instance *camera, Buffer<float> & grad_in) noexcept;
 
     public:
-        void render_backward(Stream &stream, luisa::vector<Buffer<float>> &grad_in) noexcept override;
+        luisa::vector<void *> render_backward(Stream &stream, luisa::vector<Buffer<float>> &grad_in) noexcept override;
         Instance(Pipeline &pipeline, CommandBuffer &command_buffer,
                           const DifferentiableIntegrator *integrator) noexcept;
         ~Instance() noexcept override;

src/base/pipeline.cpp

@@ -148,8 +148,8 @@ void Pipeline::render(Stream &stream) noexcept {
     _integrator->render(stream);
 }
 
-void Pipeline::render_diff(Stream &stream, luisa::vector<Buffer<float>> &grads) noexcept {
-    _integrator->render_backward(stream, grads);
+luisa::vector<void *> Pipeline::render_diff(Stream &stream, luisa::vector<Buffer<float>> &grads) noexcept {
+    return _integrator->render_backward(stream, grads);
 }
 
 luisa::vector<void*> Pipeline::render_with_return(Stream &stream) noexcept {

src/base/pipeline.h

@@ -218,7 +218,7 @@ class Pipeline {
     void update_texture(Stream &stream, uint texture_id, float4 new_value) noexcept;
     void update_mesh(uint mesh_id, uint64_t vertex_buffer) noexcept;
     void render(Stream &stream) noexcept;
-    void render_diff(Stream &stream, luisa::vector<Buffer<float>> &grads) noexcept;
+    luisa::vector<void *> render_diff(Stream &stream, luisa::vector<Buffer<float>> &grads) noexcept;
     luisa::vector<void*> render_with_return(Stream &stream) noexcept;
     [[nodiscard]] auto &printer() noexcept { return *_printer; }
     [[nodiscard]] auto &printer() const noexcept { return *_printer; }

src/films/color.cpp

@@ -73,6 +73,10 @@ class ColorFilmInstance final : public Film::Instance {
         command_buffer << compute::synchronize();
         return _converted.native_handle();
     }
+    void* export_image_origin(CommandBuffer &command_bufferr) noexcept override {
+        _check_prepared();
+        return _image.native_handle();
+    }
     void download(CommandBuffer &command_buffer, float4 *framebuffer) const noexcept override;
     [[nodiscard]] Film::Accumulation read(Expr<uint2> pixel) const noexcept override;
     void release() noexcept override;