From 27ec951d3cfa70b629194edb24f275713db36a85 Mon Sep 17 00:00:00 2001 From: JIankai Xing Date: Wed, 1 May 2024 18:44:27 +0800 Subject: [PATCH] update EPSM_path --- src/base/camera.cpp | 1 + src/base/differentiation.cpp | 42 +- src/base/differentiation.h | 8 +- src/base/film.h | 1 + src/base/geometry.cpp | 7 +- src/base/geometry.h | 2 + src/base/integrator.cpp | 7 +- src/base/integrator.h | 5 +- src/base/pipeline.cpp | 4 +- src/base/pipeline.h | 2 +- src/compute | 2 +- src/films/color.cpp | 4 + src/integrators/bkp2.cpp | 712 ++++++++++++++++++++++ src/integrators/megappm_diff.cpp | 983 ++++++++++++++++--------------- src/python/lrapi.cpp | 35 +- src/shapes/mesh.cpp | 7 +- src/tests/CMakeLists.txt | 3 + src/tests/test_ad.cpp | 181 ++++++ src/tests/test_ad_torch.py | 194 +++--- src/util/frame.cpp | 1 + 20 files changed, 1601 insertions(+), 600 deletions(-) create mode 100644 src/integrators/bkp2.cpp create mode 100644 src/tests/test_ad.cpp diff --git a/src/base/camera.cpp b/src/base/camera.cpp index 50049f85..14c93c8f 100644 --- a/src/base/camera.cpp +++ b/src/base/camera.cpp @@ -221,6 +221,7 @@ Camera::Sample Camera::Instance::generate_ray(Expr pixel_coord, Exprorigin(), 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); diff --git a/src/base/differentiation.cpp b/src/base/differentiation.cpp index aaa35c14..a78e5ec4 100644 --- a/src/base/differentiation.cpp +++ b/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,7 +464,7 @@ 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; @@ -471,17 +472,32 @@ void Differentiation::register_geometry_parameter(CommandBuffer &command_buffer, //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>(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 &constants_id, luisa::vector &constants, luisa::vector &textures_id, @@ -495,7 +511,7 @@ 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(buffer_id); auto length = buffer_view.size(); - LUISA_INFO("here length is {}, size is {}, as size is {}",length,geoms[geoms_id[i]].view().size(),geoms[geoms_id[i]].view().as().size()); + //LUISA_INFO("here length is {}, size is {}, as size is {}",length,geoms[geoms_id[i]].view().size(),geoms[geoms_id[i]].view().as().size()); stream << buffer_view.copy_from(geoms[geoms_id[i]].view().as()) << synchronize() << p.mesh()->build() << synchronize(); _is_dirty = true; } } -std::tuple, luisa::vector> Differentiation::get_parameter_from_external(Stream &stream, luisa::vector &constants_id, luisa::vector &textures_id, luisa::vector &geoms_id) noexcept { +std::tuple, luisa::vector, luisa::vector, luisa::vector> Differentiation::get_parameter_from_external(Stream &stream, luisa::vector &constants_id, luisa::vector &textures_id, luisa::vector &geoms_id) noexcept { luisa::vector geom_param{}; luisa::vector geom_size{}; + luisa::vector tri_param{}; + luisa::vector 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().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> Differentiation::get_gradients(Stream &stream) { @@ -552,8 +574,8 @@ std::tuple, luisa::vector> 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(reinterpret_cast(geom_grad_buf_view.native_handle())+geom_grad_buf_view.offset_bytes())); } return std::make_tuple(texture_res, geom_res); diff --git a/src/base/differentiation.h b/src/base/differentiation.h index 27fe9802..c536f021 100644 --- a/src/base/differentiation.h +++ b/src/base/differentiation.h @@ -66,16 +66,18 @@ class Differentiation { uint _counter_offset; float2 _range; BufferView _buffer_view; + BufferView _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 buffer_view, uint length, uint buffer_id, Mesh *mesh) noexcept + uint counter_offset, BufferView buffer_view, uint length, uint buffer_id, Mesh *mesh, BufferView 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 &constants_id, luisa::vector &constants, luisa::vector &textures_id, luisa::vector> &textures, luisa::vector &geoms_id, luisa::vector> &geoms) noexcept; - std::tuple, luisa::vector> get_parameter_from_external + std::tuple, luisa::vector,luisa::vector, luisa::vector> get_parameter_from_external (Stream &stream, luisa::vector &constants_id, luisa::vector &textures_id, luisa::vector &geoms_id) noexcept; diff --git a/src/base/film.h b/src/base/film.h index 3695d81d..70016010 100644 --- a/src/base/film.h +++ b/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; }; diff --git a/src/base/geometry.cpp b/src/base/geometry.cpp index a44fea6f..870be518 100644 --- a/src/base/geometry.cpp +++ b/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; diff --git a/src/base/geometry.h b/src/base/geometry.h index 80e320ca..a1aafd1c 100644 --- a/src/base/geometry.h +++ b/src/base/geometry.h @@ -47,11 +47,13 @@ class Geometry { Mesh *resource; uint buffer_id_base; BufferView vertices; + BufferView triangles; }; struct MeshData { Mesh *resource; BufferView vertices; + BufferView triangles; uint16_t shadow_term; uint16_t intersection_offset; uint geometry_buffer_id_base : 22; diff --git a/src/base/integrator.cpp b/src/base/integrator.cpp index 43904612..985b17e7 100644 --- a/src/base/integrator.cpp +++ b/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> &grad_in) noexcept { +luisa::vector DifferentiableIntegrator::Instance::render_backward(Stream &stream, luisa::vector> &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 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( diff --git a/src/base/integrator.h b/src/base/integrator.h index 21fc3ea5..33950666 100644 --- a/src/base/integrator.h +++ b/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> &grad_in) { + virtual luisa::vector render_backward(Stream &stream, luisa::vector> &grad_in) { LUISA_INFO("Not implemented!"); + return luisa::vector{}; }; virtual luisa::vector 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 & grad_in) noexcept; public: - void render_backward(Stream &stream, luisa::vector> &grad_in) noexcept override; + luisa::vector render_backward(Stream &stream, luisa::vector> &grad_in) noexcept override; Instance(Pipeline &pipeline, CommandBuffer &command_buffer, const DifferentiableIntegrator *integrator) noexcept; ~Instance() noexcept override; diff --git a/src/base/pipeline.cpp b/src/base/pipeline.cpp index 6f60005e..7044429f 100644 --- a/src/base/pipeline.cpp +++ b/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> &grads) noexcept { - _integrator->render_backward(stream, grads); +luisa::vector Pipeline::render_diff(Stream &stream, luisa::vector> &grads) noexcept { + return _integrator->render_backward(stream, grads); } luisa::vector Pipeline::render_with_return(Stream &stream) noexcept { diff --git a/src/base/pipeline.h b/src/base/pipeline.h index b2d57c17..f77ff5cf 100644 --- a/src/base/pipeline.h +++ b/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> &grads) noexcept; + luisa::vector render_diff(Stream &stream, luisa::vector> &grads) noexcept; luisa::vector render_with_return(Stream &stream) noexcept; [[nodiscard]] auto &printer() noexcept { return *_printer; } [[nodiscard]] auto &printer() const noexcept { return *_printer; } diff --git a/src/compute b/src/compute index f90bc063..f3154848 160000 --- a/src/compute +++ b/src/compute @@ -1 +1 @@ -Subproject commit f90bc063f3883647089ccd87f3e9e712e9b81bac +Subproject commit f31548488c914ecb345154b54742592cf8250ae8 diff --git a/src/films/color.cpp b/src/films/color.cpp index d3556cda..cf1d2671 100644 --- a/src/films/color.cpp +++ b/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 pixel) const noexcept override; void release() noexcept override; diff --git a/src/integrators/bkp2.cpp b/src/integrators/bkp2.cpp new file mode 100644 index 00000000..7a4be0ce --- /dev/null +++ b/src/integrators/bkp2.cpp @@ -0,0 +1,712 @@ +$for(id, path_size-1){ + auto normal_cur_0 = v0->normal(); + auto normal_cur_1 = v1->normal(); + auto normal_cur_2 = v2->normal(); + instance = pipeline().geometry()->instance(inst_ids[id + 1]); + triangle = pipeline().geometry()->triangle(instance, triangle_ids[id + 1]); + v_buffer = instance.vertex_buffer_id(); + v0 = pipeline().buffer(v_buffer).read(triangle.i0); + v1 = pipeline().buffer(v_buffer).read(triangle.i1); + v2 = pipeline().buffer(v_buffer).read(triangle.i2); + auto point_nxt_0 = v0->position(); + auto point_nxt_1 = v1->position(); + auto point_nxt_2 = v2->position(); + + Float3 bary_pre; + + + auto instance = pipeline().geometry()->instance(inst_ids[id]); + auto triangle = pipeline().geometry()->triangle(instance, triangle_ids[id]); + auto v_buffer = instance.vertex_buffer_id(); + + auto v0 = pipeline().buffer(v_buffer).read(triangle.i0); + auto v1 = pipeline().buffer(v_buffer).read(triangle.i1); + auto v2 = pipeline().buffer(v_buffer).read(triangle.i2); + + mat_bary[locate(id*2+0,id*2+0)] = dot(v0-v2, grad_p_cur[0]); + mat_bary[locate(id*2+0,id*2+1)] = dot(v1-v2, grad_p_cur[0]); + mat_bary[locate(id*2+1,id*2+0)] = dot(v0-v2, grad_p_cur[1]); + mat_bary[locate(id*2+1,id*2+1)] = dot(v1-v2, grad_p_cur[1]); + mat_param[id*2+0,1] = grad_p_cur[0]; + mat_param[id*2+1,1] = grad_p_cur[1]; + mat_param[id*2+0,3] = grad_n_cur[0]; + mat_param[id*2+1,3] = grad_n_cur[1]; + mat_param[id*2+0,4] = grad_eta[0]; + mat_param[id*2+1,4] = grad_eta[1]; + + ArrayFloat3<2> grad_uv_pre, grad_uv_cur, grad_uv_nxt; + + $autodiff { + Float3 bary_cur = bary_coords[id]; + Float3 bary_nxt = bary_coords[id+1]; + requires_grad(bary_pre, bary_cur, bary_nxt); + Float3 point_pre = point_pre_0*bary_pre[0]+point_pre_1*bary_pre[1]+point_pre_2*(1-bary_pre[0]-bary_pre[1]); + Float3 point_cur = point_cur_0*bary_cur[0]+point_cur_1*bary_cur[1]+point_cur_2*(1-bary_cur[0]-bary_cur[1]); + Float3 point_nxt = point_nxt_0*bary_nxt[0]+point_nxt_1*bary_nxt[1]+point_nxt_2*(1-bary_nxt[0]-bary_nxt[1]); + Float3 normal_cur = normal_cur_0*bary_cur[0]+normal_cur_1*bary_cur[1]+normal_cur_2*(1-bary_cur[0]-bary_cur[1]); + requires_grad(point_pre, point_cur, point_nxt, normal_cur); + auto wi = normalize(point_pre-point_cur); + auto wo = normalize(point_nxt-point_cur); + auto f = Frame::make(normal_cur); + auto wi_local = f.world_to_local(wi); + auto wo_local = f.world_to_local(wo); + auto res = normalize(wi_local+wo_local*etas[id]); + device_log("res is {}",res); + backward(res[0]); + grad_uv_pre[0] = grad(bary_pre); + grad_uv_cur[0] = grad(bary_cur); + grad_uv_nxt[0] = grad(bary_nxt); + mat_param[((((id+1)<<1)|0)<<2)] = grad(point_pre); + mat_param[((((id+1)<<1)|0)<<2)|1] = grad(point_cur); + mat_param[((((id+1)<<1)|0)<<2)|2] = grad(point_nxt); + mat_param[((((id+1)<<1)|0)<<2)|3] = grad(normal_cur); + }; + auto cust_normalize = [](Float3 x){ + return x/sqrt(x[0]*x[0]+x[1]*x[1]+x[2]*x[2]); + }; + $autodiff { + Float3 bary_cur = bary_coords[id]; + Float3 bary_nxt = bary_coords[id+1]; + requires_grad(bary_pre, bary_cur, bary_nxt); + Float3 point_pre = point_pre_0*bary_pre[0]+point_pre_1*bary_pre[1]+point_pre_2*(1-bary_pre[0]-bary_pre[1]); + Float3 point_cur = point_cur_0*bary_cur[0]+point_cur_1*bary_cur[1]+point_cur_2*(1-bary_cur[0]-bary_cur[1]); + Float3 point_nxt = point_nxt_0*bary_nxt[0]+point_nxt_1*bary_nxt[1]+point_nxt_2*(1-bary_nxt[0]-bary_nxt[1]); + Float3 normal_cur = normal_cur_0*bary_cur[0]+normal_cur_1*bary_cur[1]+normal_cur_2*(1-bary_cur[0]-bary_cur[1]); + device_log("point_pre {} point_cur {} point_nxt {} normal_cur {}",point_pre, point_cur, point_nxt, normal_cur); + requires_grad(point_pre, point_cur, point_nxt, normal_cur); + auto wi = normalize(point_pre-point_cur); + auto wo = normalize(point_nxt-point_cur); + auto s = normalize(make_float3(0.0f, -normal_cur[2], normal_cur[1])); + requires_grad(s); + auto ss = normalize(s - normal_cur * dot(normal_cur, s)); + requires_grad(ss); + auto tt = normalize(cross(normal_cur, ss)); + requires_grad(tt); + Frame f{ss,tt,normal_cur}; + auto wi_local = f.world_to_local(wi); + auto wo_local = f.world_to_local(wo); + backward(wi_local); + auto res = normalize(wi_local+wo_local*etas[id]); + auto grad_s = grad(s); + auto grad_ss = grad(ss); + auto grad_tt = grad(tt); + device_log("grad(s) is {}",grad_s); + device_log("grad(ss) is {}",grad_ss); + device_log("grad(tt) is {}",grad_tt); + device_log("res is {} wi {} wo {} eta {} normal {}",res, wi_local, wo_local, etas[id], normal_cur); + device_log("grad(normal_cur) is {}",grad(normal_cur)); + device_log("grad(point_pre) is {}",grad(point_pre)); + device_log("grad(point_cur) is {}",grad(point_cur)); + device_log("grad(point_nxt) is {}",grad(point_nxt)); + grad_uv_pre[1] = grad(bary_pre); + grad_uv_cur[1] = grad(bary_cur); + grad_uv_nxt[1] = grad(bary_nxt); + mat_param[((((id+1)<<1)|1)<<2)] = grad(point_pre); + mat_param[((((id+1)<<1)|1)<<2)|1] = grad(point_cur); + mat_param[((((id+1)<<1)|1)<<2)|2] = grad(point_nxt); + mat_param[((((id+1)<<1)|1)<<2)|3] = grad(normal_cur); + }; + + for(uint j=0;j<2;j++) + { + $if(id>0) + { + mat_bary[locate(id*2+2+j,id*2-2)] = grad_uv_pre[j][0]; + mat_bary[locate(id*2+2+j,id*2-1)] = grad_uv_pre[j][1]; + }; + mat_bary[locate(id*2+2+j,id*2+0)] = grad_uv_cur[j][0]; + mat_bary[locate(id*2+2+j,id*2+1)] = grad_uv_cur[j][1]; + mat_bary[locate(id*2+2+j,id*2+2)] = grad_uv_nxt[j][0]; + mat_bary[locate(id*2+2+j,id*2+3)] = grad_uv_nxt[j][1]; + } + point_pre_0 = point_cur_0; + point_pre_1 = point_cur_1; + point_pre_2 = point_cur_2; + point_cur_0 = point_nxt_0; + point_cur_1 = point_nxt_1; + point_cur_2 = point_nxt_2; + }; + inverse_matrix(); + compute_and_scatter_grad(); + + // class PhotonMappingLogger{ + // public: + // //Buffer photon_inst_ids, photon_triangle_ids, photon_scatter_events, photon_sizes; + // Buffer path_inst_ids, path_triangle_ids, path_scatter_events, path_end_type, path_sizes, path_light_insts, path_light_tris; + // Buffer path_etas, indirect_betas, path_betas, path_camera_weights, path_light_colors; + // Buffer path_camera_starts, path_light_ends; + + // Buffer path_barys, path_light_barys; + + + + // // Buffer photon_light_starts, photon_colors; + // // Buffer photon_etas; + // //Buffer path_photon_connections; + // //unique_ptr photon_matrix, photon_matrix_param, path_matrix, path_matrix_param; + // //Pipeline _pipeline; + // uint path_per_iter, max_path_depth; + // //uint photon_per_iter, max_photon_depth; + // const Spectrum::Instance *_spectrum; + // const uint _dimension; + // PhotonMappingLogger(uint path_per_iter, uint max_path_depth, const Spectrum::Instance *spectrum) : + // path_per_iter(path_per_iter), max_path_depth(max_path_depth), + // _spectrum(spectrum),_dimension(spectrum->node()->dimension()){ + // auto &&device = spectrum->pipeline().device(); + // //photon_inst_ids = device.create_buffer(photon_per_iter * max_photon_depth); + // //photon_triangle_ids = device.create_buffer(photon_per_iter * max_photon_depth); + // //photon_etas = device.create_buffer(photon_per_iter * max_photon_depth); + // //photon_light_starts = device.create_buffer(photon_per_iter); + // //photon_colors = device.create_buffer(photon_per_iter * max_photon_depth); + // //photon_sizes = device.create_buffer(photon_per_iter); + // //path_photon_connections = device.create_buffer>(path_per_iter * max_photon_per_path); + // path_inst_ids = device.create_buffer(path_per_iter * max_path_depth); + // path_triangle_ids = device.create_buffer(path_per_iter * max_path_depth); + // path_scatter_events = device.create_buffer(path_per_iter * max_path_depth); + // path_etas = device.create_buffer(path_per_iter * max_path_depth); + // path_barys = device.create_buffer(path_per_iter * max_path_depth); + + // path_light_ends = device.create_buffer(path_per_iter * max_path_depth); + // path_light_insts = device.create_buffer(path_per_iter * max_path_depth); + // path_light_tris = device.create_buffer(path_per_iter * max_path_depth); + // path_light_barys = device.create_buffer(path_per_iter * max_path_depth); + // path_end_type = device.create_buffer(path_per_iter * max_path_depth); + + // path_camera_weights = device.create_buffer(path_per_iter); + // path_camera_starts = device.create_buffer(path_per_iter); + + // path_sizes = device.create_buffer(path_per_iter); + // indirect_betas = device.create_buffer(path_per_iter * _dimension); + // path_betas = device.create_buffer(path_per_iter * max_path_depth * _dimension); + // path_light_colors = device.create_buffer(path_per_iter * max_path_depth * _dimension); + // } + + // void add_start_camera(Expr path_id, Var &ray, Float beta) { + // path_camera_starts->write(path_id, ray->origin()); + // path_camera_weights->write(path_id, beta); + // } + + // void add_path_vertex(Expr path_id, Expr path_size, shared_ptr it) { + // auto index = path_id * max_path_depth + path_size; + // path_inst_ids->write(index, it->instance_id()); + // path_triangle_ids->write(index, it->triangle_id()); + // path_barys->write(index, it->bary_coord()); + // } + + // void set_path_vertex_attrib(Expr path_id, Expr path_size, Surface::Sample &s, Float eta) { + // auto index = path_id * max_path_depth + path_size; + // path_etas->write(index, eta); + // path_scatter_events->write(index, s.event); + // for (auto i = 0u; i < _dimension; ++i) + // path_betas->write(index * _dimension + i, s.eval.f[i]); + // } + + // void add_light_end(Expr path_id, Expr path_size, shared_ptr it, SampledSpectrum L) { + // //return; + // auto index = path_id * max_path_depth + path_size; + // path_light_ends->write(index, it->p()); + // path_light_insts->write(index, it->instance_id()); + // path_light_tris->write(index, it->triangle_id()); + // path_light_barys->write(index, it->bary_coord()); + // for (auto i = 0u; i < _dimension; ++i) + // path_light_colors->write(index * _dimension + i, L[i]); + // path_end_type->write(index, 1u|path_end_type->read(index)); + // } + + // void add_envlight_end(Expr path_id, Expr path_size, Float3 end, SampledSpectrum L) { + // //return; + // auto index = path_id * max_path_depth + path_size; + // path_light_ends->write(index, end); + // for (auto i = 0u; i < _dimension; ++i) + // path_light_colors->write(index * _dimension + i, L[i]); + // path_end_type->write(index, 2u|path_end_type->read(index)); + // } + + // /* + // void add_start_light(Expr photon_id, LightSampler::Sample &light_sample) { + // auto index = photon_id; + // auto beta = light_sample.eval.L / light_sample.eval.pdf; + // auto start = light_sample.shadow_ray->origin(); + // photon_light_starts->write(index, start); + // } + // void add_photon_vertex(Expr photon_id, Expr path_size, shared_ptr it, Surface::Sample &s, Float eta) { + // auto index = photon_id * max_photon_depth + path_size; + // photon_inst_ids->write(index, it->instance_id()); + // photon_triangle_ids->write(index, it->triangle_id()); + // photon_etas->write(index, eta); + // photon_colors->write(index, s.eval.f); + // //photon_scatter_events->write(index, s.eval.events); + // } + // void connect_path_photon(UInt path_id, UInt photon_id, UInt path_photon_size, Float3 dis, Float3 Phi){ + // auto index = path_id * max_photon_per_path + path_photon_size; + // path_photon_connections->write(index, photon_id); + // } + + // void add_path_sizes(UInt path_id, UInt path_size){ + // path_sizes->write(path_id, path_size); + // } + + // void add_indirect_end(Expr path_id, Expr path_size, SampledSpectrum beta) { + // for (auto i = 0u; i < _dimension; ++i) + // indirect_betas->write(path_id * _dimension + i, beta[i]); + // //path_end_type->write(index, 4u|path_end_type->read(index)); + // } + + // auto indirect_beta(Expr path_id) { + // SampledSpectrum s{_dimension}; + // for (auto i = 0u; i < _dimension; ++i) + // s[i] = indirect_betas->read(path_id * _dimension + i); + // return s; + // } + // */ + + // }; + + [[nodiscard]] Float3 emit_viewpoint_bp(const Camera::Instance *camera, Expr frame_index, + Expr pixel_id, Expr time, Expr shutter_weight) noexcept { + + sampler()->start(pixel_id, frame_index); + auto u_filter = sampler()->generate_pixel_2d(); + auto u_lens = camera->node()->requires_lens_sampling() ? sampler()->generate_2d() : make_float2(.5f); + auto [camera_ray, _, camera_weight] = camera->generate_ray(pixel_id, time, u_filter, u_lens); + auto spectrum = pipeline().spectrum(); + auto swl = spectrum->sample(spectrum->node()->is_fixed() ? 0.f : sampler()->generate_1d()); + SampledSpectrum beta{swl.dimension(), shutter_weight * camera_weight}; + SampledSpectrum Li{swl.dimension()}; + SampledSpectrum testbeta{swl.dimension()}; + auto ray = camera_ray; + auto pdf_bsdf = def(1e16f); + + auto resolution = camera->film()->node()->resolution(); + auto pixel_id_1d = pixel_id.y*resolution.x+pixel_id.x; + + Float3 start_point = camera_ray->origin(); + Float3 end_point; + Float3 radiance; + ArrayUInt<4> triangle_ids, inst_ids; + ArrayFloat3<4> bary_coords, points, normals; + ArrayFloat<4> etas; + ArrayFloat2<4> grad_barys; + ArrayFloat3<4> grad_betas; + + logger->add_start_camera(pixel_id_1d, ray, camera_weight); + auto path_size = 0u; + $for(depth, node()->max_depth()) { + + // trace + auto wo = -ray->direction(); + auto it = pipeline().geometry()->intersect(ray); + + $if(!it->valid()) { + if (pipeline().environment()) { + auto eval = light_sampler()->evaluate_miss(ray->direction(), swl, time); + Li += beta * eval.L * balance_heuristic(pdf_bsdf, eval.pdf); + logger->add_envlight_end(pixel_id_1d, path_size, ray->direction(), eval.L); + } + $break; + }; + + // hit light + if (!pipeline().lights().empty()) { + $if(it->shape().has_light()) { + auto eval = light_sampler()->evaluate_hit(*it, ray->origin(), swl, time); + Li += beta * eval.L * balance_heuristic(pdf_bsdf, eval.pdf); + logger->add_light_end(pixel_id_1d, path_size, it, eval.L); + }; + } + + $if(!it->shape().has_surface()) { $break; }; + + // generate uniform samples + auto u_light_selection = sampler()->generate_1d(); + auto u_light_surface = sampler()->generate_2d(); + auto u_lobe = sampler()->generate_1d(); + auto u_bsdf = sampler()->generate_2d(); + auto u_rr = def(0.f); + auto rr_depth = node()->rr_depth(); + $if(depth + 1u >= rr_depth) { u_rr = sampler()->generate_1d(); }; + + // sample one light + auto light_sample = light_sampler()->sample( + *it, u_light_selection, u_light_surface, swl, time); + + // trace shadow ray + auto occluded = pipeline().geometry()->intersect_any(light_sample.shadow_ray); + + // evaluate material + auto surface_tag = it->shape().surface_tag(); + auto eta_scale = def(1.f); + Bool stop_direct = false; + auto rr_threshold = node()->rr_threshold(); + auto q = max(beta.max() * eta_scale, .05f); + $if(depth + 1u >= rr_depth) { + $if(q < rr_threshold & u_rr >= q) { stop_direct = true; }; + }; + + PolymorphicCall call; + pipeline().surfaces().dispatch(surface_tag, [&](auto surface) noexcept { + surface->closure(call, *it, swl, wo, 1.f, time); + }); + + call.execute([&](auto closure) noexcept { + // if (auto dispersive = closure->is_dispersive()) { + // $if(*dispersive) { swl.terminate_secondary(); }; + // } + // logger->add_path_vertex(pixel_id_1d, path_size, it); + // direct lighting + + $if(light_sample.eval.pdf > 0.0f & !occluded) { + auto light_pos = light_sample.shadow_ray->origin(); + auto eval = closure->evaluate(wo, wi); + auto w = balance_heuristic(light_sample.eval.pdf, eval.pdf) / + light_sample.eval.pdf; + Li += w * beta * eval.f * light_sample.eval.L; + logger->add_light_end(pixel_id_1d, path_size, it, light_sample.eval.L); + }; + + auto roughness = closure->roughness(); + Bool stop_check = (roughness.x * roughness.y > 0.16f) | stop_direct; + $if(stop_check) { + stop_direct = true; + viewpoints->push(it->p(), swl, beta, wo, pixel_id_1d, surface_tag); + }; + + // sample material + auto surface_sample = closure->sample(wo, u_lobe, u_bsdf); + ray = it->spawn_ray(surface_sample.wi); + pdf_bsdf = surface_sample.eval.pdf; + auto w = ite(surface_sample.eval.pdf > 0.f, 1.f / surface_sample.eval.pdf, 0.f); + beta *= w * surface_sample.eval.f; + + // apply eta scale + auto eta = closure->eta().value_or(1.f); + + + logger->set_path_vertex_attrib(pixel_id_1d, path_size, surface_sample, eta); + path_size += 1; + + $switch(surface_sample.event) { + $case(Surface::event_enter) { eta_scale = sqr(eta); }; + $case(Surface::event_exit) { eta_scale = sqr(1.f / eta); }; + }; + + }); + + beta = zero_if_any_nan(beta); + $if(beta.all([](auto b) noexcept { return b <= 0.f; })) { $break; }; + + $if(stop_direct) { + auto it_next = pipeline().geometry()->intersect(ray); + // logger->add_indirect_end(pixel_id_1d, path_size, beta); + $if(!it_next->valid()) { + if (pipeline().environment()) { + auto eval = light_sampler()->evaluate_miss(ray->direction(), swl, time); + Li += beta * eval.L * balance_heuristic(pdf_bsdf, eval.pdf); + logger->add_envlight_end(pixel_id_1d, path_size, ray->direction(), eval.L); + } + }; + // hit light + if (!pipeline().lights().empty()) { + $if(it_next->shape().has_light()) { + auto eval = light_sampler()->evaluate_hit(*it_next, ray->origin(), swl, time); + Li += beta * eval.L * balance_heuristic(pdf_bsdf, eval.pdf); + logger->add_light_end(pixel_id_1d, path_size, it_next, eval.L); + }; + } + $break; + }; + $if(depth + 1u >= rr_depth) { + beta *= ite(q < rr_threshold, 1.0f / q, 1.f); + }; + }; + ArrayUInt<4> triangle_ids, inst_ids; + ArrayFloat3<4> bary_coords, points, normals; + Float3 light_pos, camera_pos, radiance; + Float3 point_pre = camera_pos; + Float3 grad_rgb; + Float2 grad_xy; + ArrayFloat3<4> grad_barys; + $for(i,path_size) + { + Float3 point_cur = points[i], point_nxt; + Float3 normal_cur = normals[i]; + $if(i call; + pipeline().surfaces().dispatch(surface_tag, [&](auto surface) noexcept { + surface->closure(call, *it, swl, wo, 1.f, time); + }); + Float3 grad_b_cur, grad_b_nxt, grad_b_pre; + call.execute([&](auto closure) noexcept { + Float wi = point_nxt-point_cur; + Float wo = point_pre-point_cur; + $autodiff{ + requires_grad(bary_cur, bary_pre, bary_nxt); + auto eval = closure->evaluate(wo, wi); + auto devalf = L/detach(eval.f)*grad_rgb; + backward(devalf*eval.f); + grad_b_cur = grad(point_cur); + grad_b_nxt = grad(point_nxt); + grad_b_pre = grad(point_pre); + }; + closure->backward(wo, wi, grad_rgb * L / eval.f); + $if(i>0){ + grad_barys[i-1]+=grad_b_pre; + }; + grad_barys[i]+=grad_b_cur; + $if(isrgb(swl, Li); + } + + [[nodiscard]] void photon_tracing_bp(const Camera::Instance *camera, Expr frame_index, + Expr sampler_id, Expr time, Expr photon_id_1d, BufferFloat &grad_in) { + + sampler()->start(sampler_id, frame_index); + // generate uniform samples + auto u_light_selection = sampler()->generate_1d(); + auto u_light_surface = sampler()->generate_2d(); + auto u_direction = sampler()->generate_2d(); + auto spectrum = pipeline().spectrum(); + auto swl = spectrum->sample(spectrum->node()->is_fixed() ? 0.f : sampler()->generate_1d()); + auto light_sample = light_sampler()->sample_le( + u_light_selection, u_light_surface, u_direction, swl, time); + //cos term canceled out in pdf + SampledSpectrum beta = light_sample.eval.L / light_sample.eval.pdf; + + auto ray = light_sample.shadow_ray; + auto pdf_bsdf = def(1e16f); + + //logger->add_start_light(photon_id_1d, light_sample); + UInt path_size = 0u; + + ArrayUInt<4> triangle_ids, inst_ids; + ArrayFloat3<4> bary_coords, points, normals; + ArrayFloat<4> etas; + ArrayFloat2<4> grad_barys; + ArrayFloat3<4> grad_betas; + + + auto resolution = camera->film()->node()->resolution(); + auto max_depth = min(node()->max_depth(), 4u); + auto tot_neighbors = 0u; + + $for(depth, max_depth) { + // trace + //path_size = depth; + auto wi = -ray->direction(); + auto it = pipeline().geometry()->intersect(ray); + // miss + grad_barys[path_size] = make_float2(0.f); + grad_betas[path_size] = make_float3(0.f); + $if(!it->valid()) { + $if(node()->debug_photon()%2==1) + { + $if(photon_id_1d()->debug_photon()) + { + device_log("break it unvalid at size {} id {}", path_size, photon_id_1d); + device_log("ray {} {}", ray->origin(), ray->direction()); + }; + }; + + $break; + }; + + $if(!it->shape().has_surface()) { + $if(node()->debug_photon()%2==1) + { + $if(photon_id_1d()->debug_photon()) + { + device_log("break it non surface {} ", photon_id_1d); + }; + }; + $break; + }; + // generate uniform samples + auto u_lobe = sampler()->generate_1d(); + auto u_bsdf = sampler()->generate_2d(); + auto u_rr = def(0.f); + auto rr_depth = node()->rr_depth(); + $if(depth + 1u >= rr_depth) { u_rr = sampler()->generate_1d(); }; + $if(depth > 0) {// add diffuse constraint? + auto grid = viewpoints->point_to_grid(it->p()); + Float3 grad_beta = make_float3(0.f); + Float2 grad_bary = make_float2(0.f); + auto count_neighbors = 0u; + $for(x, grid.x - 1, grid.x + 2) { + $for(y, grid.y - 1, grid.y + 2) { + $for(z, grid.z - 1, grid.z + 2) { + Int3 check_grid{x, y, z}; + auto viewpoint_index = viewpoints->grid_head(viewpoints->grid_to_index(check_grid)); + $while(viewpoint_index != ~0u) { + auto position = viewpoints->position(viewpoint_index); + auto pixel_id = viewpoints->pixel_id(viewpoint_index); + + auto dis = distance(position, it->p()); + auto rad = indirect->radius(pixel_id); + $if(dis <= rad) { + auto viewpoint_beta = viewpoints->beta(viewpoint_index); + auto viewpoint_wo = viewpoints->wo(viewpoint_index); + auto bary = it->bary_coord(); + auto instance = pipeline().geometry()->instance(it->instance_id()); + auto triangle = pipeline().geometry()->triangle(instance, it->triangle_id()); + auto v_buffer = instance.vertex_buffer_id(); + auto v0 = pipeline().buffer(v_buffer).read(triangle.i0); + auto v1 = pipeline().buffer(v_buffer).read(triangle.i1); + auto v2 = pipeline().buffer(v_buffer).read(triangle.i2); + auto point_0 = v0->position(); + auto point_1 = v1->position(); + auto point_2 = v2->position(); + auto surface_tag = viewpoints->surface_tag(viewpoint_index); + SampledSpectrum eval_viewpoint(3u); + PolymorphicCall call; + pipeline().surfaces().dispatch(surface_tag, [&](auto surface) noexcept { + surface->closure(call, *it, swl, viewpoint_wo, 1.f, time); + }); + call.execute([&](const Surface::Closure *closure) noexcept { + eval_viewpoint = closure->evaluate(viewpoint_wo, wi).f; + }); + Float2 grad_b{0.0f, 0.0f}; + Float grad_pixel_0, grad_pixel_1, grad_pixel_2; + Float rel_dis_diff, grad_dis; + Float weight; + Float3 Phi, Phi_real; + $autodiff { + Float3 beta_diff = make_float3(beta[0u], beta[1u], beta[2u]); + requires_grad(bary, beta_diff); + Float3 photon_pos = point_0 * bary[0] + point_1 * bary[1] + point_2 * (1 - bary[0] - bary[1]); + rel_dis_diff = distance(position, photon_pos) / rad; + requires_grad(rel_dis_diff); + auto rel3 = rel_dis_diff*rel_dis_diff*rel_dis_diff; + weight = 3.5f*(1- 6*rel3*rel_dis_diff*rel_dis_diff + 15*rel3*rel_dis_diff - 10*rel3); + auto wi_local = it->shading().world_to_local(wi); + Phi = spectrum->srgb(swl, viewpoint_beta * eval_viewpoint / abs_cos_theta(wi_local)); + Phi_real = spectrum->srgb(swl, beta * viewpoint_beta * eval_viewpoint / abs_cos_theta(wi_local)); + auto Phi_beta = Phi * beta_diff * weight / 200000.0f; + auto _grad_dimension = 3u; + grad_pixel_0 = grad_in->read(pixel_id * _grad_dimension + 0); + grad_pixel_1 = grad_in->read(pixel_id * _grad_dimension + 1); + grad_pixel_2 = grad_in->read(pixel_id * _grad_dimension + 2); + auto dldPhi = (Phi_beta[0]*grad_pixel_0 + Phi_beta[1]*grad_pixel_1 + Phi_beta[2]*grad_pixel_2); + // auto dist = photon_pos[0]*photon_pos[0]+photon_pos[1]*photon_pos[1]+photon_pos[2]*photon_pos[2]; + // backward(dist); + backward(dldPhi); + grad_b = grad(bary).xy(); + grad_bary += grad_b; + grad_beta += grad(beta_diff); + grad_dis = grad(rel_dis_diff); + }; + count_neighbors+=1; + + $if((node()->debug_mode() / 2) % 2 == 1) { + $if(photon_id_1d()->debug_photon()) + { + $if(etas[0]==1.8f) + { + UInt2 pixel_id_2d = make_uint2(pixel_id % resolution.x, pixel_id / resolution.x); + camera->film()->accumulate(pixel_id_2d, make_float3(grad_dis,0.0f,0.0f)); + }; + }; + }; + + $if((node()->debug_mode() / 16) % 2 == 1) { + $if(photon_id_1d()->debug_photon()) + { + indirect->add_phi(pixel_id, Phi_real*weight); + indirect->add_cur_n(pixel_id, 1u); + }; + }; + }; + viewpoint_index = viewpoints->nxt(viewpoint_index); + }; + }; + }; + }; + $if(count_neighbors>0){ + tot_neighbors+=count_neighbors; + grad_betas[path_size] = make_float3(grad_beta[0]/count_neighbors, grad_beta[1]/count_neighbors, grad_beta[2]/count_neighbors); + grad_barys[path_size] = make_float2(grad_bary[0]/count_neighbors, grad_bary[1]/count_neighbors); + }; + }; + auto surface_tag = it->shape().surface_tag(); + auto eta_scale = def(1.f); + PolymorphicCall call; + pipeline().surfaces().dispatch(surface_tag, [&](auto surface) noexcept { + surface->closure(call, *it, swl, wi, 1.f, time); + }); + call.execute([&](auto closure) noexcept { + // apply opacity map + if (auto dispersive = closure->is_dispersive()) { + $if(*dispersive) { swl.terminate_secondary(); }; + } + // sample material + auto surface_sample = closure->sample(wi, u_lobe, u_bsdf, TransportMode::IMPORTANCE); + ray = it->spawn_ray(surface_sample.wi); + pdf_bsdf = surface_sample.eval.pdf; + auto w = ite(surface_sample.eval.pdf > 0.f, 1.f / surface_sample.eval.pdf, 0.f); + auto bnew = beta * w * surface_sample.eval.f; + // apply eta scale + auto eta = closure->eta().value_or(1.f); + auto roughness = closure->roughness(); + $switch(surface_sample.event) { + $case(Surface::event_enter) { eta_scale = sqr(eta); etas[path_size] = eta;}; + $case(Surface::event_exit) { eta_scale = sqr(1.f / eta); etas[path_size] = 1.f / eta;}; + }; + eta_scale *= ite(beta.max() < bnew.max(), 1.f, bnew.max() / beta.max()); + beta = bnew; + }); + beta = zero_if_any_nan(beta); + $if(beta.all([](auto b) noexcept { return b <= 0.f; })) { + // $if(photon_id_1d()->debug_photon()) + // { + // device_log("break beta negative {}",photon_id_1d); + // }; + $break; + }; + auto rr_threshold = node()->rr_threshold(); + auto q = max(eta_scale, .05f); + $if(depth + 1u >= rr_depth) { + $if(q < rr_threshold & u_rr >= q) { + $if(photon_id_1d()->debug_photon()) + { + device_log("break rr {}",photon_id_1d); + }; + $break; + }; + beta *= ite(q < rr_threshold, 1.0f / q, 1.f); + }; + triangle_ids[path_size] = it->triangle_id(); + inst_ids[path_size] = it->instance_id(); + bary_coords[path_size] = it->bary_coord(); + points[path_size] = it->p(); + normals[path_size] = it->ng(); + path_size = path_size+1u; + }; + //@Todo: log surface event + $if(photon_id_1d()->debug_photon()) + { + $if(path_size==2) { + $if(etas[0]==1.8f) + { + $if((node()->debug_mode() / 1) % 2 == 1) { + device_log("path size is {}",path_size); + device_log("photon_id is {}",photon_id_1d); + device_log("light sample is {}",light_sample.shadow_ray->origin()); + $for(i,path_size) + { + device_log("{} point {} normal {} eta {} inst {}, bary {} grad_bary {}",i, points[i], normals[i], etas[i], inst_ids[i], bary_coords[i], grad_barys[i]); + }; + }; + EPSM_photon(path_size, points, normals, inst_ids, triangle_ids, bary_coords, etas, light_sample.shadow_ray->origin(), grad_barys, photon_id_1d); + }; + }; + }; + } \ No newline at end of file diff --git a/src/integrators/megappm_diff.cpp b/src/integrators/megappm_diff.cpp index 909b13de..31af30f3 100644 --- a/src/integrators/megappm_diff.cpp +++ b/src/integrators/megappm_diff.cpp @@ -7,6 +7,10 @@ #include #include #include +#include + +#include +#include namespace luisa::render { @@ -31,7 +35,8 @@ class MegakernelPhotonMappingDiff: public DifferentiableIntegrator { bool _separate_direct; bool _shared_radius; int _debug_photon_id; - bool _debug; + uint _backward_iter; + int _debug_mode; public: MegakernelPhotonMappingDiff(Scene *scene, const SceneNodeDesc *desc) noexcept @@ -42,7 +47,8 @@ class MegakernelPhotonMappingDiff: public DifferentiableIntegrator { _initial_radius{std::max(desc->property_float_or_default("initial_radius", -200.f), -10000.f)},//<0 for world_size/-radius (-grid count) _photon_per_iter{std::max(desc->property_uint_or_default("photon_per_iter", 200000u), 10u)}, _debug_photon_id{desc->property_int_or_default("debug_photon_id", -1)}, - _debug{desc->property_bool_or_default("debug", false)}, + _debug_mode{desc->property_int_or_default("debug", 0)}, + _backward_iter{desc->property_uint_or_default("backward_iter", 1u)}, _separate_direct{true}, //when false, use photon mapping for all flux and gathering at first intersection. Just for debug _shared_radius{true} {};//whether or not use the shared radius trick in SPPM paper. True is better in performance. [[nodiscard]] auto max_depth() const noexcept { return _max_depth; } @@ -53,8 +59,9 @@ class MegakernelPhotonMappingDiff: public DifferentiableIntegrator { [[nodiscard]] bool is_differentiable() const noexcept override { return true; } [[nodiscard]] auto separate_direct() const noexcept { return _separate_direct; } [[nodiscard]] auto shared_radius() const noexcept { return _shared_radius; } - [[nodiscard]] auto debug() const noexcept { return _debug; } + [[nodiscard]] auto debug_mode() const noexcept { return _debug_mode; } [[nodiscard]] auto debug_photon() const noexcept { return _debug_photon_id; } + [[nodiscard]] auto backward_iter() const noexcept { return _backward_iter; } [[nodiscard]] luisa::string_view impl_type() const noexcept override { return LUISA_RENDER_PLUGIN_NAME; } [[nodiscard]] luisa::unique_ptr build(Pipeline &pipeline, CommandBuffer &command_buffer) const noexcept override; }; @@ -215,6 +222,7 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins }; //Store the information of pixel updates class PixelIndirect { + public: Buffer _radius; Buffer _cur_n; Buffer _cur_w; @@ -442,143 +450,12 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins }; } }; - - class PhotonMappingLogger{ - public: - //Buffer photon_inst_ids, photon_triangle_ids, photon_scatter_events, photon_sizes; - Buffer path_inst_ids, path_triangle_ids, path_scatter_events, path_end_type, path_sizes, path_light_insts, path_light_tris; - Buffer path_etas, indirect_betas, path_betas, path_camera_weights,path_light_colors; - Buffer path_camera_starts, path_light_ends; - - Buffer path_barys, path_light_barys; - - - - // Buffer photon_light_starts, photon_colors; - // Buffer photon_etas; - //Buffer path_photon_connections; - //unique_ptr photon_matrix, photon_matrix_param, path_matrix, path_matrix_param; - //Pipeline _pipeline; - uint path_per_iter, max_path_depth; - //uint photon_per_iter, max_photon_depth; - const Spectrum::Instance *_spectrum; - const uint _dimension; - PhotonMappingLogger(uint path_per_iter, uint max_path_depth, const Spectrum::Instance *spectrum) : - path_per_iter(path_per_iter), max_path_depth(max_path_depth), - _spectrum(spectrum),_dimension(spectrum->node()->dimension()){ - auto &&device = spectrum->pipeline().device(); - //photon_inst_ids = device.create_buffer(photon_per_iter * max_photon_depth); - //photon_triangle_ids = device.create_buffer(photon_per_iter * max_photon_depth); - //photon_etas = device.create_buffer(photon_per_iter * max_photon_depth); - //photon_light_starts = device.create_buffer(photon_per_iter); - //photon_colors = device.create_buffer(photon_per_iter * max_photon_depth); - //photon_sizes = device.create_buffer(photon_per_iter); - //path_photon_connections = device.create_buffer>(path_per_iter * max_photon_per_path); - path_inst_ids = device.create_buffer(path_per_iter * max_path_depth); - path_triangle_ids = device.create_buffer(path_per_iter * max_path_depth); - path_scatter_events = device.create_buffer(path_per_iter * max_path_depth); - path_etas = device.create_buffer(path_per_iter * max_path_depth); - path_barys = device.create_buffer(path_per_iter * max_path_depth); - - path_light_ends = device.create_buffer(path_per_iter * max_path_depth); - path_light_insts = device.create_buffer(path_per_iter * max_path_depth); - path_light_tris = device.create_buffer(path_per_iter * max_path_depth); - path_light_barys = device.create_buffer(path_per_iter * max_path_depth); - path_end_type = device.create_buffer(path_per_iter * max_path_depth); - - path_camera_weights = device.create_buffer(path_per_iter); - path_camera_starts = device.create_buffer(path_per_iter); - - path_sizes = device.create_buffer(path_per_iter); - indirect_betas = device.create_buffer(path_per_iter * _dimension); - path_betas = device.create_buffer(path_per_iter * max_path_depth * _dimension); - path_light_colors = device.create_buffer(path_per_iter * max_path_depth * _dimension); - } - - void add_start_camera(Expr path_id, Var &ray, Float beta) { - path_camera_starts->write(path_id, ray->origin()); - path_camera_weights->write(path_id, beta); - } - - void add_path_vertex(Expr path_id, Expr path_size, shared_ptr it) { - auto index = path_id * max_path_depth + path_size; - path_inst_ids->write(index, it->instance_id()); - path_triangle_ids->write(index, it->triangle_id()); - path_barys->write(index, it->bary_coord()); - } - - void set_path_vertex_attrib(Expr path_id, Expr path_size, Surface::Sample &s, Float eta) { - auto index = path_id * max_path_depth + path_size; - path_etas->write(index, eta); - path_scatter_events->write(index, s.event); - for (auto i = 0u; i < _dimension; ++i) - path_betas->write(index * _dimension + i, s.eval.f[i]); - } - void add_light_end(Expr path_id, Expr path_size, shared_ptr it, SampledSpectrum L) { - //return; - auto index = path_id * max_path_depth + path_size; - path_light_ends->write(index, it->p()); - path_light_insts->write(index, it->instance_id()); - path_light_tris->write(index, it->triangle_id()); - path_light_barys->write(index, it->bary_coord()); - for (auto i = 0u; i < _dimension; ++i) - path_light_colors->write(index * _dimension + i, L[i]); - path_end_type->write(index, 1u|path_end_type->read(index)); - } - - void add_envlight_end(Expr path_id, Expr path_size, Float3 end, SampledSpectrum L) { - //return; - auto index = path_id * max_path_depth + path_size; - path_light_ends->write(index, end); - for (auto i = 0u; i < _dimension; ++i) - path_light_colors->write(index * _dimension + i, L[i]); - path_end_type->write(index, 2u|path_end_type->read(index)); - } - - /* - void add_start_light(Expr photon_id, LightSampler::Sample &light_sample) { - auto index = photon_id; - auto beta = light_sample.eval.L / light_sample.eval.pdf; - auto start = light_sample.shadow_ray->origin(); - photon_light_starts->write(index, start); - } - void add_photon_vertex(Expr photon_id, Expr path_size, shared_ptr it, Surface::Sample &s, Float eta) { - auto index = photon_id * max_photon_depth + path_size; - photon_inst_ids->write(index, it->instance_id()); - photon_triangle_ids->write(index, it->triangle_id()); - photon_etas->write(index, eta); - photon_colors->write(index, s.eval.f); - //photon_scatter_events->write(index, s.eval.events); - } - void connect_path_photon(UInt path_id, UInt photon_id, UInt path_photon_size, Float3 dis, Float3 Phi){ - auto index = path_id * max_photon_per_path + path_photon_size; - path_photon_connections->write(index, photon_id); - } - - void add_path_sizes(UInt path_id, UInt path_size){ - path_sizes->write(path_id, path_size); - } - - void add_indirect_end(Expr path_id, Expr path_size, SampledSpectrum beta) { - for (auto i = 0u; i < _dimension; ++i) - indirect_betas->write(path_id * _dimension + i, beta[i]); - //path_end_type->write(index, 4u|path_end_type->read(index)); - } - - auto indirect_beta(Expr path_id) { - SampledSpectrum s{_dimension}; - for (auto i = 0u; i < _dimension; ++i) - s[i] = indirect_betas->read(path_id * _dimension + i); - return s; - } - */ - - }; - - luisa::unique_ptr logger; luisa::unique_ptr viewpoints; luisa::unique_ptr indirect; + const UInt max_size = 4u; + const UInt param_size_per_vert = 5u; + const UInt _grad_dimension = 5u; protected: void _render_one_camera_backward(CommandBuffer &command_buffer, uint iteration, Camera::Instance *camera, Buffer &grad_in) noexcept { @@ -615,12 +492,9 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins auto viewpoints_per_iter = resolution.x * resolution.y; - logger = make_unique(pixel_count, node()->max_depth(), spectrum); indirect = make_unique(viewpoints_per_iter, spectrum, camera->film(), clamp, node()->shared_radius()); viewpoints = make_unique(viewpoints_per_iter, spectrum); - //return; - //pathlogger = make_unique(node()->max_depth(), node()->photon_per_iter(), spectrum); - //initialize PixelIndirect + Kernel1D indirect_initialize_kernel = [&]() noexcept { auto index = dispatch_x(); @@ -645,16 +519,19 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins viewpoints->reset(index); }; - Kernel2D viewpath_construct_kernel = [&](UInt frame_index, Float time, Float shutter_weight) noexcept { + Kernel2D viewpath_construct_kernel = [&](UInt frame_index, Float time, Float shutter_weight, BufferFloat grad_in) noexcept { //Construct view path auto pixel_id = dispatch_id().xy(); - auto L = emit_viewpoint_bp(camera, frame_index, pixel_id, time, shutter_weight); - camera->film()->accumulate(pixel_id, L, 0.5f); + auto L = emit_viewpoint_bp(camera, frame_index, pixel_id, time, shutter_weight, grad_in); + $if((node()->debug_mode() / 2) % 2 == 0 ) { + $if((node()->debug_mode() / 16) % 2 == 1 ){ + camera->film()->accumulate(pixel_id, L, 0.5f); + }; + }; }; Kernel1D build_grid_kernel = [&]() noexcept { auto index = static_cast(dispatch_x()); - auto radius = node()->initial_radius(); $if(viewpoints->nxt(index) == 0u) { viewpoints->link(index); }; @@ -663,11 +540,19 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins Kernel2D emit_photons_kernel = [&](UInt frame_index, Float time, BufferFloat grad_in) noexcept { auto pixel_id = dispatch_id().xy(); auto sampler_id = UInt2(pixel_id.x + resolution.x, pixel_id.y); - $if(pixel_id.x * resolution.y + pixel_id.y < photon_per_iter) { - photon_tracing_bp(camera, frame_index, sampler_id, time, pixel_id.x * resolution.y + pixel_id.y, grad_in); + $if(pixel_id.y * resolution.x + pixel_id.x < photon_per_iter) { + photon_tracing_bp(camera, frame_index, sampler_id, time, pixel_id.y * resolution.x + pixel_id.x, grad_in); }; }; + Kernel2D indirect_draw_kernel = [&](UInt tot_photon, UInt spp) noexcept { + set_block_size(16u, 16u, 1u); + auto pixel_id = dispatch_id().xy(); + auto pixel_id_1d = pixel_id.y * resolution.x + pixel_id.x; + auto L = get_indirect(camera->pipeline().spectrum(), pixel_id_1d, tot_photon); + camera->film()->accumulate(pixel_id, L, 0.5f * spp); + }; + Kernel1D indirect_update_kernel = [&]() noexcept { set_block_size(16u, 16u, 1u); auto pixel_id = dispatch_x(); @@ -679,22 +564,12 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins viewpoints->write_grid_len(indirect->radius(0u)); }; - //accumulate the stored indirect light into final image - Kernel2D indirect_draw_kernel = [&](UInt tot_photon, UInt spp) noexcept { - set_block_size(16u, 16u, 1u); - auto pixel_id = dispatch_id().xy(); - auto pixel_id_1d = pixel_id.x * resolution.y + pixel_id.y; - auto L = get_indirect(camera->pipeline().spectrum(), pixel_id_1d, tot_photon); - camera->film()->accumulate(pixel_id, L, 0.5f * spp); - }; Clock clock_compile; - auto indirect_initialize = pipeline().device().compile(indirect_initialize_kernel); auto viewpoint_reset = pipeline().device().compile(viewpoint_reset_kernel); auto viewpath_construct = pipeline().device().compile(viewpath_construct_kernel); auto build_grid = pipeline().device().compile(build_grid_kernel); auto emit_photon = pipeline().device().compile(emit_photons_kernel); - auto indirect_draw = pipeline().device().compile(indirect_draw_kernel); auto indirect_update = pipeline().device().compile(indirect_update_kernel); auto shared_update = pipeline().device().compile(shared_update_kernel); @@ -715,11 +590,12 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins command_buffer << indirect_initialize().dispatch(viewpoints_per_iter) << synchronize(); pipeline().update(command_buffer, 0); + for (auto s : shutter_samples) { - runtime_spp+=spp; - for (auto i = 0u; i < 1u; i++) { + runtime_spp+=node()->backward_iter(); + for (auto i = 0u; i < node()->backward_iter(); i++) { command_buffer << viewpoint_reset().dispatch(viewpoints->size()); - command_buffer << viewpath_construct(sample_id++, s.point.time, s.point.weight).dispatch(resolution); + command_buffer << viewpath_construct(sample_id++, s.point.time, s.point.weight, grad_in).dispatch(resolution); command_buffer << build_grid().dispatch(viewpoints->size()); command_buffer << emit_photon(sample_id++, s.point.time, grad_in).dispatch(make_uint2(add_x, resolution.y)); command_buffer << indirect_update().dispatch(viewpoints_per_iter); @@ -728,12 +604,14 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins } } } + command_buffer << synchronize(); - LUISA_INFO("Finish tracing"); - command_buffer << indirect_draw(node()->photon_per_iter(), runtime_spp).dispatch(resolution); - command_buffer << pipeline().printer().retrieve(); - command_buffer << synchronize(); - LUISA_INFO("Finish indirect_draw"); + // LUISA_INFO("Backward Rendering Forward Finished"); + // $if((node()->debug_mode() / 16) % 2 == 1) { + // command_buffer << indirect_draw(node()->photon_per_iter(), runtime_spp).dispatch(resolution); + // command_buffer << synchronize(); + // LUISA_INFO("Backward Rendering Indirect Finished"); + // }; progress.done(); auto render_time = clock.toc(); LUISA_INFO("Backward Rendering finished in {} ms.", render_time); @@ -741,7 +619,7 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins void _render_one_camera(CommandBuffer &command_buffer, Camera::Instance *camera) noexcept override { - if (node()->debug() == true) { + if ((node()->debug_mode() / 4) % 2 == 1) { Buffer grad_in = pipeline().device().create_buffer(camera->film()->node()->resolution().x * camera->film()->node()->resolution().y * 3); _render_one_camera_backward(command_buffer, 0, camera, grad_in); return; @@ -780,7 +658,6 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins auto viewpoints_per_iter = resolution.x * resolution.y; - logger = make_unique(pixel_count, node()->max_depth(), spectrum); indirect = make_unique(viewpoints_per_iter, spectrum, camera->film(), clamp, node()->shared_radius()); viewpoints = make_unique(viewpoints_per_iter, spectrum); //return; @@ -829,8 +706,8 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins Kernel2D emit_photons_kernel = [&](UInt frame_index, Float time) noexcept { auto pixel_id = dispatch_id().xy(); auto sampler_id = UInt2(pixel_id.x + resolution.x, pixel_id.y); - $if(pixel_id.x * resolution.y + pixel_id.y < photon_per_iter) { - photon_tracing(camera, frame_index, sampler_id, time, pixel_id.x * resolution.y + pixel_id.y); + $if(pixel_id.y * resolution.x + pixel_id.x < photon_per_iter) { + photon_tracing(camera, frame_index, sampler_id, time, pixel_id.y * resolution.x + pixel_id.x); }; }; @@ -849,7 +726,7 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins Kernel2D indirect_draw_kernel = [&](UInt tot_photon, UInt spp) noexcept { set_block_size(16u, 16u, 1u); auto pixel_id = dispatch_id().xy(); - auto pixel_id_1d = pixel_id.x * resolution.y + pixel_id.y; + auto pixel_id_1d = pixel_id.y * resolution.x + pixel_id.x; // $if(pixel_id_1d == 0) { // auto cur_n_tot = indirect->cur_n(0u); // pipeline().printer().info("photon cur n is {}", cur_n_tot); @@ -904,7 +781,7 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins command_buffer << indirect_draw(node()->photon_per_iter(), runtime_spp).dispatch(resolution); command_buffer << synchronize(); command_buffer << pipeline().printer().retrieve(); - LUISA_INFO("Finishi indirect_draw"); + //LUISA_INFO("Finishi indirect_draw"); progress.done(); auto render_time = clock.toc(); LUISA_INFO("Rendering finished in {} ms.", render_time); @@ -931,10 +808,8 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins SampledSpectrum testbeta{swl.dimension()}; auto ray = camera_ray; auto pdf_bsdf = def(1e16f); - - auto resolution = camera->film()->node()->resolution(); - auto pixel_id_1d = pixel_id.x*resolution.y+pixel_id.y; + auto pixel_id_1d = pixel_id.y*resolution.x+pixel_id.x; $for(depth, node()->max_depth()) { @@ -1051,11 +926,10 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins beta *= ite(q < rr_threshold, 1.0f / q, 1.f); }; }; - //return spectrum->srgb(swl, testbeta);//DEBUG return spectrum->srgb(swl, Li); } - void photon_tracing(const Camera::Instance *camera, Expr frame_index, + [[nodiscard]] void photon_tracing(const Camera::Instance *camera, Expr frame_index, Expr sampler_id, Expr time, Expr photon_id_1d) { sampler()->start(sampler_id, frame_index); // generate uniform samples @@ -1126,7 +1000,6 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins } indirect->add_phi(pixel_id, Phi*weight); indirect->add_cur_n(pixel_id, 1u); - indirect->add_cur_w(pixel_id, weight); //pipeline().printer().info("working here!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! pixel_id and phi are {}, {} ", pixel_id, Phi); }; viewpoint_index = viewpoints->nxt(viewpoint_index); @@ -1178,12 +1051,12 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins } [[nodiscard]] Float3 emit_viewpoint_bp(const Camera::Instance *camera, Expr frame_index, - Expr pixel_id, Expr time, Expr shutter_weight) noexcept { + Expr pixel_id, Expr time, Expr shutter_weight, BufferFloat &grad_in) noexcept { sampler()->start(pixel_id, frame_index); auto u_filter = sampler()->generate_pixel_2d(); auto u_lens = camera->node()->requires_lens_sampling() ? sampler()->generate_2d() : make_float2(.5f); - auto [camera_ray, _, camera_weight] = camera->generate_ray(pixel_id, time, u_filter, u_lens); + auto [camera_ray, pixel, camera_weight] = camera->generate_ray(pixel_id, time, u_filter, u_lens); auto spectrum = pipeline().spectrum(); auto swl = spectrum->sample(spectrum->node()->is_fixed() ? 0.f : sampler()->generate_1d()); SampledSpectrum beta{swl.dimension(), shutter_weight * camera_weight}; @@ -1193,35 +1066,42 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins auto pdf_bsdf = def(1e16f); auto resolution = camera->film()->node()->resolution(); - auto pixel_id_1d = pixel_id.x*resolution.y+pixel_id.y; - logger->add_start_camera(pixel_id_1d, ray, camera_weight); + auto pixel_id_1d = pixel_id.y*resolution.x+pixel_id.x; + + Float3 start_point = camera_ray->origin(); + Float3 end_point; + Float3 pixel_world; + ArrayUInt<4> triangle_ids, inst_ids; + ArrayFloat3<4> bary_coords, points, normals; + ArrayFloat<4> etas; + ArrayFloat2<4> grad_barys; + ArrayFloat3<4> grad_betas; + Bool flag = false; + + Float3 grad_rgb = make_float3(grad_in->read(pixel_id_1d*_grad_dimension), grad_in->read(pixel_id_1d*_grad_dimension+1), grad_in->read(pixel_id_1d*_grad_dimension+2)); + Float2 grad_pixel = make_float2(grad_in->read(pixel_id_1d*_grad_dimension+3), grad_in->read(pixel_id_1d*_grad_dimension+4)); + Float3 radiance = make_float3(1.0f, 1.0f, 1.0f);// = get_radiance_cache(pixel_id, frame_index); + + auto grad_pixel_world = make_float3(camera->camera_to_world() * make_float4(grad_pixel, 0.f, 1.f)); auto path_size = 0u; - $for(depth, node()->max_depth()) { + $for(depth, node()->max_depth()) { // trace auto wo = -ray->direction(); auto it = pipeline().geometry()->intersect(ray); - - $if(!it->valid()) { - if (pipeline().environment()) { - auto eval = light_sampler()->evaluate_miss(ray->direction(), swl, time); - Li += beta * eval.L * balance_heuristic(pdf_bsdf, eval.pdf); - logger->add_envlight_end(pixel_id_1d, path_size, ray->direction(), eval.L); - } - $break; - }; - - // hit light - if (!pipeline().lights().empty()) { - $if(it->shape().has_light()) { - auto eval = light_sampler()->evaluate_hit(*it, ray->origin(), swl, time); - Li += beta * eval.L * balance_heuristic(pdf_bsdf, eval.pdf); - logger->add_light_end(pixel_id_1d, path_size, it, eval.L); - }; - } $if(!it->shape().has_surface()) { $break; }; - + + $if(depth == 0) { + $autodiff { + auto p = it->p(); + requires_grad(p); + auto pixel_pos_world = (p - start_point) * detach(distance(pixel_world, start_point) / distance(p, start_point)); + backward(dot(pixel_pos_world, grad_pixel_world)); + grad_barys[0] = get_bary_grad(grad(p), it->instance_id(), it->triangle_id()); + }; + }; + // generate uniform samples auto u_light_selection = sampler()->generate_1d(); auto u_light_surface = sampler()->generate_2d(); @@ -1254,19 +1134,6 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins }); call.execute([&](auto closure) noexcept { - if (auto dispersive = closure->is_dispersive()) { - $if(*dispersive) { swl.terminate_secondary(); }; - } - logger->add_path_vertex(pixel_id_1d, path_size, it); - // direct lighting - $if(light_sample.eval.pdf > 0.0f & !occluded) { - auto wi = light_sample.shadow_ray->direction(); - auto eval = closure->evaluate(wo, wi); - auto w = balance_heuristic(light_sample.eval.pdf, eval.pdf) / - light_sample.eval.pdf; - Li += w * beta * eval.f * light_sample.eval.L; - logger->add_light_end(pixel_id_1d, path_size, it, light_sample.eval.L); - }; auto roughness = closure->roughness(); Bool stop_check = (roughness.x * roughness.y > 0.16f) | stop_direct; @@ -1275,6 +1142,14 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins viewpoints->push(it->p(), swl, beta, wo, pixel_id_1d, surface_tag); }; + $if(light_sample.eval.pdf > 0.0f & !occluded) { + auto wi = light_sample.shadow_ray->direction(); + auto eval = closure->evaluate(wo, wi); + auto w = balance_heuristic(light_sample.eval.pdf, eval.pdf) / + light_sample.eval.pdf; + Li += w * beta * eval.f * light_sample.eval.L; + }; + // sample material auto surface_sample = closure->sample(wo, u_lobe, u_bsdf); ray = it->spawn_ray(surface_sample.wi); @@ -1284,25 +1159,31 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins // apply eta scale auto eta = closure->eta().value_or(1.f); - - logger->set_path_vertex_attrib(pixel_id_1d, path_size, surface_sample, eta); - path_size += 1; $switch(surface_sample.event) { - $case(Surface::event_enter) { eta_scale = sqr(eta); }; - $case(Surface::event_exit) { eta_scale = sqr(1.f / eta); }; + $case(Surface::event_enter) { eta_scale = sqr(eta); etas[path_size] = eta;}; + $case(Surface::event_exit) { eta_scale = sqr(1.f / eta); etas[path_size] = 1.f/eta;}; }; + + etas[path_size] = eta; + points[path_size] = it->p(); + triangle_ids[path_size] = it->triangle_id(); + inst_ids[path_size] = it->instance_id(); + bary_coords[path_size] = it->bary_coord(); + points[path_size] = it->p(); + normals[path_size] = it->ng(); + path_size++; }); + beta = zero_if_any_nan(beta); $if(beta.all([](auto b) noexcept { return b <= 0.f; })) { $break; }; + $if(stop_direct) { auto it_next = pipeline().geometry()->intersect(ray); - // logger->add_indirect_end(pixel_id_1d, path_size, beta); $if(!it_next->valid()) { if (pipeline().environment()) { auto eval = light_sampler()->evaluate_miss(ray->direction(), swl, time); Li += beta * eval.L * balance_heuristic(pdf_bsdf, eval.pdf); - logger->add_envlight_end(pixel_id_1d, path_size, ray->direction(), eval.L); } }; // hit light @@ -1310,20 +1191,102 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins $if(it_next->shape().has_light()) { auto eval = light_sampler()->evaluate_hit(*it_next, ray->origin(), swl, time); Li += beta * eval.L * balance_heuristic(pdf_bsdf, eval.pdf); - logger->add_light_end(pixel_id_1d, path_size, it_next, eval.L); }; } + flag=true; + end_point = it->p(); $break; }; $if(depth + 1u >= rr_depth) { beta *= ite(q < rr_threshold, 1.0f / q, 1.f); }; }; - //return spectrum->srgb(swl, testbeta);//DEBUG + + $if(flag==true){ + ray = camera_ray; + $for(i,path_size) + { + Float3 point_cur = points[i], point_nxt, point_pre; + Float3 normal_cur = normals[i]; + + $if(i > 0) { + point_pre = points[i - 1]; + } $else { + point_pre = start_point; + }; + + $if(i call; + + auto wo = normalize(point_pre-point_cur); + auto it = pipeline().geometry()->intersect(ray); + + ray = it->spawn_ray(normalize(point_nxt-point_cur)); + pipeline().surfaces().dispatch(it->shape().surface_tag(), [&](auto surface) noexcept { + surface->closure(call, *it, swl, wo, 1.f, time); + }); + Float3 grad_p_cur, grad_p_nxt, grad_p_pre; + call.execute([&](auto closure) noexcept { + Float3 wi = normalize(point_nxt-point_cur); + Float3 wo = normalize(point_pre-point_cur); + auto eval = closure->evaluate(wo, wi); + auto evalf = make_float3(eval.f[0u],eval.f[1u],eval.f[2u]); + auto devalf = ite(evalf == 0.f, make_float3(0.f), make_float3(radiance[0u] / evalf[0u],radiance[1u] / evalf[1u],radiance[2u] / evalf[2u])); + SampledSpectrum a{3u},b{3u}; + a[0u] = grad_rgb[0u]; + a[1u] = grad_rgb[1u]; + a[2u] = grad_rgb[2u]; + b[0u] = devalf[0u]; + b[1u] = devalf[1u]; + b[2u] = devalf[2u]; + closure->backward(wo, wi, a*b); + $autodiff{ + requires_grad(point_cur, point_nxt, point_pre); + wi = normalize(point_nxt-point_cur); + wo = normalize(point_pre-point_cur); + auto eval = closure->evaluate(wo, wi); + evalf = make_float3(eval.f[0u], eval.f[1u], eval.f[2u]); + backward(grad_rgb*devalf*evalf); + grad_p_cur = grad(point_cur); + grad_p_nxt = grad(point_nxt); + grad_p_pre = grad(point_pre); + }; + }); + + $if(i>0){ + grad_barys[i-1]+=get_bary_grad(grad_p_pre, inst_ids[i-1], triangle_ids[i-1]); + }; + grad_barys[i]+=get_bary_grad(grad_p_cur, inst_ids[i], triangle_ids[i]); + $if(isrgb(swl, Li); } - void photon_tracing_bp(const Camera::Instance *camera, Expr frame_index, + [[nodisgard]] Float2 get_bary_grad(Float3 grad_in, UInt inst_id, UInt triangle_id){ + auto instance = pipeline().geometry()->instance(inst_id); + auto triangle = pipeline().geometry()->triangle(instance, triangle_id); + auto v_buffer = instance.vertex_buffer_id(); + + auto v0 = pipeline().buffer(v_buffer).read(triangle.i0); + auto v1 = pipeline().buffer(v_buffer).read(triangle.i1); + auto v2 = pipeline().buffer(v_buffer).read(triangle.i2); + + Float2 grad_bary; + grad_bary[0] = dot(v0->position() - v2->position(), grad_in); + grad_bary[1] = dot(v1->position() - v2->position(), grad_in); + return grad_bary; + }; + + [[nodiscard]] void photon_tracing_bp(const Camera::Instance *camera, Expr frame_index, Expr sampler_id, Expr time, Expr photon_id_1d, BufferFloat &grad_in) { sampler()->start(sampler_id, frame_index); @@ -1346,36 +1309,26 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins ArrayUInt<4> triangle_ids, inst_ids; ArrayFloat3<4> bary_coords, points, normals; - ArrayFloat<4> etas; + ArrayFloat<4> etas; ArrayFloat2<4> grad_barys; ArrayFloat3<4> grad_betas; - - ArrayFloat<8*8*2> mat_bary; - ArrayFloat3<8*4> mat_param; - + + auto resolution = camera->film()->node()->resolution(); auto max_depth = min(node()->max_depth(), 4u); auto tot_neighbors = 0u; $for(depth, max_depth) { - // trace - //path_size = depth; + auto wi = -ray->direction(); auto it = pipeline().geometry()->intersect(ray); - // miss + grad_barys[path_size] = make_float2(0.f); + grad_betas[path_size] = make_float3(0.f); + $if(!it->valid()) { - $if(photon_id_1d==node()->debug_photon()) - { - device_log("break it unvalid at size {} id {}", path_size, photon_id_1d); - device_log("ray {} {}", ray->origin(), ray->direction()); - }; $break; }; $if(!it->shape().has_surface()) { - $if(photon_id_1d==node()->debug_photon()) - { - device_log("break it non surface {} ", photon_id_1d); - }; $break; }; // generate uniform samples @@ -1422,27 +1375,52 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins call.execute([&](const Surface::Closure *closure) noexcept { eval_viewpoint = closure->evaluate(viewpoint_wo, wi).f; }); + Float2 grad_b{0.0f, 0.0f}; + Float grad_pixel_0, grad_pixel_1, grad_pixel_2; + Float rel_dis_diff, grad_dis; + Float weight; + Float3 Phi, Phi_real; $autodiff { Float3 beta_diff = make_float3(beta[0u], beta[1u], beta[2u]); requires_grad(bary, beta_diff); Float3 photon_pos = point_0 * bary[0] + point_1 * bary[1] + point_2 * (1 - bary[0] - bary[1]); - auto rel_dis_diff = distance(position, photon_pos) / rad; + rel_dis_diff = distance(position, photon_pos) / rad; + requires_grad(rel_dis_diff); auto rel3 = rel_dis_diff*rel_dis_diff*rel_dis_diff; - auto weight = 3.5f*(1- 6*rel3*rel_dis_diff*rel_dis_diff + 15*rel3*rel_dis_diff - 10*rel3); + weight = 3.5f*(1- 6*rel3*rel_dis_diff*rel_dis_diff + 15*rel3*rel_dis_diff - 10*rel3); auto wi_local = it->shading().world_to_local(wi); - auto Phi = spectrum->srgb(swl, viewpoint_beta * eval_viewpoint / abs_cos_theta(wi_local)); - auto Phi_beta = Phi * beta_diff * weight / 200000.0f; - auto _grad_dimension = 3u; - auto grad_pixel_0 = grad_in->read(pixel_id * _grad_dimension + 0); - auto grad_pixel_1 = grad_in->read(pixel_id * _grad_dimension + 1); - auto grad_pixel_2 = grad_in->read(pixel_id * _grad_dimension + 2); - auto dldPhi = (Phi_beta[0]*grad_pixel_0 + Phi_beta[1]*grad_pixel_1 + Phi_beta[2]*grad_pixel_2); + Phi = spectrum->srgb(swl, viewpoint_beta * eval_viewpoint / abs_cos_theta(wi_local)); + Phi_real = spectrum->srgb(swl, beta * viewpoint_beta * eval_viewpoint / abs_cos_theta(wi_local)); + auto Phi_beta = Phi * beta_diff * weight / (indirect->_photon_per_iter*1.0f); + grad_pixel_0 = grad_in->read(pixel_id * _grad_dimension + 0); + grad_pixel_1 = grad_in->read(pixel_id * _grad_dimension + 1); + grad_pixel_2 = grad_in->read(pixel_id * _grad_dimension + 2); + auto dldPhi = (Phi_beta[0u]*grad_pixel_0 + Phi_beta[1u]*grad_pixel_1 + Phi_beta[2u]*grad_pixel_2); backward(dldPhi); - auto grad_b = grad(bary).xy(); - grad_bary += grad_b; + grad_bary += grad(bary).xy(); grad_beta += grad(beta_diff); + grad_dis = grad(rel_dis_diff); }; count_neighbors+=1; + + $if((node()->debug_mode() / 2) % 2 == 1) { + $if(photon_id_1d()->debug_photon()) + { + $if(etas[0]==1.8f) + { + UInt2 pixel_id_2d = make_uint2(pixel_id % resolution.x, pixel_id / resolution.x); + camera->film()->accumulate(pixel_id_2d, make_float3(grad_dis,0.0f,0.0f)); + }; + }; + }; + + $if((node()->debug_mode() / 16) % 2 == 1) { + $if(photon_id_1d()->debug_photon()) + { + indirect->add_phi(pixel_id, Phi_real*weight); + indirect->add_cur_n(pixel_id, 1u); + }; + }; }; viewpoint_index = viewpoints->nxt(viewpoint_index); }; @@ -1451,8 +1429,8 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins }; $if(count_neighbors>0){ tot_neighbors+=count_neighbors; - grad_betas[path_size] = make_float3(grad_beta[0]/count_neighbors, grad_beta[1]/count_neighbors, grad_beta[2]/count_neighbors); - grad_barys[path_size] = make_float2(grad_bary[0]/count_neighbors, grad_bary[1]/count_neighbors); + grad_betas[path_size] += make_float3(grad_beta[0]/count_neighbors, grad_beta[1]/count_neighbors, grad_beta[2]/count_neighbors); + grad_barys[path_size] += make_float2(grad_bary[0]/count_neighbors, grad_bary[1]/count_neighbors); }; }; auto surface_tag = it->shape().surface_tag(); @@ -1462,10 +1440,6 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins surface->closure(call, *it, swl, wi, 1.f, time); }); call.execute([&](auto closure) noexcept { - // apply opacity map - if (auto dispersive = closure->is_dispersive()) { - $if(*dispersive) { swl.terminate_secondary(); }; - } // sample material auto surface_sample = closure->sample(wi, u_lobe, u_bsdf, TransportMode::IMPORTANCE); ray = it->spawn_ray(surface_sample.wi); @@ -1484,201 +1458,288 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins }); beta = zero_if_any_nan(beta); $if(beta.all([](auto b) noexcept { return b <= 0.f; })) { - $if(photon_id_1d==node()->debug_photon()) - { - device_log("break beta negative {}",photon_id_1d); - }; $break; }; auto rr_threshold = node()->rr_threshold(); auto q = max(eta_scale, .05f); $if(depth + 1u >= rr_depth) { - $if(q < rr_threshold & u_rr >= q) { - $if(photon_id_1d==node()->debug_photon()) - { - device_log("break rr {}",photon_id_1d); - }; + $if(q < rr_threshold & u_rr >= q) { $break; }; beta *= ite(q < rr_threshold, 1.0f / q, 1.f); }; + triangle_ids[path_size] = it->triangle_id(); inst_ids[path_size] = it->instance_id(); bary_coords[path_size] = it->bary_coord(); points[path_size] = it->p(); normals[path_size] = it->ng(); - // $if(photon_id_1d==node()->debug_photon()) - // { - // device_log("adding path size {} id is {}",path_size,photon_id_1d); - // }; path_size = path_size+1u; }; - // $if(photon_id_1d==node()->debug_photon()){ - // device_log("path size is {} id is {}",path_size,photon_id_1d); - // $for(i,path_size) - // { - // device_log("{} point {} normal {} eta {}, bary {}",i, points[i], normals[i], etas[i], bary_coords[i]); - // }; - // }; - // $if(path_size>=2u) - // { - // device_log("path size is {}",path_size); - // device_log("phpton_id is {}",photon_id_1d); - // $for(i,path_size) - // { - // device_log("{} point {} normal {} eta {} inst {}, bary {}",i, points[i], normals[i], etas[i], inst_ids[i], bary_coords[i]); - // }; - // return; - // EPSM_photon(path_size, points, normals, inst_ids, triangle_ids, bary_coords, etas, light_sample, grad_betas, grad_barys, mat_bary, mat_param); - // }; + //@Todo: log surface event $if(photon_id_1d()->debug_photon()) { - $if(path_size>=2&path_size<=4) { - EPSM_photon(path_size, points, normals, inst_ids, triangle_ids, bary_coords, etas, light_sample, grad_betas, grad_barys, mat_bary, mat_param); + $if(path_size==2) { + $if(etas[0]==1.8f) + { + $if((node()->debug_mode() / 1) % 2 == 1) { + device_log("path size is {}",path_size); + device_log("photon_id is {}",photon_id_1d); + device_log("light sample is {}",light_sample.shadow_ray->origin()); + $for(i,path_size) + { + device_log("{} point {} normal {} eta {} inst {}, bary {} grad_bary {}",i, points[i], normals[i], etas[i], inst_ids[i], bary_coords[i], grad_barys[i]); + }; + }; + EPSM_photon(path_size, points, normals, inst_ids, triangle_ids, bary_coords, etas, light_sample.shadow_ray->origin(), grad_barys, photon_id_1d); + }; }; }; } - void EPSM_photon(UInt path_size, ArrayFloat3<4> &points, ArrayFloat3<4> &normals, ArrayUInt<4> &inst_ids, ArrayUInt<4> &triangle_ids, ArrayFloat3<4> &bary_coords, - ArrayFloat<4> &etas, LightSampler::Sample &light_sample, ArrayFloat3<4> grad_beta, ArrayFloat2<4> grad_bary, ArrayFloat<8 * 8 * 2> &mat_bary, ArrayFloat3<8 * 4> &mat_param){ - { - auto locate = [&](UInt i, UInt j) { - return (i<<4|j); - }; - auto locate_adj = [&](UInt i, UInt j) { - return (i<<4|8|j); + + struct GradStruct { + Float3 grad_p_pre; + Float3 grad_p_cur; + Float3 grad_p_nxt; + Float3 grad_n_cur; + Float grad_eta; + }; + + GradStruct half_vec_constraint(Float3 point_pre, Float3 point_cur, Float3 point_nxt, Float3 normal_cur, Float eta, size_t j) { + Float3 grad_p_pre{0.f, 0.f, 0.f}, grad_p_cur{0.f, 0.f, 0.f}, grad_p_nxt{0.f, 0.f, 0.f}, grad_n_cur{0.f, 0.f, 0.f}; + Float grad_eta{0.f}; + $autodiff{ + requires_grad(point_pre, point_cur, point_nxt, normal_cur, eta); + auto wi = normalize(point_cur-point_nxt); + auto wo = normalize(point_nxt-point_cur); + auto f = Frame::make(normal_cur); + auto wi_local = f.world_to_local(wi); + auto wo_local = f.world_to_local(wo); + auto res = normalize(wi_local+wo_local*eta); + backward(res[j]); + grad_p_pre = grad(point_pre); + grad_p_cur = grad(point_cur); + grad_p_nxt = grad(point_nxt); + grad_n_cur = grad(normal_cur); + grad_eta = grad(eta); }; + return {grad_p_pre, grad_p_cur, grad_p_nxt, grad_n_cur, grad_eta}; + } - Callable inverse_matrix = [&]() { - //inverse a matrix which mat->read(i,j) gives the (i,j) element - auto n = path_size*2; - $for (i,n) { - mat_bary[locate_adj(i,i)] = 1; - }; - // $for (i,8) { - // device_log("mat_bary {} {} {} {} {} {} {} {}",mat_bary[locate(i,0)],mat_bary[locate(i,1)],mat_bary[locate(i,2)],mat_bary[locate(i,3)],mat_bary[locate(i,4)],mat_bary[locate(i,5)],mat_bary[locate(i,6)],mat_bary[locate(i,7)]); - // }; - $for (i, n) { - $if (mat_bary[locate(i,i)]==0) { - auto p=n; - $for (j,i + 1, n) { - $if (mat_bary[locate(j,i)] != 0) { - p=j; - $break; - }; - }; - $for(k,n) - { - auto t = mat_bary[locate(i, k)]; - mat_bary[locate(i,k)] = mat_bary[locate(p,k)]; - mat_bary[locate(p,k)] = t; - - t = mat_bary[locate_adj(i,k)]; - mat_bary[locate_adj(i,k)] = mat_bary[locate_adj(p,k)]; - mat_bary[locate_adj(p,k)] = t; + auto locate(UInt i, UInt j) { + return (i*max_size*4u+j); + }; + + auto locate_adj(UInt i, UInt j) { + return (i*max_size*4+max_size*2+j); + }; + + void inverse_matrix(ArrayFloat<8*8*2> mat_bary, UInt path_size){ + auto n = path_size*2; + $for (i,n) { + mat_bary[locate_adj(i,i)] = 1; + }; + // ArrayFloat<8 * 8 * 2> mat_bary_copy{mat_bary}; + // $if((node()->debug_mode() / 1) % 2 == 1) + // { + // $for (i,8) { + // device_log("mat_bary {} {} {} {} {} {} {} {}",mat_bary[locate(i,0)],mat_bary[locate(i,1)],mat_bary[locate(i,2)],mat_bary[locate(i,3)],mat_bary[locate(i,4)],mat_bary[locate(i,5)],mat_bary[locate(i,6)],mat_bary[locate(i,7)]); + // }; + // }; + $for (i, n) { + $if (mat_bary[locate(i,i)]==0) { + auto p=n; + $for (j,i + 1, n) { + $if (mat_bary[locate(j,i)] != 0) { + p=j; + $break; }; }; - $for(j, n) { - $if(j!=i) - { - auto factor = mat_bary[locate(j,i)]/mat_bary[locate(i,i)]; - $for (k,i,n) { - mat_bary[locate(j,k)] = mat_bary[locate(j,k)] - factor * mat_bary[locate(i,k)]; - }; - $for (k,n) { - mat_bary[locate_adj(j,k)] = mat_bary[locate_adj(j,k)] - factor * mat_bary[locate_adj(i,k)]; - }; - }; + $for(k,n) + { + auto t = mat_bary[locate(i, k)]; + mat_bary[locate(i,k)] = mat_bary[locate(p,k)]; + mat_bary[locate(p,k)] = t; + + t = mat_bary[locate_adj(i,k)]; + mat_bary[locate_adj(i,k)] = mat_bary[locate_adj(p,k)]; + mat_bary[locate_adj(p,k)] = t; }; }; - $for(i, n) { - auto f = mat_bary[locate(i,i)]; - $for(j,i,n) { - mat_bary[locate(i,j)] /= f; - }; - $for(j, n) { - mat_bary[locate_adj(i,j)] /= f; + $for(j, n) { + $if(j!=i) + { + auto factor = mat_bary[locate(j,i)]/mat_bary[locate(i,i)]; + $for (k,i,n) { + mat_bary[locate(j,k)] = mat_bary[locate(j,k)] - factor * mat_bary[locate(i,k)]; + }; + $for (k,n) { + mat_bary[locate_adj(j,k)] = mat_bary[locate_adj(j,k)] - factor * mat_bary[locate_adj(i,k)]; + }; }; }; }; - Callable compute_and_scatter_grad = [&](){ - ArrayFloat<8> tmp; - auto n = path_size*2; - $for(i, n){ - tmp[i] = 0.0f; - $for(j, n/2){ - tmp[i] -= grad_bary[j][0] * mat_bary[locate_adj(j * 2, i)] + grad_bary[j][1] * mat_bary[locate_adj(j * 2 + 1, i)]; - }; - }; - $if(node()->debug()) + $for(i, n) { + auto f = mat_bary[locate(i,i)]; + $for(j,i,n) { + mat_bary[locate(i,j)] /= f; + }; + $for(j, n) { + mat_bary[locate_adj(i,j)] /= f; + }; + }; + // $if((node()->debug_mode() / 1) % 2 == 1) + // { + // $for (i,8){ + // $for(j,8) + // { + // mat_bary_copy[locate_adj(i,j)] = 0; + // $for (k,8) { + // mat_bary_copy[locate_adj(i,j)]+=mat_bary_copy[locate(i,k)]*mat_bary[locate_adj(k,j)]; + // }; + // }; + // device_log("should be id {} {} {} {} {} {} {} {}",mat_bary_copy[locate_adj(i,0)],mat_bary_copy[locate_adj(i,1)],mat_bary_copy[locate_adj(i,2)],mat_bary_copy[locate_adj(i,3)],mat_bary_copy[locate_adj(i,4)],mat_bary_copy[locate_adj(i,5)],mat_bary_copy[locate_adj(i,6)],mat_bary_copy[locate_adj(i,7)]); + // }; + // $for (i,8) { + // device_log("mat_bary {} {} {} {} {} {} {} {}",mat_bary[locate(i,0)],mat_bary[locate(i,1)],mat_bary[locate(i,2)],mat_bary[locate(i,3)],mat_bary[locate(i,4)],mat_bary[locate(i,5)],mat_bary[locate(i,6)],mat_bary[locate(i,7)]); + // }; + // }; + } + + void compute_and_scatter_grad(ArrayFloat<8 * 8 * 2> &mat_bary, ArrayFloat3<8 * 5 + 2 * 2> &mat_param, ArrayFloat2<4> &grad_bary, UInt path_size, ArrayUInt<4> &inst_ids, ArrayUInt<4> &triangle_ids, ArrayFloat3<4> &bary_coords) { + ArrayFloat<8> tmp; + auto n = path_size*2; + $for(i, n){ + tmp[i] = 0.0f; + $for(j, n/2){ + tmp[i] -= grad_bary[j][0] * mat_bary[locate_adj(j * 2, i)] + grad_bary[j][1] * mat_bary[locate_adj(j * 2 + 1, i)]; + }; + }; + $if((node()->debug_mode() / 1) % 2 == 1) + { + $for (i,8) { + device_log("mat_param {} {} {} {}",mat_param[i*4+0],mat_param[i*4+1],mat_param[i*4+2],mat_param[i*4+3]); + }; + $for(j, n/2){ + device_log("grad_bary {} {} ",grad_bary[j][0], grad_bary[j][1]); + }; + }; + //Todo add eta support + $for(i, n/2) { + Float3 grad_vertex = make_float3(0.0f), grad_normal = make_float3(0.0f); + grad_vertex = tmp[i*2] * mat_param[((i * 2) << 2) + 2] + tmp[i*2+1] * mat_param[((i * 2 + 1) << 2) + 2]; + $if(i()->debug_mode() / 1) % 2 == 1) + { + device_log("grad_vertex {} grad_normal {} bary {}",grad_vertex, grad_normal, bary_coords[i]); }; + pipeline().differentiation()->add_geom_gradients(grad_vertex, grad_normal, bary_coords[i], inst_ids[i], triangle_ids[i]); + }; + }; + } + + void EPSM_path(UInt path_size, ArrayFloat2<4> grad_bary, Float3 start_point, Float3 end_point, ArrayUInt<4> &inst_ids, ArrayUInt<4> &triangle_ids, ArrayFloat3<4> &bary_coords, ArrayFloat<4> &etas, ArrayFloat3<4> &points, ArrayFloat3<4> &normals) { + ArrayFloat<8*8*2> mat_bary; + ArrayFloat3<8*5+2*2> mat_param; + Float3 point_pre, point_cur, point_nxt, normal_cur; + + $for(id, path_size){ + + Float3 bary_cur = bary_coords[id]; + Float3 bary_nxt = bary_coords[id+1]; + + point_cur = points[id]; + normal_cur = normals[id]; + + $if (id>0){ + point_pre = points[id-1]; + } $else{ + point_pre = start_point; }; - $for(i, n/2) { - Float3 grad_vertex = make_float3(0.0f), grad_normal = make_float3(0.0f); - grad_vertex = tmp[(i+1) * 2 - 1] * mat_param[((i * 2 - 1) << 2) + 2] + tmp[(i+1) * 2 - 2] * mat_param[((i * 2 - 2) << 2) + 2]; - $if(i0){ + auto grad_b_pre = get_bary_grad(grad_p_pre, inst_ids[id-1], triangle_ids[id-1]); + mat_bary[locate(id*2+j,id*2-2)] = grad_b_pre[0]; + mat_bary[locate(id*2+j,id*2-1)] = grad_b_pre[1]; + mat_param[(id*2+j)*param_size_per_vert+0] = grad_p_pre; + } $else{ + mat_param[max_size*param_size_per_vert+j] = grad_p_pre; }; - $if(node()->debug()) - { - device_log("grad_vertex {} grad_normal {}",grad_vertex, grad_normal); + + $if (idadd_geom_gradients(grad_vertex, grad_normal, bary_coords[i], inst_ids[i], triangle_ids[i]); - }; - }; + } }; - // Float3 point_pre = light_sample.second->origin(); - // Float2 bary_pre = Float2(0.0f); - // Float2 bary_cur = bary_coords[0]; - // Float3 point_cur = points[0]; - // Float3 point_nxt,normal_cur, bary_nxt; - - auto instance = pipeline().geometry()->instance(inst_ids[0]); - auto triangle = pipeline().geometry()->triangle(instance, triangle_ids[0]); - auto v_buffer = instance.vertex_buffer_id(); - auto v0 = pipeline().buffer(v_buffer).read(triangle.i0); - auto v1 = pipeline().buffer(v_buffer).read(triangle.i1); - auto v2 = pipeline().buffer(v_buffer).read(triangle.i2); - - auto point_pre_0 = light_sample.shadow_ray->origin(); - auto point_pre_1 = light_sample.shadow_ray->origin(); - auto point_pre_2 = light_sample.shadow_ray->origin(); - - auto point_cur_0 = v0->position(); - auto point_cur_1 = v1->position(); - auto point_cur_2 = v2->position(); - + inverse_matrix(mat_bary, path_size); + compute_and_scatter_grad(mat_bary, mat_param, grad_bary, path_size, inst_ids, triangle_ids, bary_coords); + } + + void EPSM_photon(UInt path_size, ArrayFloat3<4> &points, ArrayFloat3<4> &normals, ArrayUInt<4> &inst_ids, ArrayUInt<4> &triangle_ids, ArrayFloat3<4> &bary_coords, ArrayFloat<4> &etas, Float3 start_point, ArrayFloat2<4> grad_bary, UInt photon_id_1d){ + { + ArrayFloat<8*8*2> mat_bary; + ArrayFloat3<8*5+2*2> mat_param; + Float3 point_pre, point_cur, point_nxt, normal_cur; $for(id, path_size-1){ - auto normal_cur_0 = v0->normal(); - auto normal_cur_1 = v1->normal(); - auto normal_cur_2 = v2->normal(); - instance = pipeline().geometry()->instance(inst_ids[id + 1]); - triangle = pipeline().geometry()->triangle(instance, triangle_ids[id + 1]); - v_buffer = instance.vertex_buffer_id(); - v0 = pipeline().buffer(v_buffer).read(triangle.i0); - v1 = pipeline().buffer(v_buffer).read(triangle.i1); - v2 = pipeline().buffer(v_buffer).read(triangle.i2); - auto point_nxt_0 = v0->position(); - auto point_nxt_1 = v1->position(); - auto point_nxt_2 = v2->position(); - // point_nxt = points[id+1]; - // bary_nxt = bary_coords[id+1]; - // normal_cur = normals[id]; - Float3 bary_pre; + + Float3 bary_cur = bary_coords[id]; + Float3 bary_nxt = bary_coords[id+1]; + + point_cur = points[id]; + normal_cur = normals[id]; + + $if (id>0){ + point_pre = points[id-1]; + } $else{ + point_pre = start_point; + }; + + point_nxt = points[id+1]; + $if(id==0) { - Float3 bary_cur = bary_coords[id]; - bary_pre = bary_cur; + Float3 bary_cur = bary_coords[0]; + auto instance = pipeline().geometry()->instance(inst_ids[id]); + auto triangle = pipeline().geometry()->triangle(instance, triangle_ids[id]); + auto v_buffer = instance.vertex_buffer_id(); + + auto v0 = pipeline().buffer(v_buffer).read(triangle.i0); + auto v1 = pipeline().buffer(v_buffer).read(triangle.i1); + auto v2 = pipeline().buffer(v_buffer).read(triangle.i2); + + auto point_cur_0 = v0->position(); + auto point_cur_1 = v1->position(); + auto point_cur_2 = v2->position(); $autodiff{ requires_grad(bary_cur); - Float3 point_pre = light_sample.shadow_ray->origin(); Float3 point_cur = point_cur_0*bary_cur[0]+point_cur_1*bary_cur[1]+point_cur_2*(1-bary_cur[0]-bary_cur[1]); requires_grad(point_cur); auto wi_diff = normalize(point_cur-point_pre); @@ -1690,7 +1751,6 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins }; $autodiff { requires_grad(bary_cur); - Float3 point_pre = light_sample.shadow_ray->origin(); Float3 point_cur = point_cur_0 * bary_cur[0] + point_cur_1 * bary_cur[1] + point_cur_2 * (1 - bary_cur[0] - bary_cur[1]); requires_grad(point_cur); auto wi_diff = normalize(point_cur - point_pre); @@ -1700,83 +1760,40 @@ class MegakernelPhotonMappingDiffInstance : public DifferentiableIntegrator::Ins mat_bary[locate(1, 0)] = grad_b[0]; mat_bary[locate(1, 1)] = grad_b[1]; }; - } $else{ - bary_pre = bary_coords[id-1]; - }; - ArrayFloat3<2> grad_uv_pre, grad_uv_cur, grad_uv_nxt; - $autodiff { - Float3 bary_cur = bary_coords[id]; - Float3 bary_nxt = bary_coords[id+1]; - requires_grad(bary_pre, bary_cur, bary_nxt); - Float3 point_pre = point_pre_0*bary_pre[0]+point_pre_1*bary_pre[1]+point_pre_2*(1-bary_pre[0]-bary_pre[1]); - Float3 point_cur = point_cur_0*bary_cur[0]+point_cur_1*bary_cur[1]+point_cur_2*(1-bary_cur[0]-bary_cur[1]); - Float3 point_nxt = point_nxt_0*bary_nxt[0]+point_nxt_1*bary_nxt[1]+point_nxt_2*(1-bary_nxt[0]-bary_nxt[1]); - Float3 normal_cur = normal_cur_0*bary_cur[0]+normal_cur_1*bary_cur[1]+normal_cur_2*(1-bary_cur[0]-bary_cur[1]); - requires_grad(point_pre, point_cur, point_nxt, normal_cur); - auto wi = normalize(point_pre-point_cur); - auto wo = normalize(point_nxt-point_cur); - auto f = Frame::make(normal_cur); - auto wi_local = f.world_to_local(wi); - auto wo_local = f.world_to_local(wo); - auto res = normalize(wi_local+wo_local*etas[id]); - backward(res[0]); - grad_uv_pre[0] = grad(bary_pre); - grad_uv_cur[0] = grad(bary_cur); - grad_uv_nxt[0] = grad(bary_nxt); - mat_param[((((id+1)<<1)|0)<<2)] = grad(point_pre); - mat_param[((((id+1)<<1)|0)<<2)|1] = grad(point_cur); - mat_param[((((id+1)<<1)|0)<<2)|2] = grad(point_nxt); - mat_param[((((id+1)<<1)|0)<<2)|3] = grad(normal_cur); }; - $autodiff { - Float3 bary_cur = bary_coords[id]; - Float3 bary_nxt = bary_coords[id+1]; - requires_grad(bary_pre, bary_cur, bary_nxt); - Float3 point_pre = point_pre_0*bary_pre[0]+point_pre_1*bary_pre[1]+point_pre_2*(1-bary_pre[0]-bary_pre[1]); - Float3 point_cur = point_cur_0*bary_cur[0]+point_cur_1*bary_cur[1]+point_cur_2*(1-bary_cur[0]-bary_cur[1]); - Float3 point_nxt = point_nxt_0*bary_nxt[0]+point_nxt_1*bary_nxt[1]+point_nxt_2*(1-bary_nxt[0]-bary_nxt[1]); - Float3 normal_cur = normal_cur_0*bary_cur[0]+normal_cur_1*bary_cur[1]+normal_cur_2*(1-bary_cur[0]-bary_cur[1]); - requires_grad(point_pre, point_cur, point_nxt, normal_cur); - auto wi = normalize(point_pre-point_cur); - auto wo = normalize(point_nxt-point_cur); - auto f = Frame::make(normal_cur); - auto wi_local = f.world_to_local(wi); - auto wo_local = f.world_to_local(wo); - auto res = normalize(wi_local+wo_local*etas[id]); - backward(res[1]); - grad_uv_pre[1] = grad(bary_pre); - grad_uv_cur[1] = grad(bary_cur); - grad_uv_nxt[1] = grad(bary_nxt); - mat_param[((((id+1)<<1)|1)<<2)] = grad(point_pre); - mat_param[((((id+1)<<1)|1)<<2)|1] = grad(point_cur); - mat_param[((((id+1)<<1)|1)<<2)|2] = grad(point_nxt); - mat_param[((((id+1)<<1)|1)<<2)|3] = grad(normal_cur); - }; - for(uint j=0;j<2;j++) - { - $if(id>0) - { - mat_bary[locate(id*2+2+j,id*2-2)] = grad_uv_pre[j][0]; - mat_bary[locate(id*2+2+j,id*2-1)] = grad_uv_pre[j][1]; + uint offset = 2u; + for(uint j=0;j<2;j++){ + auto [grad_p_pre, grad_p_cur, grad_p_nxt, grad_n_cur, grad_eta] = half_vec_constraint(point_pre, point_cur, point_nxt, normal_cur, etas[id], j); + auto grad_b_cur = get_bary_grad(grad_p_cur+grad_n_cur, inst_ids[id], triangle_ids[id]); + mat_bary[locate(offset+id*2+j,id*2+0)] = grad_b_cur[0]; + mat_bary[locate(offset+id*2+j,id*2+1)] = grad_b_cur[1]; + + mat_param[(offset+id*2+j)*param_size_per_vert+1] = grad_p_cur; + mat_param[(offset+id*2+j)*param_size_per_vert+3] = grad_n_cur; + mat_param[(offset + id * 2 + j) * param_size_per_vert + 4] = make_float3(grad_eta); + + $if (id>0){ + auto grad_b_pre = get_bary_grad(grad_p_pre, inst_ids[id-1], triangle_ids[id-1]); + mat_bary[locate(offset+id*2+j,id*2-2)] = grad_b_pre[0]; + mat_bary[locate(offset+id*2+j,id*2-1)] = grad_b_pre[1]; + mat_param[(offset+id*2+j)*param_size_per_vert+0] = grad_p_pre; + }; + + $if (id MegakernelPhotonMappingDiff::build( Pipeline &pipeline, CommandBuffer &command_buffer) const noexcept { return luisa::make_unique( diff --git a/src/python/lrapi.cpp b/src/python/lrapi.cpp index 9c6774b4..c936a969 100644 --- a/src/python/lrapi.cpp +++ b/src/python/lrapi.cpp @@ -207,14 +207,14 @@ PYBIND11_MODULE(_lrapi, m) { }); m.def("render", []() { - LUISA_INFO("LuisaRender API render_scene"); + //LUISA_INFO("LuisaRender API render_scene"); auto res = scene_python._pipeline->render_with_return(*scene_python._stream); scene_python._stream->synchronize(); std::vector res_vec(res.size()); for (int i = 0; i < res.size(); i++) { res_vec[i] = reinterpret_cast(res[i]); } - LUISA_INFO("res_vec: {}",res_vec[0]); + //LUISA_INFO("res_vec: {}",res_vec[0]); return res_vec; }); @@ -225,7 +225,7 @@ PYBIND11_MODULE(_lrapi, m) { m.def("update_scene", [](std::vector params) { - LUISA_INFO("LuisaRender Update Scene"); + //LUISA_INFO("LuisaRender Update Scene"); luisa::vector constants{}; luisa::vector> textures{}; luisa::vector> geoms{}; @@ -233,7 +233,7 @@ PYBIND11_MODULE(_lrapi, m) { luisa::vector textures_id{}; luisa::vector geoms_id{}; for (auto param: params) { - LUISA_INFO("Param: {} {} {} {} {}", param.type, param.id, param.size, param.buffer_ptr, param.value); + //LUISA_INFO("Param: {} {} {} {} {}", param.type, param.id, param.size, param.buffer_ptr, param.value); if(param.type == "constant") { constants_id.push_back(param.id); constants.push_back(param.value); @@ -250,12 +250,12 @@ PYBIND11_MODULE(_lrapi, m) { geoms.push_back(std::move(buffer)); } } - LUISA_INFO("geom_id_size is {}", geoms_id.size()); + //LUISA_INFO("geom_id_size is {}", geoms_id.size()); scene_python._pipeline->differentiation()->update_parameter_from_external(*scene_python._stream, constants_id, constants, textures_id, textures, geoms_id, geoms); }); m.def("get_scene_param", [](std::vector params) { - LUISA_INFO("LuisaRender API get_parameter start"); + //LUISA_INFO("LuisaRender API get_parameter start"); luisa::vector constants_id{}; luisa::vector textures_id{}; luisa::vector geoms_id{}; @@ -270,12 +270,14 @@ PYBIND11_MODULE(_lrapi, m) { geoms_id.push_back(param.id); } } - auto [geom_param, geom_size] = scene_python._pipeline->differentiation()->get_parameter_from_external(*scene_python._stream, constants_id, textures_id, geoms_id); + auto [geom_param, geom_size, tri_param, tri_size] = scene_python._pipeline->differentiation()->get_parameter_from_external(*scene_python._stream, constants_id, textures_id, geoms_id); // std::vector ret_con_param(constants_id.size()); // std::vector ret_tex_param(textures_id.size()); // std::vector ret_tex_size(textures_id.size()); std::vector ret_geom_param(geoms_id.size()); std::vector ret_geom_size(geoms_id.size()); + std::vector ret_tri_param(geoms_id.size()); + std::vector ret_tri_size(geoms_id.size()); // for (int i = 0; i < ret_con_param.size(); i++) { // ret_con_param[i] = constant_param[i]; // } @@ -286,21 +288,30 @@ PYBIND11_MODULE(_lrapi, m) { for (int i = 0; i < ret_geom_param.size(); i++) { ret_geom_param[i] = reinterpret_cast(geom_param[i]); ret_geom_size[i] = geom_size[i]; - LUISA_INFO("LuisaRender API get_parameter {} {} {}", i, ret_geom_size[i], ret_geom_param[i]); + ret_tri_param[i] = reinterpret_cast(tri_param[i]); + ret_tri_size[i] = tri_size[i]; + //LUISA_INFO("LuisaRender API get_parameter {} {} {}", i, ret_geom_size[i], ret_geom_param[i]); } - LUISA_INFO("LuisaRender API get_parameter finish"); - return std::make_tuple(ret_geom_param, ret_geom_size); + //LUISA_INFO("LuisaRender API get_parameter finish"); + return std::make_tuple(ret_geom_param, ret_geom_size, ret_tri_param, ret_tri_size); }); m.def("render_backward", [](std::vector grad_ptr,std::vector sizes){ - LUISA_INFO("LuisaRender API render_backward"); + + //LUISA_INFO("LuisaRender API render_backward"); //scene_python._pipeline->differentiation()->clear_gradients(*scene_python._stream); luisa::vector> grad_buffer{grad_ptr.size()}; for (int i = 0; i < grad_ptr.size(); i++) { auto buffer = scene_python._pipeline->device().import_external_buffer(reinterpret_cast(grad_ptr[i]),sizes[i]); grad_buffer[i] = std::move(buffer); } - scene_python._pipeline->render_diff(*scene_python._stream, grad_buffer); + auto res = scene_python._pipeline->render_diff(*scene_python._stream, grad_buffer); + scene_python._stream->synchronize(); + std::vector res_vec(res.size()); + for (int i = 0; i < res.size(); i++) { + res_vec[i] = reinterpret_cast(res[i]); + } + return res_vec; }); m.def("get_gradients", [](){ diff --git a/src/shapes/mesh.cpp b/src/shapes/mesh.cpp index 69ec8aea..74258415 100644 --- a/src/shapes/mesh.cpp +++ b/src/shapes/mesh.cpp @@ -46,11 +46,10 @@ class MeshLoader { importer.SetPropertyInteger( AI_CONFIG_PP_SBP_REMOVE, aiPrimitiveType_LINE | aiPrimitiveType_POINT); importer.SetPropertyFloat(AI_CONFIG_PP_GSN_MAX_SMOOTHING_ANGLE, 45.f); - auto import_flags = aiProcess_RemoveComponent | aiProcess_SortByPType | - aiProcess_ValidateDataStructure | aiProcess_ImproveCacheLocality | - aiProcess_PreTransformVertices | aiProcess_FindInvalidData | + + uint import_flags = aiProcess_PreTransformVertices | aiProcess_JoinIdenticalVertices; - + //import_flags = 0u; auto remove_flags = aiComponent_ANIMATIONS | aiComponent_BONEWEIGHTS | aiComponent_CAMERAS | aiComponent_LIGHTS | aiComponent_MATERIALS | aiComponent_TEXTURES | diff --git a/src/tests/CMakeLists.txt b/src/tests/CMakeLists.txt index 14d668de..0becf073 100644 --- a/src/tests/CMakeLists.txt +++ b/src/tests/CMakeLists.txt @@ -10,3 +10,6 @@ target_link_libraries(test_sky PRIVATE luisa-render-texture-sky-precompute) add_executable(test_sphere test_sphere.cpp) target_link_libraries(test_sphere PRIVATE luisa::render) + +add_executable(test_ad test_ad.cpp) +target_link_libraries(test_ad PRIVATE luisa::render) diff --git a/src/tests/test_ad.cpp b/src/tests/test_ad.cpp new file mode 100644 index 00000000..37d2c34e --- /dev/null +++ b/src/tests/test_ad.cpp @@ -0,0 +1,181 @@ + +// Created by Mike on 2021/12/7. +// + +#include +#include + +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace luisa; +using namespace luisa::render; +using namespace luisa::compute; +// + +[[nodiscard]] auto parse_cli_options(int argc, const char *const *argv) noexcept { + cxxopts::Options cli{"luisa-render-cli"}; + cli.add_option("", "b", "backend", "Compute backend name", cxxopts::value(), ""); + cli.add_option("", "d", "device", "Compute device index", cxxopts::value()->default_value("-1"), ""); + cli.add_option("", "", "scene", "Path to scene description file", cxxopts::value(), ""); + cli.add_option("", "D", "define", "Parameter definitions to override scene description macros.", + cxxopts::value>()->default_value(""), "="); + cli.add_option("", "h", "help", "Display this help message", cxxopts::value()->default_value("false"), ""); + cli.allow_unrecognised_options(); + cli.positional_help(""); + cli.parse_positional("scene"); + auto options = [&] { + try { + return cli.parse(argc, argv); + } catch (const std::exception &e) { + LUISA_WARNING_WITH_LOCATION( + "Failed to parse command line arguments: {}.", + e.what()); + std::cout << cli.help() << std::endl; + exit(-1); + } + }(); + if (options["help"].as()) { + std::cout << cli.help() << std::endl; + exit(0); + } + if (options["scene"].count() == 0u) [[unlikely]] { + LUISA_WARNING_WITH_LOCATION("Scene file not specified."); + std::cout << cli.help() << std::endl; + exit(-1); + } + if (auto unknown = options.unmatched(); !unknown.empty()) [[unlikely]] { + luisa::string opts{unknown.front()}; + for (auto &&u : luisa::span{unknown}.subspan(1)) { + opts.append("; ").append(u); + } + LUISA_WARNING_WITH_LOCATION( + "Unrecognized options: {}", opts); + } + return options; +} + +using namespace luisa; +using namespace luisa::compute; +using namespace luisa::render; + +[[nodiscard]] auto parse_cli_macros(int &argc, char *argv[]) { + SceneParser::MacroMap macros; + + auto parse_macro = [¯os](luisa::string_view d) noexcept { + if (auto p = d.find('='); p == luisa::string::npos) [[unlikely]] { + LUISA_WARNING_WITH_LOCATION( + "Invalid definition: {}", d); + } else { + auto key = d.substr(0, p); + auto value = d.substr(p + 1); + LUISA_VERBOSE_WITH_LOCATION("Parameter definition: {} = '{}'", key, value); + if (auto iter = macros.find(key); iter != macros.end()) { + LUISA_WARNING_WITH_LOCATION( + "Duplicate definition: {} = '{}'. " + "Ignoring the previous one: {} = '{}'.", + key, value, key, iter->second); + iter->second = value; + } else { + macros.emplace(key, value); + } + } + }; + // parse all options starting with '-D' or '--define' + for (int i = 1; i < argc; i++) { + auto arg = luisa::string_view{argv[i]}; + if (arg == "-D" || arg == "--define") { + if (i + 1 == argc) { + LUISA_WARNING_WITH_LOCATION( + "Missing definition after {}.", arg); + // remove the option + argv[i] = nullptr; + } else { + parse_macro(argv[i + 1]); + // remove the option and its argument + argv[i] = nullptr; + argv[++i] = nullptr; + } + } else if (arg.starts_with("-D")) { + parse_macro(arg.substr(2)); + // remove the option + argv[i] = nullptr; + } + } + // remove all nullptrs + auto new_end = std::remove(argv, argv + argc, nullptr); + argc = static_cast(new_end - argv); + return macros; +} + +int main(int argc, char *argv[]) { + + log_level_info(); + luisa::log_level_verbose(); + luisa::compute::Context context{argv[0]}; + auto macros = parse_cli_macros(argc, argv); + for (auto &&[k, v] : macros) { + LUISA_INFO("Found CLI Macro: {} = {}", k, v); + } + + auto options = parse_cli_options(argc, argv); + auto backend = options["backend"].as(); + auto index = options["device"].as(); + auto path = options["scene"].as(); + compute::DeviceConfig config; + config.device_index = index; + config.inqueue_buffer_limit = false;// Do not limit the number of in-queue buffers --- we are doing offline rendering! + auto device = context.create_device(backend, &config); + + Clock clock; + auto scene_desc = SceneParser::parse(path, macros); + auto parse_time = clock.toc(); + + LUISA_INFO("Parsed scene description file '{}' in {} ms.", + path.string(), parse_time); + auto scene = Scene::create(context, scene_desc.get()); + auto stream = device.create_stream(StreamTag::GRAPHICS); + auto pipeline = Pipeline::create(device, stream, *scene); + + Kernel1D ad_kernel = [] () { + $autodiff { + Float3 point_pre = make_float3(-0.021873882, 1.9799696, -0.018273553); + Float3 point_cur = make_float3(0.5625003, 1, -0.52317667); + Float3 point_nxt = make_float3(0.8394947, 0, -0.7626182); + Float3 normal_cur = make_float3(0.,1.,0.); + device_log("point_pre {} point_cur {} point_nxt {} normal_cur {}",point_pre, point_cur, point_nxt, normal_cur); + requires_grad(point_pre, point_cur, point_nxt, normal_cur); + auto wi = normalize(point_pre-point_cur); + auto wo = normalize(point_nxt-point_cur); + auto s = normalize(make_float3(0.0f, -normal_cur[2], normal_cur[1])); + // auto ss = normalize(s - n * dot(n, s)); + // auto tt = normalize(cross(n, ss)); + auto f = Frame::make(normal_cur, s); + auto wi_local = f.world_to_local(wi); + auto wo_local = f.world_to_local(wo); + backward(wi); + auto res = normalize(wi_local+wo_local*1.8f); + device_log("res is {} wi {} wo {} eta {} normal {}",res, wi_local, wo_local, 1.8, normal_cur); + device_log("grad(normal_cur) is {}",grad(normal_cur)); + device_log("grad(point_pre) is {}",grad(point_pre)); + device_log("grad(point_cur) is {}",grad(point_cur)); + device_log("grad(point_nxt) is {}",grad(point_nxt)); + }; + }; + stream << device.compile(ad_kernel)().dispatch(1u) << synchronize(); +} diff --git a/src/tests/test_ad_torch.py b/src/tests/test_ad_torch.py index c35b0fad..e3b5c5c2 100644 --- a/src/tests/test_ad_torch.py +++ b/src/tests/test_ad_torch.py @@ -33,29 +33,124 @@ def cu_device_ptr_to_torch_tensor(ptr, shape, dtype=cupy.float32): # Convert the CuPy ndarray to a DLPack tensor and then to a PyTorch tensor return torch.utils.dlpack.from_dlpack(array.toDlpack()) + +def compute_vertex_normals(vertex_pos, face_ids): + v0 = vertex_pos[face_ids[:, 0]] + v1 = vertex_pos[face_ids[:, 1]] + v2 = vertex_pos[face_ids[:, 2]] + face_normals = torch.cross(v1 - v0, v2 - v0) + vertex_normals = torch.zeros_like(vertex_pos) + vertex_normals[face_ids[:, 0]] += face_normals + vertex_normals[face_ids[:, 1]] += face_normals + vertex_normals[face_ids[:, 2]] += face_normals + vertex_normals_norm = vertex_normals/torch.norm(vertex_normals, dim=-1, keepdim=True) + return vertex_normals_norm + + +gt_args = ["C:/Users/jiankai/anaconda3/Lib/site-packages/luisarender/dylibs","-b","cuda", "D:/Code/LuisaRender2/data/scenes/cbox_caustic.luisa"] +init_args = ["C:/Users/jiankai/anaconda3/Lib/site-packages/luisarender/dylibs","-b","cuda", "D:/Code/LuisaRender2/data/scenes/cbox_caustic.luisa"] + +luisarender.load_scene(gt_args) +x = luisarender.ParamStruct() +x.type = 'geom' +x.id = 0 +[geom_ptr,geom_size,face_ids,face_sizes] = luisarender.get_scene_param([x]) +geom_ptr_torch = cu_device_ptr_to_torch_tensor(geom_ptr[0], (geom_size[0]//8,8), dtype=cupy.float32) +face_ids_torch= cu_device_ptr_to_torch_tensor(face_ids[0], (face_sizes[0]//3,3), dtype=cupy.int32) +vertex_pos = geom_ptr_torch.clone()[...,:3] +vertex = geom_ptr_torch.clone() + +# print(face_ids_torch,vertex) +# exit() + +vertex_height = torch.ones((vertex_pos.shape[0]),dtype=torch.float32).cuda() +optimizer = torch.optim.Adam([vertex_height], lr=0.001) +vertex_height.requires_grad_() + + +vertex_pos[...,1] = ((vertex_pos[...,1]-1)*0.05)+1 +gt_param = vertex_pos[...,1].clone() +vertex_normal = compute_vertex_normals(vertex_pos, face_ids_torch) +vertex[...,0:3] = vertex_pos +vertex[...,3:6] = vertex_normal + +pos_ptr = vertex.contiguous().data_ptr() +pos_size = np.prod(vertex.shape) +pos_dtype=float + +x.size = pos_size +x.buffer_ptr = pos_ptr + +torch.cuda.synchronize() +luisarender.update_scene([x]) +target_img = cu_device_ptr_to_torch_tensor(luisarender.render()[0], (512, 512, 4)).clone() +imageio.imwrite("gt_0.05.exr",target_img.detach().cpu().numpy()[...,:3]) +imageio.imwrite("gt.png",target_img.detach().cpu().numpy()[...,:3]) + +#print(torch.max(target_img), torch.min(target_img), torch.sum(target_img)) + + + + +torch.cuda.synchronize() +vertex_pos[...,1] = vertex_height +vertex_normal = compute_vertex_normals(vertex_pos, face_ids_torch) +vertex[...,0:3] = vertex_pos +vertex[...,3:6] = vertex_normal +luisarender.update_scene([x]) + +render_img = cu_device_ptr_to_torch_tensor(luisarender.render()[0], (512, 512,4)).clone() +imageio.imwrite("init.exr",render_img.detach().cpu().numpy()[...,:3]) +imageio.imwrite("init.png",render_img.detach().cpu().numpy()[...,:3]) + +#exit() +cm = plt.get_cmap('viridis') +# Apply the colormap like a function to any array: + +for i in range(500): + + vertex_pos = geom_ptr_torch.clone()[...,:3] + vertex_pos[...,1] = vertex_height + vertex_normal = compute_vertex_normals(vertex_pos, face_ids_torch) + vertex[...,0:3] = vertex_pos + vertex[...,3:6] = vertex_normal + torch.cuda.synchronize() + luisarender.update_scene([x]) + + render_img = cu_device_ptr_to_torch_tensor(luisarender.render()[0], (512, 512, 4)).clone() + imageio.imwrite(f"outputs/render{i}.exr",render_img.detach().cpu().numpy()[...,:3]) + render_img.requires_grad_() + #loss = loss_func(render_img,target_img) + loss = torch.sum((render_img[400:,...]-target_img[400:,...])**2) + loss.backward() + grad = render_img.grad[...,:3] + #print("debug pixel", render_img[458, 210,:3], grad[458, 210,:3]) + #exit() + aux_buffer = luisarender.render_backward([grad.contiguous().data_ptr()],[np.prod(grad.shape)]) + aux_buffer_torch = cu_device_ptr_to_torch_tensor(aux_buffer[0], (512, 512, 4), dtype=cupy.float32) + aux_buffer_numpy = aux_buffer_torch.cpu().numpy() + + imageio.imwrite(f"outputs/backwarde_render_{i}.exr",aux_buffer_numpy[...,:3]) + + graddis_avg = aux_buffer_numpy[...,0] + mx = np.max(abs(graddis_avg)) + print(mx, np.max(graddis_avg), np.min(graddis_avg)) + grad_reldis_vis = cm(graddis_avg/mx) + imageio.imwrite(f"outputs/grad_vis_{i}.png",grad_reldis_vis) + + exit() + + tex_grad, geom_grad = luisarender.get_gradients() + geom_grad_torch = cu_device_ptr_to_torch_tensor(geom_grad[0], vertex.shape, dtype=cupy.float32) + #print(loss, geom_grad_torch) + + vertex_height.grad = geom_grad_torch[...,1] + #vertex_normal.grad = geom_grad_torch[...,3:6] + vertex_normal.backward(geom_grad_torch[...,3:6]) + optimizer.step() + print(loss, torch.sum((vertex_height-gt_param)**2)) + #print(vertex_height) -# def torch_to_lc_buffer(tensor): -# assert tensor.dtype is torch.float32 # TODO -# size = np.prod(tensor.shape) -# buf = luisa.Buffer.import_external_memory( -# tensor.contiguous().data_ptr(), -# size, dtype=float) -# return buf - -# def lc_buffer_to_torch(buf): -# assert buf.dtype is float # TODO -# shape = (buf.size,) -# return cu_device_ptr_to_torch_tensor(buf.native_handle, shape) - -def is_torch_tensor(a): - return getattr(a, '__module__', None) == 'torch' \ - and type(a).__name__ == 'Tensor' - -def torch_ensure_grad_shape(a, b): - if is_torch_tensor(a) and a.dtype in [torch.float, torch.float32, torch.float64]: - return a.reshape(b.shape) - else: - return a # def torch_to_luisa_scene(args): # return tuple(torch_to_lc_buffer(a) if is_torch_tensor(a) else a for a in args) @@ -86,61 +181,6 @@ def torch_ensure_grad_shape(a, b): # uint param_size; # uint64_t param_buffer_ptr; # float4 param_value; - - -gt_args = ["C:/Users/jiankai/anaconda3/Lib/site-packages/luisarender/dylibs","-b","cuda", "D:/Code/LuisaRender2/data/scenes/cbox_caustic.luisa"] -init_args = ["C:/Users/jiankai/anaconda3/Lib/site-packages/luisarender/dylibs","-b","cuda", "D:/Code/LuisaRender2/data/scenes/cbox_caustic.luisa"] - -differentiable_params_list = [ - {"type":"mesh","idx":0,"param":"vertex_position"}, - #{"type":"texture","idx":0,"param":"base_color"} -] - -luisarender.load_scene(gt_args) -target_img = cu_device_ptr_to_torch_tensor(luisarender.render()[0], (512, 512, 4)).clone() -imageio.imwrite("gt.exr",target_img.detach().cpu().numpy()[...,:3]) - -#print(torch.max(target_img), torch.min(target_img), torch.sum(target_img)) - -x = luisarender.ParamStruct() -x.type = 'geom' -x.id = 0 -[geom_ptr,geom_size] = luisarender.get_scene_param([x]) -geom_ptr_torch = cu_device_ptr_to_torch_tensor(geom_ptr[0], (geom_size[0]//8,8), dtype=cupy.float32) -vertex_pos = geom_ptr_torch.clone() -vertex_pos[...,1]=1.0 -vertex_pos[...,3]=0.0 -vertex_pos[...,4]=1.0 -vertex_pos[...,5]=0.0 -pos_ptr = vertex_pos.contiguous().data_ptr() -pos_size = np.prod(vertex_pos.shape) -pos_dtype=float - -optimizer = torch.optim.Adam([vertex_pos], lr=0.001) -x.size = pos_size -x.buffer_ptr = pos_ptr -luisarender.update_scene([x]) - -render_img = cu_device_ptr_to_torch_tensor(luisarender.render()[0], (512, 512,4)).clone() -imageio.imwrite("init.exr",render_img.detach().cpu().numpy()[...,:3]) - -loss_func = torch.nn.MSELoss() - -for i in range(500): - render_img = cu_device_ptr_to_torch_tensor(luisarender.render()[0], (512, 512, 4)).clone() - imageio.imwrite(f"render{i}.exr",render_img.detach().cpu().numpy()[...,:3]) - render_img.requires_grad_() - #loss = loss_func(render_img,target_img) - loss = torch.sum((render_img-target_img)**2) - loss.backward() - grad = render_img.grad[...,:3] - luisarender.render_backward([grad.contiguous().data_ptr()],[np.prod(grad.shape)]) - tex_grad, geom_grad = luisarender.get_gradients() - geom_grad_torch = cu_device_ptr_to_torch_tensor(geom_grad[0], vertex_pos.shape, dtype=cupy.float32) - print(loss, torch.max(geom_grad_torch), torch.min(geom_grad_torch), geom_grad_torch.shape) - #exit() - luisarender.update_scene([x]) - #exit() #optimizer.zero_grad() #tex.grad = tex_grad_torch diff --git a/src/util/frame.cpp b/src/util/frame.cpp index 81091f93..60702193 100644 --- a/src/util/frame.cpp +++ b/src/util/frame.cpp @@ -33,6 +33,7 @@ Frame Frame::make(Expr n, Expr s) noexcept { return {ss, tt, n}; } + Float3 Frame::local_to_world(Expr d) const noexcept { return normalize(d.x * _s + d.y * _t + d.z * _n); }