Add homegrown vector lib, log callbacks instead of iostream, misc refactor

Author: Tom94

CMakeLists.txt

@@ -105,17 +105,7 @@ function(TCNN_AUTODETECT_CUDA_ARCHITECTURES OUT_VARIABLE)
 			"}\n"
 		)
 
-		if (CMAKE_CUDA_COMPILER_LOADED) # CUDA as a language
-			try_run(run_result compile_result ${PROJECT_BINARY_DIR} ${file} RUN_OUTPUT_VARIABLE compute_capabilities)
-		else()
-			try_run(
-				run_result compile_result ${PROJECT_BINARY_DIR} ${file}
-				CMAKE_FLAGS "-DINCLUDE_DIRECTORIES=${CUDA_INCLUDE_DIRS}"
-				LINK_LIBRARIES ${CUDA_LIBRARIES}
-				RUN_OUTPUT_VARIABLE compute_capabilities
-			)
-		endif()
-
+		try_run(run_result compile_result ${PROJECT_BINARY_DIR} ${file} RUN_OUTPUT_VARIABLE compute_capabilities)
 		if (run_result EQUAL 0)
 			# If the user has multiple GPUs with the same compute capability installed, list that capability only once.
 			list(REMOVE_DUPLICATES compute_capabilities)
@@ -257,8 +247,7 @@ set(CMAKE_RUNTIME_OUTPUT_DIRECTORY_MINSIZEREL ${CMAKE_BINARY_DIR})
 set(CMAKE_RUNTIME_OUTPUT_DIRECTORY_DEBUG ${CMAKE_BINARY_DIR})
 
 set(TCNN_SOURCES
-	src/common.cu
-	src/common_device.cu
+	src/common_host.cu
 	src/cpp_api.cu
 	src/cutlass_mlp.cu
 	src/encoding.cu

README.md

@@ -45,7 +45,7 @@ using namespace tcnn;
 
 auto model = create_from_config(n_input_dims, n_output_dims, config);
 
-// Train the model (batch_size must be a multiple of tcnn::batch_size_granularity)
+// Train the model (batch_size must be a multiple of tcnn::BATCH_SIZE_GRANULARITY)
 GPUMatrix<float> training_batch_inputs(n_input_dims, batch_size);
 GPUMatrix<float> training_batch_targets(n_output_dims, batch_size);
 

bindings/torch/setup.py

@@ -4,8 +4,10 @@
 from setuptools import setup
 from pkg_resources import parse_version
 import subprocess
+import shutil
 import sys
 import torch
+from glob import glob
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
 SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
@@ -123,15 +125,21 @@ def find_cl_path():
 # List of sources.
 bindings_dir = os.path.dirname(__file__)
 root_dir = os.path.abspath(os.path.join(bindings_dir, "../.."))
-base_definitions = []
+
+base_definitions = [
+	# PyTorch-supplied parameters may be unaligned. TCNN must be made aware of this such that
+	# it does not optimize for aligned memory accesses.
+	"-DTCNN_PARAMS_UNALIGNED",
+]
+
 base_source_files = [
 	"tinycudann/bindings.cpp",
 	"../../dependencies/fmt/src/format.cc",
 	"../../dependencies/fmt/src/os.cc",
 	"../../src/cpp_api.cu",
-	"../../src/common.cu",
-	"../../src/common_device.cu",
+	"../../src/common_host.cu",
 	"../../src/encoding.cu",
+	"../../src/object.cu",
 ]
 
 if include_networks:

bindings/torch/tinycudann/bindings.cpp

@@ -44,6 +44,7 @@
 #include <json/json.hpp>
 
 #include <pybind11_json/pybind11_json.hpp>
+#include <pybind11/functional.h>
 
 #include <tiny-cuda-nn/cpp_api.h>
 
@@ -53,10 +54,10 @@
 #define CHECK_THROW(x) \
 	do { if (!(x)) throw std::runtime_error(std::string(FILE_LINE " check failed " #x)); } while(0)
 
-c10::ScalarType torch_type(tcnn::cpp::EPrecision precision) {
+c10::ScalarType torch_type(tcnn::cpp::Precision precision) {
 	switch (precision) {
-		case tcnn::cpp::EPrecision::Fp32: return torch::kFloat32;
-		case tcnn::cpp::EPrecision::Fp16: return torch::kHalf;
+		case tcnn::cpp::Precision::Fp32: return torch::kFloat32;
+		case tcnn::cpp::Precision::Fp16: return torch::kHalf;
 		default: throw std::runtime_error{"Unknown precision tcnn->torch"};
 	}
 }
@@ -246,41 +247,19 @@ class Module {
 		return output;
 	}
 
-	uint32_t n_input_dims() const {
-		return m_module->n_input_dims();
-	}
+	uint32_t n_input_dims() const { return m_module->n_input_dims(); }
 
-	uint32_t n_params() const {
-		return (uint32_t)m_module->n_params();
-	}
+	uint32_t n_params() const { return (uint32_t)m_module->n_params(); }
+	tcnn::cpp::Precision param_precision() const { return m_module->param_precision(); }
+	c10::ScalarType c10_param_precision() const { return torch_type(param_precision()); }
 
-	tcnn::cpp::EPrecision param_precision() const {
-		return m_module->param_precision();
-	}
+	uint32_t n_output_dims() const { return m_module->n_output_dims(); }
+	tcnn::cpp::Precision output_precision() const { return m_module->output_precision(); }
+	c10::ScalarType c10_output_precision() const { return torch_type(output_precision()); }
 
-	c10::ScalarType c10_param_precision() const {
-		return torch_type(param_precision());
-	}
+	nlohmann::json hyperparams() const { return m_module->hyperparams(); }
+	std::string name() const { return m_module->name(); }
 
-	uint32_t n_output_dims() const {
-		return m_module->n_output_dims();
-	}
-
-	tcnn::cpp::EPrecision output_precision() const {
-		return m_module->output_precision();
-	}
-
-	c10::ScalarType c10_output_precision() const {
-		return torch_type(output_precision());
-	}
-
-	nlohmann::json hyperparams() const {
-		return m_module->hyperparams();
-	}
-
-	std::string name() const {
-		return m_module->name();
-	}
 
 private:
 	std::unique_ptr<tcnn::cpp::Module> m_module;
@@ -296,22 +275,34 @@ Module create_network(uint32_t n_input_dims, uint32_t n_output_dims, const nlohm
 }
 #endif
 
-Module create_encoding(uint32_t n_input_dims, const nlohmann::json& encoding, tcnn::cpp::EPrecision requested_precision) {
+Module create_encoding(uint32_t n_input_dims, const nlohmann::json& encoding, tcnn::cpp::Precision requested_precision) {
 	return Module{tcnn::cpp::create_encoding(n_input_dims, encoding, requested_precision)};
 }
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-	py::enum_<tcnn::cpp::EPrecision>(m, "Precision")
-		.value("Fp32", tcnn::cpp::EPrecision::Fp32)
-		.value("Fp16", tcnn::cpp::EPrecision::Fp16)
+	py::enum_<tcnn::cpp::LogSeverity>(m, "LogSeverity")
+		.value("Info", tcnn::cpp::LogSeverity::Info)
+		.value("Debug", tcnn::cpp::LogSeverity::Debug)
+		.value("Warning", tcnn::cpp::LogSeverity::Warning)
+		.value("Error", tcnn::cpp::LogSeverity::Error)
+		.value("Success", tcnn::cpp::LogSeverity::Success)
+		.export_values()
+		;
+
+	py::enum_<tcnn::cpp::Precision>(m, "Precision")
+		.value("Fp32", tcnn::cpp::Precision::Fp32)
+		.value("Fp16", tcnn::cpp::Precision::Fp16)
 		.export_values()
 		;
 
 	m.def("batch_size_granularity", &tcnn::cpp::batch_size_granularity);
+	m.def("default_loss_scale", &tcnn::cpp::default_loss_scale);
 	m.def("free_temporary_memory", &tcnn::cpp::free_temporary_memory);
 	m.def("has_networks", &tcnn::cpp::has_networks);
 	m.def("preferred_precision", &tcnn::cpp::preferred_precision);
 
+	m.def("set_log_callback", &tcnn::cpp::set_log_callback);
+
 	// Encapsulates an abstract context of an operation
 	// (commonly the forward pass) to be passed on to other
 	// operations (commonly the backward pass).

bindings/torch/tinycudann/modules.py

@@ -8,6 +8,7 @@
 
 import gc
 import importlib
+import os
 import warnings
 
 import torch
@@ -57,6 +58,14 @@ def _get_system_compute_capability():
 if _C is None:
 	raise EnvironmentError(f"Could not find compatible tinycudann extension for compute capability {system_compute_capability}.")
 
+# Pipe tcnn warnings and errors into Python
+# def _log(severity, msg):
+# 	if severity == _C.LogSeverity.Warning:
+# 		warnings.warn(f"tinycudann warning: {msg}")
+# 	elif severity == _C.LogSeverity.Error:
+# 		warnings.warn(f"tinycudann error: {msg}")
+
+# _C.set_log_callback(_log)
 def _torch_precision(tcnn_precision):
 	if tcnn_precision == _C.Precision.Fp16:
 		return torch.half
@@ -162,7 +171,7 @@ def __init__(self, seed=1337):
 		self.params = torch.nn.Parameter(initial_params, requires_grad=True)
 		self.register_parameter(name="params", param=self.params)
 
-		self.loss_scale = 128.0 if self.native_tcnn_module.param_precision() == _C.Precision.Fp16 else 1.0
+		self.loss_scale = _C.default_loss_scale(self.native_tcnn_module.param_precision())
 
 	def forward(self, x):
 		if not x.is_cuda: