| 23 | |
| 24 | |
| 25 | class Splat2D: |
| 26 | def __init__(self, device: falcor.Device): |
| 27 | self.device = device |
| 28 | |
| 29 | self.blobs_buf = device.create_structured_buffer( |
| 30 | struct_size=32, |
| 31 | element_count=BLOB_COUNT, |
| 32 | bind_flags=falcor.ResourceBindFlags.ShaderResource |
| 33 | | falcor.ResourceBindFlags.UnorderedAccess |
| 34 | | falcor.ResourceBindFlags.Shared, |
| 35 | ) |
| 36 | |
| 37 | self.grad_blobs_buf = device.create_structured_buffer( |
| 38 | struct_size=32, |
| 39 | element_count=BLOB_COUNT, |
| 40 | bind_flags=falcor.ResourceBindFlags.ShaderResource |
| 41 | | falcor.ResourceBindFlags.UnorderedAccess |
| 42 | | falcor.ResourceBindFlags.Shared, |
| 43 | ) |
| 44 | |
| 45 | self.image_buf = device.create_structured_buffer( |
| 46 | struct_size=12, |
| 47 | element_count=RESOLUTION * RESOLUTION, |
| 48 | bind_flags=falcor.ResourceBindFlags.ShaderResource |
| 49 | | falcor.ResourceBindFlags.UnorderedAccess |
| 50 | | falcor.ResourceBindFlags.Shared, |
| 51 | ) |
| 52 | |
| 53 | self.grad_image_buf = device.create_structured_buffer( |
| 54 | struct_size=12, |
| 55 | element_count=RESOLUTION * RESOLUTION, |
| 56 | bind_flags=falcor.ResourceBindFlags.ShaderResource |
| 57 | | falcor.ResourceBindFlags.UnorderedAccess |
| 58 | | falcor.ResourceBindFlags.Shared, |
| 59 | ) |
| 60 | |
| 61 | self.forward_pass = falcor.ComputePass( |
| 62 | device, file=DIR / "splat2d.cs.slang", cs_entry="forward_main" |
| 63 | ) |
| 64 | |
| 65 | self.backward_pass = falcor.ComputePass( |
| 66 | device, file=DIR / "splat2d.cs.slang", cs_entry="backward_main" |
| 67 | ) |
| 68 | |
| 69 | def forward(self, blobs): |
| 70 | self.blobs_buf.from_torch(blobs.detach()) |
| 71 | self.device.render_context.wait_for_cuda() |
| 72 | vars = self.forward_pass.globals.forward |
| 73 | vars.blobs = self.blobs_buf |
| 74 | vars.blob_count = BLOB_COUNT |
| 75 | vars.image = self.image_buf |
| 76 | vars.resolution = falcor.uint2(RESOLUTION, RESOLUTION) |
| 77 | self.forward_pass.execute(threads_x=RESOLUTION, threads_y=RESOLUTION) |
| 78 | self.device.render_context.wait_for_falcor() |
| 79 | return self.image_buf.to_torch([RESOLUTION, RESOLUTION, 3], falcor.float32) |
| 80 | |
| 81 | def backward(self, blobs, grad_intensities): |
| 82 | self.grad_blobs_buf.from_torch(torch.zeros([BLOB_COUNT, 8]).cuda()) |