Merge pull request #45 from cuteday/vdbtype

Support vec3f (rgb) openvdb data for volume rendering

src/core/common.h

@@ -12,9 +12,6 @@
 #include <sstream>
 #include <vector>
 
-#include <cuda.h>
-#include <cuda_runtime.h>
-#include <cuda_runtime_api.h>
 #include <json.hpp>
 
 #include "config.h"	
@@ -82,6 +79,8 @@ typedef unsigned char uchar;
 # define KRR_DEVICE_LAMBDA(...) [ =, *this ] KRR_DEVICE(__VA_ARGS__) mutable 
 
 #if !defined(__CUDA_ARCH__)
+struct uint3;
+struct dim3;
 extern const uint3 threadIdx, blockIdx;
 extern const dim3 blockDim, gridDim;
 #endif	// eliminate intellisense warnings for these kernel built-in variables

src/core/device/buffer.h

@@ -6,6 +6,7 @@
 #include <thrust/execution_policy.h>
 
 #include "common.h"
+#include "device/cuda.h"
 #include "util/check.h"
 
 KRR_NAMESPACE_BEGIN

src/core/device/cuda.h

@@ -1,10 +1,13 @@
 #pragma once
 
-#include <cuda.h>
 #include <map>
 #include <typeinfo>
 #include <typeindex>
 
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cuda_runtime_api.h>
+
 #include "common.h"
 #include "render/color.h"
 #include "device/taggedptr.h"

src/core/device/gpustd.h

@@ -17,6 +17,7 @@
 #include <utility>
 
 #include "util/check.h"
+#include "cuda.h"
 #include "logger.h"
 #include "common.h"
 

src/core/device/scene.cpp

@@ -154,27 +154,42 @@ void RTScene::uploadSceneMediumData() {
 			HomogeneousMedium gMedium(m->sigma_a, m->sigma_s, m->Le, m->g, KRR_DEFAULT_COLORSPACE);
 			mHomogeneousMedium.push_back(gMedium);
 		} else if (auto m = std::dynamic_pointer_cast<VDBVolume>(medium)) {
-			float maxDensity{0};
-			mNanoVDBMedium.emplace_back(m->getNode()->getGlobalTransform(), m->sigma_a, m->sigma_s,
-										m->g, std::move(*m->densityGrid),
-										m->temperatureGrid ? std::move(*m->temperatureGrid)
-														   : NanoVDBGrid{},
-										m->LeScale, m->temperatureScale, m->temperatureOffset, 
-										KRR_DEFAULT_COLORSPACE);
+			if (std::dynamic_pointer_cast<NanoVDBGrid<float>>(m->densityGrid)) {
+				mNanoVDBMedium.emplace_back(
+					m->getNode()->getGlobalTransform(), m->sigma_a, m->sigma_s, m->g,
+					std::move(*std::dynamic_pointer_cast<NanoVDBGrid<float>>(m->densityGrid)),
+					m->temperatureGrid ? std::move(*m->temperatureGrid) : NanoVDBGrid<float>{},
+					m->LeScale, m->temperatureScale, m->temperatureOffset, KRR_DEFAULT_COLORSPACE);
+			} else if (std::dynamic_pointer_cast<NanoVDBGrid<Array3f>>(m->densityGrid)) {
+				mNanoVDBRGBMedium.emplace_back(
+					m->getNode()->getGlobalTransform(), m->sigma_a, m->sigma_s, m->g,
+					std::move(*std::dynamic_pointer_cast<NanoVDBGrid<Array3f>>(m->densityGrid)),
+					m->temperatureGrid ? std::move(*m->temperatureGrid) : NanoVDBGrid<float>{},
+					m->LeScale, m->temperatureScale, m->temperatureOffset, KRR_DEFAULT_COLORSPACE);
+			} else {
+				Log(Error, "Unsupported heterogeneous VDB medium data type");
+				continue;
+			}
 			mNanoVDBMedium.back().initializeFromHost();
-		}
+		} else
+			Log(Error, "Unknown medium type not uploaded to device memory.");
 	}
 	mHomogeneousMediumBuffer.alloc_and_copy_from_host(mHomogeneousMedium);
 	mNanoVDBMediumBuffer.alloc_and_copy_from_host(mNanoVDBMedium);
+	mNanoVDBRGBMediumBuffer.alloc_and_copy_from_host(mNanoVDBRGBMedium);
 
 	size_t homogeneousId = 0;
 	size_t nanoVDBId	 = 0;
+	size_t nanoVDBRGBId	 = 0;
 	for (auto medium : mScene.lock()->getMedia()) {
 		if (auto m = std::dynamic_pointer_cast<HomogeneousVolume>(medium)) 
 			mMedium.push_back(Medium(&mHomogeneousMediumBuffer[homogeneousId++]));
-		else if (auto m = std::dynamic_pointer_cast<VDBVolume>(medium))
-			mMedium.push_back(Medium(&mNanoVDBMediumBuffer[nanoVDBId++]));
-		else Log(Error, "Unknown medium type not uploaded to device memory.");
+		else if (auto m = std::dynamic_pointer_cast<VDBVolume>(medium)) {
+			if (std::dynamic_pointer_cast<NanoVDBGrid<float>>(m->densityGrid))
+				mMedium.push_back(Medium(&mNanoVDBMediumBuffer[nanoVDBId++]));
+			else if (std::dynamic_pointer_cast<NanoVDBGrid<Array3f>>(m->densityGrid))
+				mMedium.push_back(Medium(&mNanoVDBRGBMediumBuffer[nanoVDBRGBId++]));
+		}
 	}
 	mMediumBuffer.alloc_and_copy_from_host(mMedium);
 }

src/core/device/scene.h

@@ -1,11 +1,14 @@
 #pragma once
+#include <optix.h>
+#include <optix_stubs.h>
+
 #include "common.h"
 #include "mesh.h"
 #include "light.h"
 #include "camera.h"
 #include "texture.h"
-#include "device/gpustd.h"
 
+#include "device/gpustd.h"
 #include "device/buffer.h"
 #include "device/memory.h"
 #include "render/lightsampler.h"
@@ -77,8 +80,10 @@ class RTScene {
 
 	std::vector<HomogeneousMedium> mHomogeneousMedium;
 	TypedBuffer<HomogeneousMedium> mHomogeneousMediumBuffer;
-	std::vector<NanoVDBMedium> mNanoVDBMedium;
-	TypedBuffer<NanoVDBMedium> mNanoVDBMediumBuffer;
+	std::vector<NanoVDBMedium<float>> mNanoVDBMedium;
+	TypedBuffer<NanoVDBMedium<float>> mNanoVDBMediumBuffer;
+	std::vector<NanoVDBMedium<Array3f>> mNanoVDBRGBMedium;
+	TypedBuffer<NanoVDBMedium<Array3f>> mNanoVDBRGBMediumBuffer;
 
 	std::vector<Medium> mMedium;
 	TypedBuffer<Medium> mMediumBuffer;

src/core/medium.h

@@ -11,9 +11,9 @@ class Ray;
 class HGPhaseFunction;
 class PhaseFunctionSample;
 class HomogeneousMedium;
-class NanoVDBMedium;
 class MediumProperties;
 class MajorantIterator;
+template <typename DataType> class NanoVDBMedium;
 
 class PhaseFunction : public TaggedPointer<HGPhaseFunction> {
 public:
@@ -32,7 +32,7 @@ struct MediumProperties {
 	Spectrum Le;
 };
 
-class Medium : public TaggedPointer<HomogeneousMedium, NanoVDBMedium> {
+class Medium : public TaggedPointer<HomogeneousMedium, NanoVDBMedium<float>, NanoVDBMedium<Array3f>> {
 public:
 	using TaggedPointer::TaggedPointer;
 

src/core/scenegraph.cpp

@@ -611,12 +611,11 @@ void HomogeneousVolume::renderUI() {
 	if (isEmissive()) ui::Text("Le: %s", Le.string().c_str());
 }
 
-void VDBVolume::renderUI() {
-	ui::Text("VDB volume"); 
+void Volume::renderUI() {
+	ui::Text("Volume Data");
 	ui::Text(("Sigma_a: " + sigma_a.string()).c_str());
 	ui::Text(("Sigma_s: " + sigma_s.string()).c_str());
 	ui::Text("g: %f", g);
-	ui::Text("Max density: %f", densityGrid->getMaxDensity());
 }
 
 void SceneGraphNode::renderUI() { 

src/core/scenegraph.h

@@ -131,10 +131,15 @@ class Volume : public SceneGraphLeaf {
 public:
 	using SharedPtr = std::shared_ptr<Volume>;
 	Volume()		= default;
+	Volume(RGB sigma_a, RGB sigma_s, float g) : sigma_a(sigma_a), sigma_s(sigma_s), g(g) {}
 
 	int getMediumId() const { return mediumId; }
 	void setMediumId(int id) { mediumId = id; }
+	virtual void renderUI() override;
 
+	RGB sigma_a;
+	RGB sigma_s;
+	float g;
 protected:
 	friend class SceneGraph;
 	int mediumId{-1};
@@ -145,18 +150,14 @@ class HomogeneousVolume : public Volume {
 	using SharedPtr = std::shared_ptr<HomogeneousVolume>;
 
 	HomogeneousVolume(RGB sigma_a, RGB sigma_s, float g, RGB Le = RGB::Zero()) :
-		sigma_a(sigma_a), sigma_s(sigma_s), g(g), Le(Le) {}
+		Volume(sigma_a, sigma_s, g), Le(Le) {}
 
 	virtual std::shared_ptr<SceneGraphLeaf> clone() override;
 	virtual void renderUI() override;
 
 	bool isEmissive() const { return !Le.isZero(); }
 
-	RGB sigma_a;
-	RGB sigma_s;
 	RGB Le;
-	float g;
-
 private:
 	friend class SceneGraph;
 };
@@ -165,24 +166,20 @@ class VDBVolume : public Volume {
 public:
 	using SharedPtr = std::shared_ptr<VDBVolume>;
 
-	VDBVolume(RGB sigma_a, RGB sigma_s, float g, 
-		NanoVDBGrid::SharedPtr density, NanoVDBGrid::SharedPtr temperature = nullptr, 
-		float LeScale = 1, float temperatureScale = 1, float temperatureOffset = 0) :
-		sigma_a(sigma_a), sigma_s(sigma_s), g(g), densityGrid(density), temperatureGrid(temperature), 
+	VDBVolume(RGB sigma_a, RGB sigma_s, float g, NanoVDBGridBase::SharedPtr density,
+			  NanoVDBGrid<float>::SharedPtr temperature = nullptr, float LeScale = 1,
+			  float temperatureScale = 1, float temperatureOffset = 0) :
+		Volume(sigma_a, sigma_s, g), densityGrid(density), temperatureGrid(temperature),
 		LeScale(LeScale), temperatureScale(temperatureScale), temperatureOffset(temperatureOffset) {}
 
-	virtual std::shared_ptr<SceneGraphLeaf> clone() override;
+	virtual SceneGraphLeaf::SharedPtr clone() override;
 	virtual AABB getLocalBoundingBox() const override { return densityGrid->getBounds(); }
-	virtual void renderUI() override;
 
-	RGB sigma_a;
-	RGB sigma_s;
-	float g;
 	float LeScale;
 	float temperatureScale;
 	float temperatureOffset;
-	NanoVDBGrid::SharedPtr densityGrid;
-	NanoVDBGrid::SharedPtr temperatureGrid;
+	NanoVDBGridBase::SharedPtr densityGrid;
+	NanoVDBGrid<float>::SharedPtr temperatureGrid;
 
 protected:
 	friend class SceneGraph;

src/core/vulkan/cuvk.h

@@ -16,6 +16,7 @@
 
 #include <common.h>
 #include <logger.h>
+#include <device/cuda.h>
 #include <util/check.h>
 
 KRR_NAMESPACE_BEGIN

src/render/color.h

@@ -1,6 +1,5 @@
 #pragma once
 
-#include "common.h"
 #include "krrmath/math.h"
 #include "util/math_utils.h"
 

src/render/media.cpp

@@ -8,8 +8,8 @@
 KRR_NAMESPACE_BEGIN
 /* Note that the function qualifier (e.g. inline) should be consistent between declaration and definition. */
 
-void initializeMajorantGrid(MajorantGrid& majorantGrid,
-							nanovdb::FloatGrid *floatGrid) {
+template <typename GridType>
+void initializeMajorantGrid(MajorantGrid &majorantGrid, GridType *densityGrid) {
 	auto res = majorantGrid.res;
 	cudaDeviceSynchronize();
 	// [TODO] This device memory is not properly freed
@@ -30,14 +30,14 @@ void initializeMajorantGrid(MajorantGrid& majorantGrid,
 														float(z + 1) / majorantGrid.res.z())));
 
 			// Compute corresponding NanoVDB index-space bounds in floating-point.
-			nanovdb::Vec3R i0 = floatGrid->worldToIndexF(
+			nanovdb::Vec3R i0 = densityGrid->worldToIndexF(
 				nanovdb::Vec3R(wb.min().x(), wb.min().y(), wb.min().z()));
-			nanovdb::Vec3R i1 = floatGrid->worldToIndexF(
+			nanovdb::Vec3R i1 = densityGrid->worldToIndexF(
 				nanovdb::Vec3R(wb.max().x(), wb.max().y(), wb.max().z()));
 
 			// Now find integer index-space bounds, accounting for both
 			// filtering and the overall index bounding box.
-			auto bbox	= floatGrid->indexBBox();
+			auto bbox	= densityGrid->indexBBox();
 			float delta = 1.f; // Filter slop
 			int nx0		= max(int(i0[0] - delta), bbox.min()[0]);
 			int nx1		= min(int(i1[0] + delta), bbox.max()[0]);
@@ -47,33 +47,30 @@ void initializeMajorantGrid(MajorantGrid& majorantGrid,
 			int nz1		= min(int(i1[2] + delta), bbox.max()[2]);
 
 			float maxValue = 0;
-			auto accessor  = floatGrid->getAccessor();
+			auto accessor  = densityGrid->getAccessor();
 
 			for (int nz = nz0; nz <= nz1; ++nz)
 				for (int ny = ny0; ny <= ny1; ++ny)
 					for (int nx = nx0; nx <= nx1; ++nx)
-						maxValue = max(maxValue, accessor.getValue({nx, ny, nz}));
-
+						if constexpr (std::is_same_v<GridType, nanovdb::FloatGrid>) {
+							maxValue = max(maxValue, accessor.getValue({nx, ny, nz}));
+						} else if constexpr (std::is_same_v<GridType, nanovdb::Vec3fGrid>) {
+							auto value = accessor.getValue({nx, ny, nz});
+							RGB color  = {value[0], value[1], value[2]};
+							RGBUnboundedSpectrum spectrum(color, *KRR_DEFAULT_COLORSPACE_GPU);
+							maxValue = max(maxValue, spectrum.maxValue());
+						} else {
+							static_assert(false, "Unsupported grid type!");
+						}
 			majorantGrid.set(x, y, z, maxValue);
 		}, 0);
 }
 
-NanoVDBMedium::NanoVDBMedium(const Affine3f &transform, RGB sigma_a, RGB sigma_s, float g,
-							 NanoVDBGrid density, NanoVDBGrid temperature, float LeScale,
-							 float temperatureScale, float temperatureOffset, 
-							 const RGBColorSpace *colorSpace) :
-	transform(transform), phase(g), sigma_a(sigma_a), sigma_s(sigma_s), densityGrid(std::move(density)), 
-	temperatureGrid(std::move(temperature)), LeScale(LeScale), temperatureScale(temperatureScale), 
-	temperatureOffset(temperatureOffset), colorSpace(colorSpace) {
-	inverseTransform = transform.inverse();
-	const Vector3f majorantGridRes{64, 64, 64};
-	majorantGrid	 = MajorantGrid(densityGrid.getBounds(), majorantGridRes);
-}
-
-void NanoVDBMedium::initializeFromHost() {
-	densityGrid.toDevice();
-	initializeMajorantGrid(majorantGrid, densityGrid.getFloatGrid());
-}
+// explicit instantiation of enable data types of vdb grids
+template void initializeMajorantGrid<nanovdb::FloatGrid>(MajorantGrid &majorantGrid,
+														 nanovdb::FloatGrid *densityGrid);
+template void initializeMajorantGrid<nanovdb::Vec3fGrid>(MajorantGrid &majorantGrid,
+														 nanovdb::Vec3fGrid *densityGrid);
 
 KRR_HOST_DEVICE PhaseFunctionSample HGPhaseFunction::sample(const Vector3f &wo,
 														 const Vector2f &u) const {