Repository: NVlabs/instant-ngp
File: src/testbed_nerf.cu (CUDA)

Code before:
	float* envmap_gradient = m_nerf.training.train_envmap ? m_envmap.envmap->gradients() : nullptr;

	bool sample_focal_plane_proportional_to_error = m_nerf.training.error_map.is_cdf_valid &&
		m_nerf.training.sample_focal_plane_proportional_to_error;
	bool sample_image_proportional_to_error = m_nerf.training.error_map.is_cdf_valid && m_nerf.training.sample_image_proportional_to_error;
	bool include_sharpness_in_error = m_nerf.training.include_sharpness_in_error;
	// This is low-overhead enough to warrant always being on.
	// It makes for useful visualizations of the training error.
	bool accumulate_error = true;

	CUDA_CHECK_THROW(cudaMemsetAsync(ray_counter, 0, sizeof(uint32_t), stream));

	auto hg_enc = dynamic_cast<MultiLevelEncoding<network_precision_t>*>(m_encoding.get());

	// RFL training only implemented in JIT training mode.
	if (!m_jit_fusion && m_nerf.training.train_mode != ETrainMode::Nerf) {
		m_nerf.training.train_mode = ETrainMode::Nerf;
		tlog::warning() << "JIT fusion is disabled, switching to NeRF training mode.";
	}

	//  TODO: the below fused kernel is actually slower than the unfused alternative for NeRF training.
	//        look into optimizing it until it is faster.
	if (m_jit_fusion &&
		(m_nerf.training.train_mode == ETrainMode::Rfl ||
		 (m_nerf.training.train_mode == ETrainMode::RflRelax &&
		  m_training_step >= 15000 &&
		  m_training_step < 30000))) {
		// JIT training, supports Nerf/Rfl/RflRelax training mode. But for pure
		// NeRF training, it is faster to use the unfused kernel.
		auto jit_guard = m_nerf_network->jit_guard(stream, false);

		if (!m_nerf.training.fused_kernel) {
			try {
				m_nerf.training.fused_kernel = std::make_unique<CudaRtcKernel>(
					"train_nerf",
					fmt::format(
						"{MODEL_BODY}\n"
						"using GRID_T = {GRID_T};\n"
						"static constexpr uint32_t N_EXTRA_DIMS = {N_EXTRA_DIMS};\n"
						"#include <neural-graphics-primitives/fused_kernels/train_nerf.cuh>\n",
						fmt::arg("MODEL_BODY", m_nerf_network->generate_device_function("eval_nerf")),
						fmt::arg("GRID_T", type_to_string<network_precision_t>()),
						fmt::arg("N_EXTRA_DIMS", m_nerf.training.dataset.n_extra_dims())
					),
					all_files(cmrc::ngp::get_filesystem())
				);
			} catch (const std::runtime_error& e) {
				tlog::warning() << e.what();
				tlog::warning() << "Disabling JIT fusion.";
				m_jit_fusion = false;
			}

Review comment (refactor):
I feel like it is too restrictive to hardcode the 15k-30k limit here. If a user selects `Rfl` or `RflRelaxed` training that should be what they're getting.

Could you instead move this logic into `run.py` and gate it behind a (default True) CLI flag called `--rfl_training_schedule`. Also please make sure that the `--nerf_compatibility` parameter disables all rfl-related things. Thanks!

Code after:
	float* envmap_gradient = m_nerf.training.train_envmap ? m_envmap.envmap->gradients() : nullptr;

	bool sample_focal_plane_proportional_to_error = m_nerf.training.error_map.is_cdf_valid &&
		m_nerf.training.sample_focal_plane_proportional_to_error;
	bool sample_image_proportional_to_error = m_nerf.training.error_map.is_cdf_valid && m_nerf.training.sample_image_proportional_to_error;
	bool include_sharpness_in_error = m_nerf.training.include_sharpness_in_error;
	// This is low-overhead enough to warrant always being on.
	// It makes for useful visualizations of the training error.
	bool accumulate_error = true;

	CUDA_CHECK_THROW(cudaMemsetAsync(ray_counter, 0, sizeof(uint32_t), stream));

	auto hg_enc = dynamic_cast<MultiLevelEncoding<network_precision_t>*>(m_encoding.get());

	// RFL training only implemented in JIT training mode.
	if (!m_jit_fusion && m_nerf.training.train_mode != ETrainMode::Nerf) {
		m_nerf.training.train_mode = ETrainMode::Nerf;
		tlog::warning() << "JIT fusion is disabled, switching to NeRF training mode.";
	}

	//  TODO: the below fused kernel is actually slower than the unfused alternative for NeRF training.
	//        look into optimizing it until it is faster.
	if (m_jit_fusion && (m_nerf.training.train_mode != ETrainMode::Nerf)) {
		// JIT training, supports Nerf/Rfl/RflRelax training mode. But for pure
		// NeRF training, it is faster to use the unfused kernel.
		auto jit_guard = m_nerf_network->jit_guard(stream, false);

		if (!m_nerf.training.fused_kernel) {
			try {
				m_nerf.training.fused_kernel = std::make_unique<CudaRtcKernel>(
					"train_nerf",
					fmt::format(
						"{MODEL_BODY}\n"
						"using GRID_T = {GRID_T};\n"
						"static constexpr uint32_t N_EXTRA_DIMS = {N_EXTRA_DIMS};\n"
						"#include <neural-graphics-primitives/fused_kernels/train_nerf.cuh>\n",
						fmt::arg("MODEL_BODY", m_nerf_network->generate_device_function("eval_nerf")),
						fmt::arg("GRID_T", type_to_string<network_precision_t>()),
						fmt::arg("N_EXTRA_DIMS", m_nerf.training.dataset.n_extra_dims())
					),
					all_files(cmrc::ngp::get_filesystem())
				);
			} catch (const std::runtime_error& e) {
				tlog::warning() << e.what();
				tlog::warning() << "Disabling JIT fusion.";
				m_jit_fusion = false;
			}
		}

		if (m_nerf.training.fused_kernel) {
			m_nerf.training.fused_kernel->launch(