-
Notifications
You must be signed in to change notification settings - Fork 2.3k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Voxel Downsample for Tensor interface (#6249)
* infra for tensor index-reduction and interface for voxeldownsample * basic sum reduction * temp fix for workload nums * temp fix with contiguous input * separate index tensor * add CPU counterpart * clean up cpp part * add unit test for index_add_ * fix point attribute shape * fix unit tests * fix doc, also fix several doc issues in t.pointcloud * fix benchmark
- Loading branch information
Showing
14 changed files
with
478 additions
and
83 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,49 @@ | ||
| // ---------------------------------------------------------------------------- | ||
| // - Open3D: www.open3d.org - | ||
| // ---------------------------------------------------------------------------- | ||
| // Copyright (c) 2018-2023 www.open3d.org | ||
| // SPDX-License-Identifier: MIT | ||
| // ---------------------------------------------------------------------------- | ||
|
|
||
| #include "open3d/core/kernel/IndexReduction.h" | ||
|
|
||
| #include "open3d/utility/Logging.h" | ||
|
|
||
| namespace open3d { | ||
| namespace core { | ||
| namespace kernel { | ||
|
|
||
| void IndexAdd_(int64_t dim, | ||
| const Tensor& index, | ||
| const Tensor& src, | ||
| Tensor& dst) { | ||
| // Permute the reduction dimension to the first. | ||
| SizeVector permute = {}; | ||
| for (int64_t d = 0; d <= dim; ++d) { | ||
| if (d == 0) { | ||
| permute.push_back(dim); | ||
| } else { | ||
| permute.push_back(d - 1); | ||
| } | ||
| } | ||
| for (int64_t d = dim + 1; d < src.NumDims(); ++d) { | ||
| permute.push_back(d); | ||
| } | ||
|
|
||
| auto src_permute = src.Permute(permute); | ||
| auto dst_permute = dst.Permute(permute); | ||
|
|
||
| if (dst.IsCPU()) { | ||
| IndexAddCPU_(dim, index, src_permute, dst_permute); | ||
| } else if (dst.IsCUDA()) { | ||
| #ifdef BUILD_CUDA_MODULE | ||
| IndexAddCUDA_(dim, index, src_permute, dst_permute); | ||
| #endif | ||
| } else { | ||
| utility::LogError("IndexAdd_: Unimplemented device"); | ||
| } | ||
| } | ||
|
|
||
| } // namespace kernel | ||
| } // namespace core | ||
| } // namespace open3d |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,36 @@ | ||
| // ---------------------------------------------------------------------------- | ||
| // - Open3D: www.open3d.org - | ||
| // ---------------------------------------------------------------------------- | ||
| // Copyright (c) 2018-2023 www.open3d.org | ||
| // SPDX-License-Identifier: MIT | ||
| // ---------------------------------------------------------------------------- | ||
|
|
||
| #pragma once | ||
|
|
||
| #include "open3d/core/Tensor.h" | ||
| #include "open3d/utility/Logging.h" | ||
|
|
||
| namespace open3d { | ||
| namespace core { | ||
| namespace kernel { | ||
|
|
||
| void IndexAdd_(int64_t dim, | ||
| const Tensor& index, | ||
| const Tensor& src, | ||
| Tensor& dst); | ||
|
|
||
| void IndexAddCPU_(int64_t dim, | ||
| const Tensor& index, | ||
| const Tensor& src, | ||
| Tensor& dst); | ||
|
|
||
| #ifdef BUILD_CUDA_MODULE | ||
| void IndexAddCUDA_(int64_t dim, | ||
| const Tensor& index, | ||
| const Tensor& src, | ||
| Tensor& dst); | ||
| #endif | ||
|
|
||
| } // namespace kernel | ||
| } // namespace core | ||
| } // namespace open3d |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,79 @@ | ||
| // ---------------------------------------------------------------------------- | ||
| // - Open3D: www.open3d.org - | ||
| // ---------------------------------------------------------------------------- | ||
| // Copyright (c) 2018-2023 www.open3d.org | ||
| // SPDX-License-Identifier: MIT | ||
| // ---------------------------------------------------------------------------- | ||
|
|
||
| #include "open3d/core/Dispatch.h" | ||
| #include "open3d/core/Indexer.h" | ||
| #include "open3d/core/Tensor.h" | ||
| #include "open3d/utility/Logging.h" | ||
|
|
||
| namespace open3d { | ||
| namespace core { | ||
| namespace kernel { | ||
|
|
||
| template <typename func_t> | ||
| void LaunchIndexReductionKernel(int64_t dim, | ||
| const Device& device, | ||
| const Tensor& index, | ||
| const Tensor& src, | ||
| Tensor& dst, | ||
| const func_t& element_kernel) { | ||
| // index: [N,], src: [N, D], dst: [M, D] | ||
| // In Indexer, output shape defines the actual master strides. | ||
| // However, in IndexAdd_, input dominates the iterations. | ||
| // So put dst (output) at indexer's input, and src (input) at output. | ||
| Indexer indexer({dst}, src, DtypePolicy::NONE); | ||
|
|
||
| // Index is simply a 1D contiguous tensor, with a different stride | ||
| // behavior to src. So use raw pointer for simplicity. | ||
| auto index_ptr = index.GetDataPtr<int64_t>(); | ||
|
|
||
| int64_t broadcasting_elems = 1; | ||
| for (int64_t d = 1; d < src.NumDims(); ++d) { | ||
| broadcasting_elems *= src.GetShape(d); | ||
| } | ||
| auto element_func = [=](int64_t workload_idx) { | ||
| int reduction_idx = workload_idx / broadcasting_elems; | ||
| int broadcasting_idx = workload_idx % broadcasting_elems; | ||
|
|
||
| const int64_t idx = index_ptr[reduction_idx]; | ||
| int64_t dst_idx = idx * broadcasting_elems + broadcasting_idx; | ||
|
|
||
| void* src_ptr = indexer.GetOutputPtr(0, workload_idx); | ||
| void* dst_ptr = indexer.GetInputPtr(0, dst_idx); | ||
| // Note input and output is switched here to adapt to the indexer | ||
| element_kernel(src_ptr, dst_ptr); | ||
| }; | ||
|
|
||
| // TODO: check in detail | ||
| // No OpenMP could be faster, otherwise there would be thousands of atomics. | ||
| for (int64_t d = 0; d < indexer.NumWorkloads(); ++d) { | ||
| element_func(d); | ||
| } | ||
| } | ||
|
|
||
| template <typename scalar_t> | ||
| static OPEN3D_HOST_DEVICE void CPUSumKernel(const void* src, void* dst) { | ||
| scalar_t* dst_s_ptr = static_cast<scalar_t*>(dst); | ||
| const scalar_t* src_s_ptr = static_cast<const scalar_t*>(src); | ||
| *dst_s_ptr += *src_s_ptr; | ||
| } | ||
|
|
||
| void IndexAddCPU_(int64_t dim, | ||
| const Tensor& index, | ||
| const Tensor& src, | ||
| Tensor& dst) { | ||
| DISPATCH_FLOAT_DTYPE_TO_TEMPLATE(src.GetDtype(), [&]() { | ||
| LaunchIndexReductionKernel(dim, src.GetDevice(), index, src, dst, | ||
| [](const void* src, void* dst) { | ||
| CPUSumKernel<scalar_t>(src, dst); | ||
| }); | ||
| }); | ||
| } | ||
|
|
||
| } // namespace kernel | ||
| } // namespace core | ||
| } // namespace open3d |
Oops, something went wrong.