| 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()) |