1 // =============================================================================
2 // PROJECT CHRONO - http://projectchrono.org
3 //
4 // Copyright (c) 2019 projectchrono.org
5 // All rights reserved.
6 //
7 // Use of this source code is governed by a BSD-style license that can be found
8 // in the LICENSE file at the top level of the distribution and at
9 // http://projectchrono.org/license-chrono.txt.
10 //
11 // =============================================================================
12 // Authors: Asher Elmquist
13 // =============================================================================
14 //
15 //
16 // =============================================================================
17 
18 #ifndef CHOPTIXPIPELINE_H
19 #define CHOPTIXPIPELINE_H
20 
21 #include <optix.h>
22 #include <cuda.h>
23 #include <cuda_runtime.h>
24 #include <cuda_runtime_api.h>
25 
26 #include "chrono_sensor/optix/ChOptixDefinitions.h"
27 #include "chrono/assets/ChVisualMaterial.h"
28 #include "chrono/assets/ChTriangleMeshShape.h"
29 #include "chrono_sensor/optix/scene/ChScene.h"
30 
31 #include "chrono_sensor/ChApiSensor.h"
32 
33 #include <vector>
34 #include <memory>
35 
36 namespace chrono {
37 namespace sensor {
38 
39 /// @addtogroup sensor_optix
40 /// @{
41 
42 template <typename T>
43 struct Record {
44     __align__(OPTIX_SBT_RECORD_ALIGNMENT) char header[OPTIX_SBT_RECORD_HEADER_SIZE];
45     T data;
46 };
47 
48 /// The type of ray tracing used to model the sensor
49 enum class PipelineType {
50     CAMERA_PINHOLE,         ///< pinhole camera model
51     CAMERA_FOV_LENS,        ///< FOV lens model
52     SEGMENTATION_PINHOLE,   ///< pinhole segmentation camera
53     SEGMENTATION_FOV_LENS,  ///< FOV lens segmentation camera
54     LIDAR_SINGLE,           ///< single sample lidar
55     LIDAR_MULTI,            ///< multi sample lidar
56     RADAR                   ///< radar model
57 };
58 // TODO: how do we allow custom ray gen programs? (Is that ever going to be a thing?)
59 
60 /// Class to hold all the Shader Binding Table parameters adnd manage the ray tracing pipeline, materials, ray gen
61 /// programs
62 class CH_SENSOR_API ChOptixPipeline {
63   public:
64     /// Class constructor. Creates a ChOptixPipeline that can manage sensors, materials, and SBT information
65     /// @param context The optix context to which this pipeline will be attached
66     /// @param trace_depth the maximum trace depth for ray recusions in this pipeline (contant for all sensors in the
67     /// pipeline)
68     /// @param debug whether the pipeline should be run in debug mode. Only use for development as this slows down the
69     /// ray tracing
70     ChOptixPipeline(OptixDeviceContext context, unsigned int trace_depth, bool debug);
71 
72     /// Class destructor. Cleans up all data associated with the pipeline
73     ~ChOptixPipeline();
74 
75     /// Create a new pipeline for a specific type of sensor
76     /// @param type the type of sensor/ray tracing to add to the pipeline manager
77     void SpawnPipeline(PipelineType type);
78 
79     /// Creates a new box material
80     /// @param mat the chrono material from which to create an optix material
81     /// @returns an id pointing to the material that was created
82     unsigned int GetBoxMaterial(std::shared_ptr<ChVisualMaterial> mat = nullptr);
83 
84     /// Creates a new sphere material
85     /// @param mat the chrono material from which to create an optix material
86     /// @returns an id pointing to the material that was created
87     unsigned int GetSphereMaterial(std::shared_ptr<ChVisualMaterial> mat = nullptr);
88 
89     /// Creates a new cylinder material
90     /// @param mat the chrono material from which to create an optix material
91     /// @returns an id pointing to the material that was created
92     unsigned int GetCylinderMaterial(std::shared_ptr<ChVisualMaterial> mat = nullptr);
93 
94     /// Creates a new rigid material/mesh in optix
95     /// @param[out] d_vertices a device pointer where the mesh's vertices will be stored
96     /// @param[out] d_indices a device pointer where the mesh's indices will be stored
97     /// @param[in] sphere_shape the chrono mesh to add to the optix scene
98     /// @param[in] mat_list the chrono materials from which to create an optix material
99     /// @returns an id pointing to the material that was created
100     unsigned int GetRigidMeshMaterial(CUdeviceptr& d_vertices,
101                                       CUdeviceptr& d_indices,
102                                       std::shared_ptr<ChTriangleMeshShape> sphere_shape,
103                                       std::vector<std::shared_ptr<ChVisualMaterial>> mat_list);
104 
105     /// Creates a new deformable material/mesh in optix
106     /// @param[out] d_vertices a device pointer where the mesh's vertices will be stored
107     /// @param[out] d_indices a device pointer where the mesh's indices will be stored
108     /// @param[in] sphere_shape the chrono mesh to add to the optix scene
109     /// @param[in] mat_list the chrono materials from which to create an optix material
110     /// @returns an id pointing to the material that was created
111     unsigned int GetDeformableMeshMaterial(CUdeviceptr& d_vertices,
112                                            CUdeviceptr& d_indices,
113                                            std::shared_ptr<ChTriangleMeshShape> sphere_shape,
114                                            std::vector<std::shared_ptr<ChVisualMaterial>> mat_list);
115 
116     /// Function to update all the deformable meshes in the optix scene based on their chrono meshes
117     void UpdateDeformableMeshes();
118 
119     /// Function to update all the shader binding tables associated with this optix scene
120     void UpdateAllSBTs();
121 
122     /// Function to update all the pipeline associated with this optix scene
123     void UpdateAllPipelines();
124 
125     /// Function to update all object velocities. Only required for sensors that measure velocity with is a dynamic
126     /// property
127     void UpdateObjectVelocity();
128 
129     /// Function to access the mesh pool on the device
130     /// @returns a device pointer to the collection of meshes
131     CUdeviceptr GetMeshPool();
132 
133     /// Function to access the material pool on the device
134     /// @returns a device pointer to the collection of materials
135     CUdeviceptr GetMaterialPool();
136 
137     /// Add a new body to the optix scene based on a chrono body
138     /// @param body the Chrono body from which to create an optix object
AddBody(std::shared_ptr<ChBody> body)139     void AddBody(std::shared_ptr<ChBody> body) { m_bodies.push_back(body); }
140 
141     /// OptixPipeline accessor function
142     /// @param id the id of the desired pipeline
143     /// @returns a reference to the specified OptixPipeline
144     OptixPipeline& GetPipeline(unsigned int id);
145 
146     /// SBT accessor
147     /// @param id the id of the pipeline whose shader binding table should be returned
148     /// @returns a shared pointer to the specified shader binding table
149     std::shared_ptr<OptixShaderBindingTable> GetSBT(unsigned int id);
150 
151     /// Raygen record accessor
152     /// @param id the id of the pipline whose raygen record should be returned
153     /// @returns a shared pointer to the specified raygen record
154     std::shared_ptr<Record<RaygenParameters>> GetRayGenRecord(unsigned int id);
155 
156     /// Explicit cleanup function for freeing memory associated with this pipeline
157     /// which should be freed before reconstructing the pipeline. Any reusable varaibles
158     /// will not be cleaned up here
159     void Cleanup();
160 
161     /// Explicit material cleanup function for freeing all memory associate with materials
162     /// in the optix scene
163     void CleanMaterials();
164 
165     /// Updates the background of the scene. Only used for cameras
166     /// @param b a new background to use for the scene. The old background will be removed
167     void UpdateBackground(Background b);
168 
169   private:
170     void CompileBaseShaders();
171     void AssembleBaseProgramGroups();
172     void CreateBaseSBT();
173     void CreateOptixProgramGroup(OptixProgramGroup& group,
174                                  OptixProgramGroupKind k,
175                                  OptixModule is_module,
176                                  const char* is_name,
177                                  OptixModule ch_module,
178                                  const char* ch_name);
179 
180     unsigned int GetMaterial(std::shared_ptr<ChVisualMaterial> mat = nullptr);
181 
182     void CreateDeviceTexture(cudaTextureObject_t& d_tex_sampler,
183                              cudaArray_t& d_img_array,
184                              std::string file_name,
185                              bool mirror = false,
186                              bool exclude_from_material_cleanup = false);
187 
188     // Private class variables - DO NOT MODIFY THE CONTEXT
189     OptixDeviceContext m_context;
190 
191     unsigned int m_trace_depth;
192     bool m_debug;
193 
194     // modules - we only make one of each - do not clear when rebuilding root
195     OptixModule m_box_intersection_module = 0;     // box.cu file
196     OptixModule m_sphere_intersection_module = 0;  // sphere.cu file
197     OptixModule m_cyl_intersection_module = 0;     // cylinder.cu file
198     OptixModule m_camera_raygen_module = 0;        // camera.cu file
199     OptixModule m_lidar_raygen_module = 0;         // lidar.cu file
200     OptixModule m_radar_raygen_module = 0;         // lidar.cu file
201     OptixModule m_material_shading_module = 0;     // material shader file
202     OptixModule m_miss_module = 0;                 // miss.cu
203 
204     // program groups - we only make one of each - do not clear when rebuilding root
205     OptixProgramGroup m_camera_pinhole_raygen_group = 0;
206     OptixProgramGroup m_camera_fov_lens_raygen_group = 0;
207     OptixProgramGroup m_segmentation_pinhole_raygen_group = 0;
208     OptixProgramGroup m_segmentation_fov_lens_raygen_group = 0;
209     OptixProgramGroup m_lidar_single_raygen_group = 0;
210     OptixProgramGroup m_lidar_multi_raygen_group = 0;
211     OptixProgramGroup m_radar_raygen_group = 0;
212 
213     OptixProgramGroup m_hit_box_group = 0;
214     OptixProgramGroup m_hit_sphere_group = 0;
215     OptixProgramGroup m_hit_cyl_group = 0;
216     OptixProgramGroup m_hit_mesh_group = 0;
217     OptixProgramGroup m_miss_group = 0;
218 
219     // compile options - TODO: should probably depend on the pipeline - do not clear for now
220     OptixPipelineCompileOptions m_pipeline_compile_options;
221 
222     // the pipelines that are in use - do not clear when rebuilding root
223     std::vector<OptixPipeline> m_pipelines;
224     std::vector<std::shared_ptr<Record<RaygenParameters>>> m_raygen_records;
225 
226     // refresh this with new material information without clearing when root get rebuilt
227     std::vector<std::shared_ptr<OptixShaderBindingTable>> m_sbts;
228 
229     // miss record - do not clear
230     CUdeviceptr md_miss_record = {};  ///< handle to the miss record on the device
231 
232     // material records - clear when rebuilding root
233     std::vector<Record<MaterialRecordParameters>> m_material_records;  ///< vector of all material records
234     CUdeviceptr md_material_records = {};                              ///< handle to the material records on the device
235 
236     // material pool (host and device)
237     std::vector<MaterialParameters> m_material_pool;
238     CUdeviceptr md_material_pool = {};
239 
240     // mesh data pool
241     std::vector<MeshParameters> m_mesh_pool;        ///< mesh pool on the device
242     CUdeviceptr md_mesh_pool = {};                  ///< struct to know which buffers make up which meshes
243     std::vector<CUdeviceptr> m_mesh_buffers_dptrs;  ///< for keeping a handle to free later
244 
245     // texture and image handles
246     std::unordered_map<std::string, cudaTextureObject_t> m_texture_samplers;  ///< for keeping a handle to free later
247     std::unordered_map<std::string, cudaArray_t> m_img_textures;              ///< for keeping a handle to free later
248 
249     cudaTextureObject_t md_miss_texture_sampler = {};  ///< handle to the environment texture sampler
250     cudaArray_t md_miss_img_texture = {};              ///< handle to the environment image texture
251 
252     /// keep track of chrono meshes we've added and their corresponding mesh pool id
253     std::vector<std::tuple<std::shared_ptr<geometry::ChTriangleMeshConnected>, unsigned int>> m_known_meshes;
254 
255     /// list of deformable meshes <mesh shape, dvertices, dnormals, num prev triangles>
256     std::vector<std::tuple<std::shared_ptr<ChTriangleMeshShape>, CUdeviceptr, CUdeviceptr, unsigned int>>
257         m_deformable_meshes;
258 
259     // default material in the material pool
260     bool m_default_material_inst = false;
261     unsigned int m_default_material_id;
262 
263     // bodies in simulation
264     std::vector<std::shared_ptr<ChBody>> m_bodies;
265 };
266 
267 /// @} sensor_optix
268 
269 }  // namespace sensor
270 }  // namespace chrono
271 
272 #endif