diff --git a/src/fvdb/detail/io/LoadNanovdb.cpp b/src/fvdb/detail/io/LoadNanovdb.cpp index 55d55a6b..fe34fca5 100644 --- a/src/fvdb/detail/io/LoadNanovdb.cpp +++ b/src/fvdb/detail/io/LoadNanovdb.cpp @@ -198,7 +198,8 @@ nanovdbTensorGridToFVDBGrid(const nanovdb::NanoGrid *sourceGrid) { copyIndexGridToHandle(sourceGrid); // Check if this grid has FVDB blind data attached to it - bool foundFVDB = false; + bool foundFVDB = false; + unsigned fvdbBlindIndex = 0; torch::Dtype blindDtype; for (unsigned i = 0; i < sourceGrid->blindDataCount(); i += 1) { const nanovdb::GridBlindMetaData &blindMetadata = sourceGrid->blindMetaData(i); @@ -206,16 +207,16 @@ nanovdbTensorGridToFVDBGrid(const nanovdb::NanoGrid *sourceGrid) { if (blindMetadata.mDataClass == nanovdb::GridBlindDataClass::GridName) { continue; } - std::tuple> isFvdb = - isFvdbBlindData(sourceGrid->blindMetaData(0)); + std::tuple> isFvdb = isFvdbBlindData(blindMetadata); if (std::get<0>(isFvdb)) { TORCH_CHECK( !foundFVDB, "Internal Error: Grid has multiple FVDB blind data tensors. Only one is supported."); TORCH_CHECK(std::get<1>(isFvdb).has_value(), "Invalid blind metadata for nanovdb Tensor grid."); - foundFVDB = true; - blindDtype = std::get<1>(isFvdb).value(); + foundFVDB = true; + fvdbBlindIndex = i; + blindDtype = std::get<1>(isFvdb).value(); } else { TORCH_WARN( "Grid has blind data, but it is not valid FVDB blind data. Blind data will be ignored."); @@ -233,7 +234,7 @@ nanovdbTensorGridToFVDBGrid(const nanovdb::NanoGrid *sourceGrid) { } // Pointer to actual blind data - uint8_t *readHead = (uint8_t *)(sourceGrid->blindMetaData(0).blindData()); + uint8_t *readHead = (uint8_t *)(sourceGrid->blindMetaData(fvdbBlindIndex).blindData()); // Read the shape of the tensor const int64_t ndim = *reinterpret_cast(readHead); @@ -363,11 +364,11 @@ nanovdbGridToFvdbGrid(const nanovdb::NanoGrid *sourceGrid) { "Invalid FVDB blind metadata for nanovdb grid. Should not have extra type."); // Pointer to actual blind data - uint8_t *readHead = (uint8_t *)(sourceGrid->blindMetaData(0).blindData()); + uint8_t *readHead = (uint8_t *)(blindMetadata.blindData()); // Read the shape of the tensor const int64_t ndim = *reinterpret_cast(readHead); - TORCH_CHECK(sourceGrid->blindMetaData(0).blindDataSize() == + TORCH_CHECK(blindMetadata.blindDataSize() == nanovdb::math::AlignUp<32U>(sizeof(int64_t) * (ndim + 1)), "Invalid FVDB blind data for nanovdb grid. Unexpected size."); readHead += sizeof(int64_t); diff --git a/src/tests/CMakeLists.txt b/src/tests/CMakeLists.txt index 25a2a6c0..b8f13591 100644 --- a/src/tests/CMakeLists.txt +++ b/src/tests/CMakeLists.txt @@ -137,6 +137,7 @@ endfunction() # Configure an example test ConfigureTest(ExampleTest "ExampleTest.cpp") +ConfigureTest(LoadNanovdbTest "../../tests/cpp/LoadNanovdbTest.cpp") # Configure unit tests ConfigureTest(JaggedTensorTest "JaggedTensorTest.cpp") diff --git a/tests/cpp/LoadNanovdbTest.cpp b/tests/cpp/LoadNanovdbTest.cpp new file mode 100644 index 00000000..512e4e2c --- /dev/null +++ b/tests/cpp/LoadNanovdbTest.cpp @@ -0,0 +1,204 @@ +// Copyright Contributors to the OpenVDB Project +// SPDX-License-Identifier: Apache-2.0 + +#include +#include + +#include + +#include + +#include + +#include +#include +#include +#include +#include + +namespace { + +constexpr const char *kGridName = "tensor-grid-with-leading-blind-name"; + +void +copyFixedString(char *target, const size_t maxSize, const std::string &source) { + TORCH_CHECK_VALUE(source.size() < maxSize, "String is too long for fixed NanoVDB buffer."); + std::memset(target, 0, maxSize); + std::memcpy(target, source.data(), source.size()); +} + +nanovdb::GridHandle +prependGridNameBlindData(nanovdb::GridHandle &sourceHandle, + const std::string &gridName) { + constexpr uint64_t metadataBytes = sizeof(nanovdb::GridBlindMetaData); + + const nanovdb::GridData *sourceGridData = sourceHandle.gridData(0); + TORCH_CHECK(sourceGridData != nullptr, "Expected a valid source grid."); + TORCH_CHECK(sourceGridData->mBlindMetadataCount == 1, + "Test fixture expects exactly one FVDB blind metadata entry."); + + const uint8_t *sourceBytes = static_cast(sourceHandle.buffer().data()); + const uint64_t oldGridBytes = sourceGridData->mGridSize; + const uint64_t blindMetadataOffset = + static_cast(sourceGridData->mBlindMetadataOffset); + const nanovdb::GridBlindMetaData *oldBlindMetadata = + reinterpret_cast(sourceBytes + blindMetadataOffset); + const uint64_t oldBlindDataOffset = + blindMetadataOffset + static_cast(oldBlindMetadata->mDataOffset); + const uint64_t oldBlindDataBytes = oldGridBytes - oldBlindDataOffset; + + const uint64_t gridNameBytes = gridName.size() + 1; + const uint64_t paddedGridNameBytes = nanovdb::math::AlignUp<32UL>(gridNameBytes); + const uint64_t newGridBytes = oldGridBytes + metadataBytes + paddedGridNameBytes; + + nanovdb::HostBuffer outBuffer(newGridBytes); + uint8_t *outBytes = static_cast(outBuffer.data()); + std::memset(outBytes, 0, newGridBytes); + + std::memcpy(outBytes, sourceBytes, blindMetadataOffset); + + nanovdb::GridData *outGridData = reinterpret_cast(outBytes); + outGridData->mGridSize = newGridBytes; + outGridData->mBlindMetadataCount = 2; + outGridData->mBlindMetadataOffset = static_cast(blindMetadataOffset); + + nanovdb::GridBlindMetaData *gridNameMetadata = + reinterpret_cast(outBytes + blindMetadataOffset); + gridNameMetadata->mDataOffset = static_cast(2 * metadataBytes); + gridNameMetadata->mValueCount = paddedGridNameBytes; + gridNameMetadata->mValueSize = 1; + gridNameMetadata->mSemantic = nanovdb::GridBlindDataSemantic::Unknown; + gridNameMetadata->mDataClass = nanovdb::GridBlindDataClass::GridName; + gridNameMetadata->mDataType = nanovdb::GridType::Unknown; + copyFixedString(gridNameMetadata->mName, nanovdb::GridBlindMetaData::MaxNameSize, "grid_name"); + + nanovdb::GridBlindMetaData *fvdbMetadata = reinterpret_cast( + outBytes + blindMetadataOffset + metadataBytes); + *fvdbMetadata = *oldBlindMetadata; + fvdbMetadata->mDataOffset = + static_cast(oldBlindMetadata->mDataOffset + paddedGridNameBytes); + + uint8_t *gridNameData = outBytes + blindMetadataOffset + 2 * metadataBytes; + std::memcpy(gridNameData, gridName.c_str(), gridNameBytes); + + uint8_t *fvdbBlindData = gridNameData + paddedGridNameBytes; + std::memcpy(fvdbBlindData, sourceBytes + oldBlindDataOffset, oldBlindDataBytes); + + return nanovdb::GridHandle(std::move(outBuffer)); +} + +nanovdb::GridHandle +makeTensorGridBlindDataHandle(const fvdb::GridBatchData &gridBatchData, + const torch::Tensor &sourceData, + const std::string &gridName) { + TORCH_CHECK(gridBatchData.batchSize() == 1, "Test fixture expects a single source grid."); + TORCH_CHECK(gridBatchData.device().is_cpu(), "Test fixture expects a CPU source grid."); + TORCH_CHECK(sourceData.device().is_cpu(), "Test fixture expects CPU tensor data."); + TORCH_CHECK(sourceData.scalar_type() == torch::kFloat32, + "Test fixture expects float32 tensor data."); + TORCH_CHECK(sourceData.dim() >= 1, "Test fixture expects at least one tensor dimension."); + TORCH_CHECK(gridBatchData.numVoxelsAt(0) == sourceData.size(0), + "Test fixture expects one tensor row per voxel."); + + torch::Tensor contiguousData = sourceData.contiguous(); + + const nanovdb::GridData *sourceGridData = gridBatchData.nanoGridHandle().gridData(0); + TORCH_CHECK(sourceGridData != nullptr, "Expected a valid source grid."); + + constexpr uint64_t metadataBytes = sizeof(nanovdb::GridBlindMetaData); + const uint64_t sourceGridBytes = sourceGridData->mGridSize; + const uint64_t shapeBytes = sizeof(int64_t) * static_cast(contiguousData.dim() + 1); + const uint64_t tensorBytes = + static_cast(contiguousData.numel() * contiguousData.element_size()); + const uint64_t blindDataBytes = shapeBytes + tensorBytes; + const uint64_t paddedBlindDataBytes = nanovdb::math::AlignUp<32UL>(blindDataBytes); + const uint64_t totalBytes = sourceGridBytes + metadataBytes + paddedBlindDataBytes; + + nanovdb::HostBuffer outBuffer(totalBytes); + uint8_t *outBytes = static_cast(outBuffer.data()); + std::memset(outBytes, 0, totalBytes); + std::memcpy(outBytes, gridBatchData.nanoGridHandle().buffer().data(), sourceGridBytes); + + nanovdb::GridData *outGridData = reinterpret_cast(outBytes); + outGridData->mGridSize = totalBytes; + outGridData->mGridClass = nanovdb::GridClass::TensorGrid; + outGridData->mGridType = nanovdb::GridType::OnIndex; + outGridData->mBlindMetadataCount = 1; + outGridData->mBlindMetadataOffset = static_cast(sourceGridBytes); + copyFixedString(outGridData->mGridName, nanovdb::GridData::MaxNameSize, gridName); + + nanovdb::GridBlindMetaData *blindMetadata = + reinterpret_cast(outBytes + sourceGridBytes); + blindMetadata->mDataOffset = static_cast(metadataBytes); + blindMetadata->mValueCount = paddedBlindDataBytes; + blindMetadata->mValueSize = 1; + blindMetadata->mSemantic = nanovdb::GridBlindDataSemantic::Unknown; + blindMetadata->mDataClass = nanovdb::GridBlindDataClass::Unknown; + blindMetadata->mDataType = nanovdb::GridType::Unknown; + copyFixedString( + blindMetadata->mName, nanovdb::GridBlindMetaData::MaxNameSize, "fvdb_jdatafloat32"); + + uint8_t *writeHead = outBytes + sourceGridBytes + metadataBytes; + *reinterpret_cast(writeHead) = contiguousData.dim(); + writeHead += sizeof(int64_t); + for (int64_t di = 0; di < contiguousData.dim(); ++di) { + *reinterpret_cast(writeHead) = contiguousData.size(di); + writeHead += sizeof(int64_t); + } + std::memcpy(writeHead, contiguousData.data_ptr(), tensorBytes); + + return nanovdb::GridHandle(std::move(outBuffer)); +} + +c10::intrusive_ptr +makeTestGrid() { + torch::Tensor ijk = torch::tensor({{0, 0, 0}, {1, 0, 0}, {0, 1, 0}, {0, 0, 1}}, + torch::TensorOptions().dtype(torch::kInt32)); + fvdb::JaggedTensor jaggedIJK(std::vector{ijk}); + return fvdb::gridbatch_from_ijk( + jaggedIJK, {nanovdb::Vec3d(1.0, 1.0, 1.0)}, {nanovdb::Vec3d(0.0, 0.0, 0.0)}); +} + +void +expectRoundTripWithLeadingGridNameBlindData(nanovdb::GridHandle &sourceHandle, + const torch::Tensor &sourceData) { + nanovdb::GridHandle patchedHandle = + prependGridNameBlindData(sourceHandle, kGridName); + + const nanovdb::GridData *gridData = patchedHandle.gridData(0); + ASSERT_EQ(gridData->mBlindMetadataCount, 2u); + const auto *loadedGrid = patchedHandle.grid(0); + if (loadedGrid != nullptr) { + ASSERT_EQ(loadedGrid->blindMetaData(0).mDataClass, nanovdb::GridBlindDataClass::GridName); + ASSERT_EQ(std::string(loadedGrid->blindMetaData(1).mName).rfind("fvdb_jdata", 0), 0u); + } + + auto loaded = fvdb::from_nanovdb(patchedHandle); + const fvdb::JaggedTensor &loadedData = std::get<1>(loaded); + const std::vector &loadedNames = std::get<2>(loaded); + + ASSERT_EQ(loadedNames.size(), 1u); + EXPECT_EQ(loadedNames[0], kGridName); + EXPECT_TRUE(torch::equal(loadedData.jdata(), sourceData)); +} + +} // namespace + +TEST(LoadNanovdb, TensorGridBlindDataCanFollowGridNameBlindData) { + torch::Tensor sourceData = + torch::arange(16, torch::TensorOptions().dtype(torch::kFloat32)).reshape({4, 2, 2}); + auto grid = makeTestGrid(); + nanovdb::GridHandle sourceHandle = + makeTensorGridBlindDataHandle(*grid, sourceData, kGridName); + expectRoundTripWithLeadingGridNameBlindData(sourceHandle, sourceData); +} + +TEST(LoadNanovdb, TensorGridShapeBlindDataCanFollowGridNameBlindData) { + torch::Tensor sourceData = + torch::arange(12, torch::TensorOptions().dtype(torch::kFloat32)).reshape({4, 3}); + auto grid = makeTestGrid(); + fvdb::JaggedTensor jaggedData(std::vector{sourceData}); + nanovdb::GridHandle sourceHandle = + fvdb::to_nanovdb(*grid, std::optional(jaggedData), {kGridName}); + expectRoundTripWithLeadingGridNameBlindData(sourceHandle, sourceData); +}