Skip to content

Commit

Permalink
update EPSM_path
Browse files Browse the repository at this point in the history
jkxing committed May 1, 2024
1 parent 2d69442 commit 27ec951
Showing 20 changed files with 1,601 additions and 600 deletions.
1 change: 1 addition & 0 deletions src/base/camera.cpp
Original file line number Diff line number Diff line change
@@ -221,6 +221,7 @@ Camera::Sample Camera::Instance::generate_ray(Expr<uint2> pixel_coord, Expr<floa
weight *= filter_weight;
auto c2w = camera_to_world();
auto o = make_float3(c2w * make_float4(ray->origin(), 1.f));
auto pixel_world = make_float3(c2w * make_float4(ray->direction() / ray->direction().z, 1.f));
auto d = normalize(make_float3x3(c2w) * ray->direction());
ray->set_origin(o);
ray->set_direction(d);
42 changes: 32 additions & 10 deletions src/base/differentiation.cpp
Original file line number Diff line number Diff line change
@@ -405,7 +405,8 @@ void Differentiation::add_geom_gradients(Float3 grad_v, Float3 grad_n, Float3 we
_grad_buffer.value()->atomic(real_grad_offset + triangle.i1 * 8 + 3).fetch_add(grad_n[0] * weight[1]);
_grad_buffer.value()->atomic(real_grad_offset + triangle.i1 * 8 + 4).fetch_add(grad_n[1] * weight[1]);
_grad_buffer.value()->atomic(real_grad_offset + triangle.i1 * 8 + 5).fetch_add(grad_n[2] * weight[1]);


//device_log("weight2 {}, grad_v {}, grad_offset {} triangle i2 {}",weight[2], grad_v, real_grad_offset, triangle.i2);
_grad_buffer.value()->atomic(real_grad_offset + triangle.i2 * 8 + 0).fetch_add(grad_v[0] * weight[2]);
_grad_buffer.value()->atomic(real_grad_offset + triangle.i2 * 8 + 1).fetch_add(grad_v[1] * weight[2]);
_grad_buffer.value()->atomic(real_grad_offset + triangle.i2 * 8 + 2).fetch_add(grad_v[2] * weight[2]);
@@ -463,25 +464,40 @@ void Differentiation::register_optimizer(Optimizer::Instance *optimizer) noexcep

void Differentiation::register_geometry_parameter(CommandBuffer &command_buffer, Geometry::MeshData& mesh, Accel& accel, uint instance_id) noexcept {

LUISA_INFO("register geometry parameter start");
//LUISA_INFO("register geometry parameter start");
auto param_offset = _param_buffer_size;
auto counter_offset = _counter_size;
auto grad_offset = _gradient_buffer_size;
auto param_index = static_cast<uint>(_geometry_params.size());
//auto channels = 3u;
//uint buffer_id = mesh.geometry_buffer_id_base;
auto buffer_view = mesh.vertices;

auto tri_buffer_view = mesh.triangles;
auto triangle_buffer_new = _pipeline.create<Buffer<int>>(tri_buffer_view.size() * 3);
auto triangle_buffer_new_view = triangle_buffer_new->view();

Kernel1D write_tri = [&]() noexcept {
auto x = dispatch_x();
auto tri = tri_buffer_view->read(x);
triangle_buffer_new_view->write(x * 3, (Int)(tri.i0));
triangle_buffer_new_view->write(x * 3 + 1, (Int)(tri.i1));
triangle_buffer_new_view->write(x * 3 + 2, (Int)(tri.i2));
};
auto write_tri_shader = pipeline().device().compile(write_tri);
command_buffer << write_tri_shader().dispatch(tri_buffer_view.size()) << synchronize();

auto length = buffer_view.size() * 8;
_counter_size = (_counter_size + length + 8u) & ~0b11u;
_param_buffer_size = (_param_buffer_size + length + 8u) & ~0b11u;
_gradient_buffer_size = (_gradient_buffer_size + length + 8u) & ~0b11u;
_geometry_params.emplace_back(param_index, instance_id, grad_offset, param_offset, counter_offset, buffer_view, length, 0u, mesh.resource);
_geometry_params.emplace_back(param_index, instance_id, grad_offset, param_offset, counter_offset, buffer_view, length, 0u, mesh.resource, triangle_buffer_new_view);
Kernel1D write_i2o = [&](UInt instance_id, UInt grad_offset) noexcept {
instance2offset.value()->write(instance_id, grad_offset+1u);
};
auto write_i2o_shader = pipeline().device().compile(write_i2o);
command_buffer << write_i2o_shader(instance_id, grad_offset).dispatch(1u) << synchronize();
LUISA_INFO("register geometry parameter finished");
//LUISA_INFO("register geometry parameter finished");
}

void Differentiation::update_parameter_from_external(Stream &stream, luisa::vector<uint> &constants_id, luisa::vector<float4> &constants, luisa::vector<uint> &textures_id,
@@ -495,33 +511,39 @@ void Differentiation::update_parameter_from_external(Stream &stream, luisa::vect
stream << image.copy_from(textures[textures_id[i]]);
}

LUISA_INFO("working on update paramaters {}",geoms_id.size());
//LUISA_INFO("working on update paramaters {}",geoms_id.size());
// apply geometry parameters
for (auto i=0;i<geoms_id.size();i++) {
auto &&p = _geometry_params[geoms_id[i]];
auto param_offset = p.param_offset();
auto buffer_view = p.buffer();
//auto [buffer_view, bindlessbuffer_id] = _pipeline.bindless_arena_buffer<Vertex>(buffer_id);
auto length = buffer_view.size();
LUISA_INFO("here length is {}, size is {}, as<vertex> size is {}",length,geoms[geoms_id[i]].view().size(),geoms[geoms_id[i]].view().as<Vertex>().size());
//LUISA_INFO("here length is {}, size is {}, as<vertex> size is {}",length,geoms[geoms_id[i]].view().size(),geoms[geoms_id[i]].view().as<Vertex>().size());
stream << buffer_view.copy_from(geoms[geoms_id[i]].view().as<Vertex>()) << synchronize() << p.mesh()->build() << synchronize();
_is_dirty = true;
}
}

std::tuple<luisa::vector<void *>, luisa::vector<uint>> Differentiation::get_parameter_from_external(Stream &stream, luisa::vector<uint> &constants_id, luisa::vector<uint> &textures_id, luisa::vector<uint> &geoms_id) noexcept {
std::tuple<luisa::vector<void *>, luisa::vector<uint>, luisa::vector<void *>, luisa::vector<uint>> Differentiation::get_parameter_from_external(Stream &stream, luisa::vector<uint> &constants_id, luisa::vector<uint> &textures_id, luisa::vector<uint> &geoms_id) noexcept {
luisa::vector<void*> geom_param{};
luisa::vector<uint> geom_size{};
luisa::vector<void*> tri_param{};
luisa::vector<uint> tri_size{};
// apply geometry parameters
for (auto i: geoms_id) {
auto p = _geometry_params[i];
auto param_offset = p.param_offset();
auto buffer_view = p.buffer();
auto tri_buffer_view = p.tri_buffer();
auto length = buffer_view.size();
auto tri_length = tri_buffer_view.size();
geom_param.push_back(buffer_view.as<float>().native_handle());
geom_size.push_back(length*8u);
tri_param.push_back(tri_buffer_view.native_handle());
tri_size.push_back(tri_length);
}
return std::make_tuple(geom_param, geom_size);
return std::make_tuple(geom_param, geom_size, tri_param, tri_size);
}

std::tuple<luisa::vector<void *>, luisa::vector<void *>> Differentiation::get_gradients(Stream &stream) {
@@ -552,8 +574,8 @@ std::tuple<luisa::vector<void *>, luisa::vector<void *>> Differentiation::get_gr
auto grad_offset = p.gradient_buffer_offset();
auto buffer_view = p.buffer();
auto length = buffer_view.size();
LUISA_INFO("here length is {}",length);
auto geom_grad_buf_view = _grad_buffer->subview(grad_offset, length);
//LUISA_INFO("here length is {}",length);
auto geom_grad_buf_view = _grad_buffer->subview(grad_offset, length*8);
geom_res.push_back(reinterpret_cast<void*>(reinterpret_cast<uint64_t>(geom_grad_buf_view.native_handle())+geom_grad_buf_view.offset_bytes()));
}
return std::make_tuple(texture_res, geom_res);
8 changes: 5 additions & 3 deletions src/base/differentiation.h
Original file line number Diff line number Diff line change
@@ -66,16 +66,18 @@ class Differentiation {
uint _counter_offset;
float2 _range;
BufferView<Vertex> _buffer_view;
BufferView<int> _tri_buffer_view;
uint _length;
uint _buffer_id;
Mesh *_mesh;
public:
GeometryParameter(uint index, uint instance_id, uint grad_offset, uint param_offset,
uint counter_offset, BufferView<Vertex> buffer_view, uint length, uint buffer_id, Mesh *mesh) noexcept
uint counter_offset, BufferView<Vertex> buffer_view, uint length, uint buffer_id, Mesh *mesh, BufferView<int> tri_buffer_view) noexcept
: _index(index), _instance_id{instance_id}, _grad_offset{grad_offset}, _param_offset{param_offset},
_counter_offset{counter_offset}, _buffer_view{buffer_view}, _length(length), _buffer_id(buffer_id), _mesh(mesh) {}
_counter_offset{counter_offset}, _buffer_view{buffer_view}, _length(length), _buffer_id(buffer_id), _mesh(mesh), _tri_buffer_view(tri_buffer_view) {}
[[nodiscard]] auto index() const noexcept { return _index; }
[[nodiscard]] auto buffer() const noexcept { return _buffer_view; }
[[nodiscard]] auto tri_buffer() const noexcept { return _tri_buffer_view; }
[[nodiscard]] auto buffer_id() const noexcept { return _buffer_id; }
[[nodiscard]] auto instance_id() const noexcept { return _instance_id; }
[[nodiscard]] auto gradient_buffer_offset() const noexcept { return _grad_offset; }
@@ -176,7 +178,7 @@ class Differentiation {
void update_parameter_from_external(Stream &stream, luisa::vector<uint> &constants_id, luisa::vector<float4> &constants, luisa::vector<uint> &textures_id,
luisa::vector<Buffer<float4>> &textures, luisa::vector<uint> &geoms_id, luisa::vector<Buffer<float>> &geoms) noexcept;

std::tuple<luisa::vector<void *>, luisa::vector<uint>> get_parameter_from_external
std::tuple<luisa::vector<void *>, luisa::vector<uint>,luisa::vector<void *>, luisa::vector<uint>> get_parameter_from_external
(Stream &stream, luisa::vector<uint> &constants_id, luisa::vector<uint> &textures_id, luisa::vector<uint> &geoms_id) noexcept;


1 change: 1 addition & 0 deletions src/base/film.h
Original file line number Diff line number Diff line change
@@ -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;
};

7 changes: 4 additions & 3 deletions src/base/geometry.cpp
Original file line number Diff line number Diff line change
@@ -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;
2 changes: 2 additions & 0 deletions src/base/geometry.h
Original file line number Diff line number Diff line change
@@ -47,11 +47,13 @@ class Geometry {
Mesh *resource;
uint buffer_id_base;
BufferView<Vertex> vertices;
BufferView<Triangle> triangles;
};

struct MeshData {
Mesh *resource;
BufferView<Vertex> vertices;
BufferView<Triangle> triangles;
uint16_t shadow_term;
uint16_t intersection_offset;
uint geometry_buffer_id_base : 22;
7 changes: 5 additions & 2 deletions src/base/integrator.cpp
Original file line number Diff line number Diff line change
@@ -169,17 +169,20 @@ DifferentiableIntegrator::Instance::Instance(

DifferentiableIntegrator::Instance::~Instance() noexcept = default;

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

void DifferentiableIntegrator::Instance::_render_one_camera_backward(
5 changes: 3 additions & 2 deletions src/base/integrator.h
Original file line number Diff line number Diff line change
@@ -44,8 +44,9 @@ class Integrator : public SceneNode {
[[nodiscard]] auto light_sampler() noexcept { return _light_sampler.get(); }
[[nodiscard]] auto light_sampler() const noexcept { return _light_sampler.get(); }
virtual void render(Stream &stream) noexcept = 0;
virtual void render_backward(Stream &stream, luisa::vector<Buffer<float>> &grad_in) {
virtual luisa::vector<void *> render_backward(Stream &stream, luisa::vector<Buffer<float>> &grad_in) {
LUISA_INFO("Not implemented!");
return luisa::vector<void *>{};
};
virtual luisa::vector<void*> render_with_return(Stream &stream) {
LUISA_INFO("Not implemented!");
@@ -121,7 +122,7 @@ class DifferentiableIntegrator : public ProgressiveIntegrator {
virtual void _render_one_camera_backward(CommandBuffer &command_buffer, uint iteration, Camera::Instance *camera, Buffer<float> & grad_in) noexcept;

public:
void render_backward(Stream &stream, luisa::vector<Buffer<float>> &grad_in) noexcept override;
luisa::vector<void *> render_backward(Stream &stream, luisa::vector<Buffer<float>> &grad_in) noexcept override;
Instance(Pipeline &pipeline, CommandBuffer &command_buffer,
const DifferentiableIntegrator *integrator) noexcept;
~Instance() noexcept override;
4 changes: 2 additions & 2 deletions src/base/pipeline.cpp
Original file line number Diff line number Diff line change
@@ -148,8 +148,8 @@ void Pipeline::render(Stream &stream) noexcept {
_integrator->render(stream);
}

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

luisa::vector<void*> Pipeline::render_with_return(Stream &stream) noexcept {
2 changes: 1 addition & 1 deletion src/base/pipeline.h
Original file line number Diff line number Diff line change
@@ -218,7 +218,7 @@ class Pipeline {
void update_texture(Stream &stream, uint texture_id, float4 new_value) noexcept;
void update_mesh(uint mesh_id, uint64_t vertex_buffer) noexcept;
void render(Stream &stream) noexcept;
void render_diff(Stream &stream, luisa::vector<Buffer<float>> &grads) noexcept;
luisa::vector<void *> render_diff(Stream &stream, luisa::vector<Buffer<float>> &grads) noexcept;
luisa::vector<void*> render_with_return(Stream &stream) noexcept;
[[nodiscard]] auto &printer() noexcept { return *_printer; }
[[nodiscard]] auto &printer() const noexcept { return *_printer; }
2 changes: 1 addition & 1 deletion src/compute
Submodule compute updated 465 files
4 changes: 4 additions & 0 deletions src/films/color.cpp
Original file line number Diff line number Diff line change
@@ -73,6 +73,10 @@ class ColorFilmInstance final : public Film::Instance {
command_buffer << compute::synchronize();
return _converted.native_handle();
}
void* export_image_origin(CommandBuffer &command_bufferr) noexcept override {
_check_prepared();
return _image.native_handle();
}
void download(CommandBuffer &command_buffer, float4 *framebuffer) const noexcept override;
[[nodiscard]] Film::Accumulation read(Expr<uint2> pixel) const noexcept override;
void release() noexcept override;
712 changes: 712 additions & 0 deletions src/integrators/bkp2.cpp

Large diffs are not rendered by default.

983 changes: 500 additions & 483 deletions src/integrators/megappm_diff.cpp

Large diffs are not rendered by default.

35 changes: 23 additions & 12 deletions src/python/lrapi.cpp
Original file line number Diff line number Diff line change
@@ -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<uint64_t> res_vec(res.size());
for (int i = 0; i < res.size(); i++) {
res_vec[i] = reinterpret_cast<uint64_t>(res[i]);
}
LUISA_INFO("res_vec: {}",res_vec[0]);
//LUISA_INFO("res_vec: {}",res_vec[0]);
return res_vec;
});

@@ -225,15 +225,15 @@ PYBIND11_MODULE(_lrapi, m) {


m.def("update_scene", [](std::vector<ParamStruct> params) {
LUISA_INFO("LuisaRender Update Scene");
//LUISA_INFO("LuisaRender Update Scene");
luisa::vector<float4> constants{};
luisa::vector<Buffer<float4>> textures{};
luisa::vector<Buffer<float>> geoms{};
luisa::vector<uint> constants_id{};
luisa::vector<uint> textures_id{};
luisa::vector<uint> 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<ParamStruct> params) {
LUISA_INFO("LuisaRender API get_parameter start");
//LUISA_INFO("LuisaRender API get_parameter start");
luisa::vector<uint> constants_id{};
luisa::vector<uint> textures_id{};
luisa::vector<uint> 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<float> ret_con_param(constants_id.size());
// std::vector<uint64_t> ret_tex_param(textures_id.size());
// std::vector<uint> ret_tex_size(textures_id.size());
std::vector<uint64_t> ret_geom_param(geoms_id.size());
std::vector<uint> ret_geom_size(geoms_id.size());
std::vector<uint64_t> ret_tri_param(geoms_id.size());
std::vector<uint> 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<uint64_t>(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<uint64_t>(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<uint64_t> grad_ptr,std::vector<uint> sizes){
LUISA_INFO("LuisaRender API render_backward");

//LUISA_INFO("LuisaRender API render_backward");
//scene_python._pipeline->differentiation()->clear_gradients(*scene_python._stream);
luisa::vector<Buffer<float>> grad_buffer{grad_ptr.size()};
for (int i = 0; i < grad_ptr.size(); i++) {
auto buffer = scene_python._pipeline->device().import_external_buffer<float>(reinterpret_cast<void *>(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<uint64_t> res_vec(res.size());
for (int i = 0; i < res.size(); i++) {
res_vec[i] = reinterpret_cast<uint64_t>(res[i]);
}
return res_vec;
});

m.def("get_gradients", [](){
7 changes: 3 additions & 4 deletions src/shapes/mesh.cpp
Original file line number Diff line number Diff line change
@@ -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 |
3 changes: 3 additions & 0 deletions src/tests/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -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)
181 changes: 181 additions & 0 deletions src/tests/test_ad.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,181 @@

// Created by Mike on 2021/12/7.
//

#include <span>
#include <iostream>

#include <cxxopts.hpp>

#include <core/stl/format.h>
#include <sdl/scene_desc.h>
#include <sdl/scene_parser.h>
#include <base/scene.h>
#include <base/pipeline.h>

#include <random>
#include <luisa-compute.h>
#include <util/sampling.h>
#include <util/medium_tracker.h>
#include <util/progress_bar.h>
#include <base/pipeline.h>
#include <base/integrator.h>
#include <sdl/scene_node_desc.h>
#include <limits>

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<luisa::string>(), "<backend>");
cli.add_option("", "d", "device", "Compute device index", cxxopts::value<int32_t>()->default_value("-1"), "<index>");
cli.add_option("", "", "scene", "Path to scene description file", cxxopts::value<std::filesystem::path>(), "<file>");
cli.add_option("", "D", "define", "Parameter definitions to override scene description macros.",
cxxopts::value<std::vector<luisa::string>>()->default_value("<none>"), "<key>=<value>");
cli.add_option("", "h", "help", "Display this help message", cxxopts::value<bool>()->default_value("false"), "");
cli.allow_unrecognised_options();
cli.positional_help("<file>");
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<bool>()) {
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 = [&macros](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<int>(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<luisa::string>();
auto index = options["device"].as<int32_t>();
auto path = options["scene"].as<std::filesystem::path>();
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();
}
194 changes: 117 additions & 77 deletions src/tests/test_ad_torch.py
Original file line number Diff line number Diff line change
@@ -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
1 change: 1 addition & 0 deletions src/util/frame.cpp
Original file line number Diff line number Diff line change
@@ -33,6 +33,7 @@ Frame Frame::make(Expr<float3> n, Expr<float3> s) noexcept {
return {ss, tt, n};
}


Float3 Frame::local_to_world(Expr<float3> d) const noexcept {
return normalize(d.x * _s + d.y * _t + d.z * _n);
}

0 comments on commit 27ec951

Please sign in to comment.