Function frostbite_fwd

model/render/renderutils/c_src/torch_bindings.cpp:295–328

// frostbite diffuse

torch::Tensor frostbite_fwd(torch::Tensor nrm, torch::Tensor wi, torch::Tensor wo, torch::Tensor linearRoughness, bool fp16)
{
    CHECK_TENSOR(nrm, 4, 3);
    CHECK_TENSOR(wi, 4, 3);
    CHECK_TENSOR(wo, 4, 3);
    CHECK_TENSOR(linearRoughness, 4, 1);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    // Extract input parameters.
    FrostbiteDiffuseKernelParams p;
    p.out.fp16 = fp16;
    update_grid(p.gridSize, nrm, wi, wo, linearRoughness);

    // Allocate output tensors.
    torch::TensorOptions opts = torch::TensorOptions().dtype(p.out.fp16 ? torch::kBFloat16 : torch::kFloat32).device(torch::kCUDA);
    torch::Tensor out = torch::empty({ p.gridSize.z, p.gridSize.y, p.gridSize.x, 1 }, opts);

    // Choose launch parameters.
    dim3 blockSize = getLaunchBlockSize(BLOCK_X, BLOCK_Y, p.gridSize);
    dim3 gridSize = getLaunchGridSize(blockSize, p.gridSize);

    p.nrm = make_cuda_tensor(nrm, p.gridSize);
    p.wi = make_cuda_tensor(wi, p.gridSize);
    p.wo = make_cuda_tensor(wo, p.gridSize);
    p.linearRoughness = make_cuda_tensor(linearRoughness, p.gridSize);
    p.out = make_cuda_tensor(out, p.gridSize);

    // Launch CUDA kernel.
    void* args[] = { &p };
    NVDR_CHECK_CUDA_ERROR(cudaLaunchKernel((const void*)FrostbiteDiffuseFwdKernel, gridSize, blockSize, args, 0, stream));

    return out;
}

std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> frostbite_bwd(torch::Tensor nrm, torch::Tensor wi, torch::Tensor wo, torch::Tensor linearRoughness, torch::Tensor grad)
{