OptixRenderer.cpp

#ifdef _WIN32
#define NOMINMAX
#endif

#include "OptixRenderer.h"
#include "ptx_string.h"
#include <iostream>
#include <optix_function_table_definition.h>
#include <optix_stubs.h>


#define CUDA_CHECK(call)                                                       \
  {                                                                            \
    cudaError_t res = call;                                                    \
    if (res != cudaSuccess) {                                                  \
      std::cerr << "CUDA Error: " << cudaGetErrorString(res) << std::endl;     \
      exit(1);                                                                 \
    }                                                                          \
  }
#define OPTIX_CHECK(call)                                                      \
  {                                                                            \
    OptixResult res = call;                                                    \
    if (res != OPTIX_SUCCESS) {                                                \
      std::cerr << "OptiX Error: " << res << std::endl;                        \
      exit(1);                                                                 \
    }                                                                          \
  }

template <typename T> struct SbtRecord {
  __align__(
      OPTIX_SBT_RECORD_ALIGNMENT) char header[OPTIX_SBT_RECORD_HEADER_SIZE];
  T data;
};
typedef SbtRecord<int> EmptyRecord;

struct Params {
  unsigned int *results;
  float3 *vertices;
  unsigned int num_points;
  OptixTraversableHandle handle;
};

OptixRenderer::OptixRenderer() {
  initOptix();
  createPipeline();
}

OptixRenderer::~OptixRenderer() {
  // Only free internally managed memory
  if (d_gas_output)
    cudaFree(reinterpret_cast<void *>(d_gas_output));
  if (d_rg)
    cudaFree(reinterpret_cast<void *>(d_rg));
  if (d_ms)
    cudaFree(reinterpret_cast<void *>(d_ms));
  if (d_ht)
    cudaFree(reinterpret_cast<void *>(d_ht));

  if (pipeline)
    optixPipelineDestroy(pipeline);
  if (raygenPG)
    optixProgramGroupDestroy(raygenPG);
  if (missPG)
    optixProgramGroupDestroy(missPG);
  if (hitPG)
    optixProgramGroupDestroy(hitPG);
  if (module)
    optixModuleDestroy(module);
  if (context)
    optixDeviceContextDestroy(context);
}

void OptixRenderer::initOptix() {
  CUcontext cuCtx = 0;
  OPTIX_CHECK(optixInit());
  OPTIX_CHECK(optixDeviceContextCreate(cuCtx, nullptr, &context));
}

void OptixRenderer::createPipeline() {
  OptixModuleCompileOptions moduleOptions = {};
  // Level 0 for fast JIT compilation during development
  moduleOptions.optLevel = OPTIX_COMPILE_OPTIMIZATION_LEVEL_0;
  moduleOptions.debugLevel = OPTIX_COMPILE_DEBUG_LEVEL_NONE;

  OptixPipelineCompileOptions pipelineOptions = {};
  pipelineOptions.usesMotionBlur = false;
  pipelineOptions.traversableGraphFlags =
      OPTIX_TRAVERSABLE_GRAPH_FLAG_ALLOW_SINGLE_GAS;
  pipelineOptions.numPayloadValues = 1;
  pipelineOptions.numAttributeValues = 2;
  pipelineOptions.exceptionFlags = OPTIX_EXCEPTION_FLAG_NONE;
  pipelineOptions.pipelineLaunchParamsVariableName = "params";

  char log[2048];
  size_t logSize = sizeof(log);
  OPTIX_CHECK(optixModuleCreate(context, &moduleOptions, &pipelineOptions,
                                reinterpret_cast<const char *>(ptx_code),
                                sizeof(ptx_code) - 1, log, &logSize, &module));

  OptixProgramGroupOptions pgOptions = {};
  OptixProgramGroupDesc rgDesc = {};
  rgDesc.kind = OPTIX_PROGRAM_GROUP_KIND_RAYGEN;
  rgDesc.raygen.module = module;
  rgDesc.raygen.entryFunctionName = "__raygen__rg";
  OPTIX_CHECK(optixProgramGroupCreate(context, &rgDesc, 1, &pgOptions, log,
                                      &logSize, &raygenPG));

  OptixProgramGroupDesc msDesc = {};
  msDesc.kind = OPTIX_PROGRAM_GROUP_KIND_MISS;
  msDesc.miss.module = module;
  msDesc.miss.entryFunctionName = "__miss__ms";
  OPTIX_CHECK(optixProgramGroupCreate(context, &msDesc, 1, &pgOptions, log,
                                      &logSize, &missPG));

  OptixProgramGroupDesc htDesc = {};
  htDesc.kind = OPTIX_PROGRAM_GROUP_KIND_HITGROUP;
  htDesc.hitgroup.moduleCH = module;
  htDesc.hitgroup.entryFunctionNameCH = "__closesthit__ch";
  OPTIX_CHECK(optixProgramGroupCreate(context, &htDesc, 1, &pgOptions, log,
                                      &logSize, &hitPG));

  OptixProgramGroup programGroups[] = {raygenPG, missPG, hitPG};
  OptixPipelineLinkOptions pipelineLinkOptions = {};
  pipelineLinkOptions.maxTraceDepth = 1;
  OPTIX_CHECK(optixPipelineCreate(context, &pipelineOptions,
                                  &pipelineLinkOptions, programGroups, 3, log,
                                  &logSize, &pipeline));
}

void OptixRenderer::buildGeometry(float3 *d_vertices, unsigned int num_vertices,
                                  uint3 *d_indices, unsigned int num_indices) {
  CUdeviceptr cu_vertices = reinterpret_cast<CUdeviceptr>(d_vertices);
  CUdeviceptr cu_indices = reinterpret_cast<CUdeviceptr>(d_indices);

  OptixBuildInput buildInput = {};
  buildInput.type = OPTIX_BUILD_INPUT_TYPE_TRIANGLES;
  buildInput.triangleArray.vertexFormat = OPTIX_VERTEX_FORMAT_FLOAT3;
  buildInput.triangleArray.numVertices = num_vertices;
  buildInput.triangleArray.vertexBuffers = &cu_vertices;
  buildInput.triangleArray.indexFormat = OPTIX_INDICES_FORMAT_UNSIGNED_INT3;
  buildInput.triangleArray.numIndexTriplets = num_indices;
  buildInput.triangleArray.indexBuffer = cu_indices;

  unsigned int triangleFlags = OPTIX_GEOMETRY_FLAG_NONE;
  buildInput.triangleArray.flags = &triangleFlags;
  buildInput.triangleArray.numSbtRecords = 1;

  OptixAccelBuildOptions accelOptions = {};
  accelOptions.buildFlags = OPTIX_BUILD_FLAG_ALLOW_COMPACTION;
  accelOptions.operation = OPTIX_BUILD_OPERATION_BUILD;

  OptixAccelBufferSizes bufferSizes;
  OPTIX_CHECK(optixAccelComputeMemoryUsage(context, &accelOptions, &buildInput,
                                           1, &bufferSizes));

  CUdeviceptr d_temp_buffer_gas;
  CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&d_temp_buffer_gas),
                        bufferSizes.tempSizeInBytes));
  CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&d_gas_output),
                        bufferSizes.outputSizeInBytes));

  OPTIX_CHECK(optixAccelBuild(context, 0, &accelOptions, &buildInput, 1,
                              d_temp_buffer_gas, bufferSizes.tempSizeInBytes,
                              d_gas_output, bufferSizes.outputSizeInBytes,
                              &gas_handle, nullptr, 0));

  CUDA_CHECK(cudaFree(reinterpret_cast<void *>(d_temp_buffer_gas)));

  createSBT();
  CUDA_CHECK(cudaDeviceSynchronize());
}

void OptixRenderer::createSBT() {
  EmptyRecord rgRecord, msRecord, htRecord;
  OPTIX_CHECK(optixSbtRecordPackHeader(raygenPG, &rgRecord));
  OPTIX_CHECK(optixSbtRecordPackHeader(missPG, &msRecord));
  OPTIX_CHECK(optixSbtRecordPackHeader(hitPG, &htRecord));

  CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&d_rg), sizeof(EmptyRecord)));
  CUDA_CHECK(cudaMemcpy(reinterpret_cast<void *>(d_rg), &rgRecord,
                        sizeof(EmptyRecord), cudaMemcpyHostToDevice));
  CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&d_ms), sizeof(EmptyRecord)));
  CUDA_CHECK(cudaMemcpy(reinterpret_cast<void *>(d_ms), &msRecord,
                        sizeof(EmptyRecord), cudaMemcpyHostToDevice));
  CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&d_ht), sizeof(EmptyRecord)));
  CUDA_CHECK(cudaMemcpy(reinterpret_cast<void *>(d_ht), &htRecord,
                        sizeof(EmptyRecord), cudaMemcpyHostToDevice));

  sbt.raygenRecord = d_rg;
  sbt.missRecordBase = d_ms;
  sbt.missRecordStrideInBytes = sizeof(EmptyRecord);
  sbt.missRecordCount = 1;
  sbt.hitgroupRecordBase = d_ht;
  sbt.hitgroupRecordStrideInBytes = sizeof(EmptyRecord);
  sbt.hitgroupRecordCount = 1;
}

void OptixRenderer::render(float3 *d_vertices, unsigned int num_points,
                           unsigned int *d_results) {
  Params params = {d_results, d_vertices, num_points, gas_handle};

  CUdeviceptr d_params;
  CUDA_CHECK(cudaMalloc(reinterpret_cast<void **>(&d_params), sizeof(Params)));
  CUDA_CHECK(cudaMemcpy(reinterpret_cast<void *>(d_params), &params,
                        sizeof(Params), cudaMemcpyHostToDevice));

  OPTIX_CHECK(optixLaunch(pipeline, 0, d_params, sizeof(Params), &sbt,
                          num_points, 1, 1));
  CUDA_CHECK(cudaDeviceSynchronize());

  CUDA_CHECK(cudaFree(reinterpret_cast<void *>(d_params)));
}