diff --git a/monai/csrc/filtering/bilateral/bilateral.cpp b/monai/csrc/filtering/bilateral/bilateral.cpp new file mode 100644 index 0000000000..2720d312e2 --- /dev/null +++ b/monai/csrc/filtering/bilateral/bilateral.cpp @@ -0,0 +1,49 @@ +/* +Copyright 2020 - 2021 MONAI Consortium +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + http://www.apache.org/licenses/LICENSE-2.0 +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +#include +#include +#include + +#include "bilateral.h" +#include "utils/common_utils.h" + +torch::Tensor BilateralFilter(torch::Tensor input, float spatial_sigma, float color_sigma, bool usePHL) { + torch::Tensor (*filterFunction)(torch::Tensor, float, float); + +#ifdef WITH_CUDA + + if (torch::cuda::is_available() && input.is_cuda()) { + CHECK_CONTIGUOUS_CUDA(input); + + if (input.size(1) > BF_CUDA_MAX_CHANNELS) { + throw std::runtime_error( + "Bilateral filtering not implemented for channel count > " + std::to_string(BF_CUDA_MAX_CHANNELS)); + } + + if (input.dim() - 2 > BF_CUDA_MAX_SPATIAL_DIMENSION) { + throw std::runtime_error( + "Bilateral filtering not implemented for spatial dimension > " + + std::to_string(BF_CUDA_MAX_SPATIAL_DIMENSION)); + } + + filterFunction = usePHL ? &BilateralFilterPHLCuda : &BilateralFilterCuda; + } else { + filterFunction = usePHL ? &BilateralFilterPHLCpu : &BilateralFilterCpu; + } +#else + filterFunction = usePHL ? &BilateralFilterPHLCpu : &BilateralFilterCpu; +#endif + + return filterFunction(input, spatial_sigma, color_sigma); +} diff --git a/monai/csrc/filtering/bilateral/bilateral.h b/monai/csrc/filtering/bilateral/bilateral.h index 1c16373fa9..c7a68d7457 100644 --- a/monai/csrc/filtering/bilateral/bilateral.h +++ b/monai/csrc/filtering/bilateral/bilateral.h @@ -14,7 +14,9 @@ limitations under the License. #pragma once #include -#include "utils/common_utils.h" + +#define BF_CUDA_MAX_CHANNELS 16 +#define BF_CUDA_MAX_SPATIAL_DIMENSION 3 torch::Tensor BilateralFilterCpu(torch::Tensor input, float spatial_sigma, float color_sigma); torch::Tensor BilateralFilterPHLCpu(torch::Tensor input, float spatial_sigma, float color_sigma); @@ -24,19 +26,4 @@ torch::Tensor BilateralFilterCuda(torch::Tensor input, float spatial_sigma, floa torch::Tensor BilateralFilterPHLCuda(torch::Tensor input, float spatial_sigma, float color_sigma); #endif -torch::Tensor BilateralFilter(torch::Tensor input, float spatial_sigma, float color_sigma, bool usePHL) { - torch::Tensor (*filterFunction)(torch::Tensor, float, float); - -#ifdef WITH_CUDA - if (torch::cuda::is_available() && input.is_cuda()) { - CHECK_CONTIGUOUS_CUDA(input); - filterFunction = usePHL ? &BilateralFilterPHLCuda : &BilateralFilterCuda; - } else { - filterFunction = usePHL ? &BilateralFilterPHLCpu : &BilateralFilterCpu; - } -#else - filterFunction = usePHL ? &BilateralFilterPHLCpu : &BilateralFilterCpu; -#endif - - return filterFunction(input, spatial_sigma, color_sigma); -} +torch::Tensor BilateralFilter(torch::Tensor input, float spatial_sigma, float color_sigma, bool usePHL); diff --git a/monai/csrc/filtering/bilateral/bilateralfilter_cuda.cu b/monai/csrc/filtering/bilateral/bilateralfilter_cuda.cu index 4477ce5845..f73ae19ac9 100644 --- a/monai/csrc/filtering/bilateral/bilateralfilter_cuda.cu +++ b/monai/csrc/filtering/bilateral/bilateralfilter_cuda.cu @@ -15,6 +15,7 @@ limitations under the License. #include #include +#include "bilateral.h" #include "utils/meta_macros.h" #include "utils/tensor_description.h" @@ -253,7 +254,7 @@ torch::Tensor BilateralFilterCuda(torch::Tensor inputTensor, float spatialSigma, torch::Tensor outputTensor = torch::zeros_like(inputTensor); #define CASE(c, d) BilateralFilterCuda(inputTensor, outputTensor, spatialSigma, colorSigma); - SWITCH_AB(CASE, 16, 3, inputTensor.size(1), inputTensor.dim() - 2); + SWITCH_AB(CASE, BF_CUDA_MAX_CHANNELS, BF_CUDA_MAX_SPATIAL_DIMENSION, inputTensor.size(1), inputTensor.dim() - 2); return outputTensor; } diff --git a/monai/csrc/filtering/bilateral/bilateralfilter_cuda_phl.cu b/monai/csrc/filtering/bilateral/bilateralfilter_cuda_phl.cu index 17dc9e7ebd..719d1643d3 100644 --- a/monai/csrc/filtering/bilateral/bilateralfilter_cuda_phl.cu +++ b/monai/csrc/filtering/bilateral/bilateralfilter_cuda_phl.cu @@ -15,6 +15,7 @@ limitations under the License. #include #include +#include "bilateral.h" #include "filtering/permutohedral/permutohedral.h" #include "utils/meta_macros.h" #include "utils/tensor_description.h" @@ -135,7 +136,7 @@ torch::Tensor BilateralFilterPHLCuda(torch::Tensor inputTensor, float spatialSig inputTensor, outputTensor, spatialSigma, colorSigma); \ })); - SWITCH_AB(CASE, 16, 3, inputTensor.size(1), inputTensor.dim() - 2); + SWITCH_AB(CASE, BF_CUDA_MAX_CHANNELS, BF_CUDA_MAX_SPATIAL_DIMENSION, inputTensor.size(1), inputTensor.dim() - 2); return outputTensor; } diff --git a/monai/csrc/filtering/permutohedral/permutohedral.cpp b/monai/csrc/filtering/permutohedral/permutohedral.cpp index 04ef6fa4da..d8fd3eaaeb 100644 --- a/monai/csrc/filtering/permutohedral/permutohedral.cpp +++ b/monai/csrc/filtering/permutohedral/permutohedral.cpp @@ -11,6 +11,9 @@ See the License for the specific language governing permissions and limitations under the License. */ +#include +#include + #include "utils/common_utils.h" #include "utils/meta_macros.h" @@ -46,6 +49,16 @@ torch::Tensor PermutohedralFilter(torch::Tensor input, torch::Tensor features) { if (torch::cuda::is_available() && data.is_cuda()) { CHECK_CONTIGUOUS_CUDA(data); + if (channelCount > PHL_CUDA_MAX_CHANNELS) { + throw std::runtime_error( + "PHL filtering not implemented for channel count > " + std::to_string(PHL_CUDA_MAX_CHANNELS)); + } + + if (featureCount > PHL_CUDA_MAX_FEATURES) { + throw std::runtime_error( + "PHL filtering not implemented for feature count > " + std::to_string(PHL_CUDA_MAX_FEATURES)); + } + #define CASE(dc, fc) \ AT_DISPATCH_FLOATING_TYPES(data.scalar_type(), "PermutohedralCuda", ([&] { \ for (int batchIndex = 0; batchIndex < batchCount; batchIndex++) { \ @@ -55,7 +68,7 @@ torch::Tensor PermutohedralFilter(torch::Tensor input, torch::Tensor features) { PermutohedralCuda(offsetData, offsetFeatures, elementCount, true); \ } \ })); - SWITCH_AB(CASE, 16, 19, channelCount, featureCount); + SWITCH_AB(CASE, PHL_CUDA_MAX_CHANNELS, PHL_CUDA_MAX_FEATURES, channelCount, featureCount); } else { #endif diff --git a/monai/csrc/filtering/permutohedral/permutohedral.h b/monai/csrc/filtering/permutohedral/permutohedral.h index 27b0ff4859..32ffee83e5 100644 --- a/monai/csrc/filtering/permutohedral/permutohedral.h +++ b/monai/csrc/filtering/permutohedral/permutohedral.h @@ -11,9 +11,13 @@ See the License for the specific language governing permissions and limitations under the License. */ +#pragma once + #include -#pragma once +#define PHL_CUDA_MAX_CHANNELS 16 +#define PHL_CUDA_MAX_FEATURES 19 + template void PermutohedralCPU(scalar_t* data, scalar_t* features, int dataChannels, int featureChannels, int elementCount); #ifdef WITH_CUDA diff --git a/monai/csrc/filtering/permutohedral/permutohedral_cuda.cu b/monai/csrc/filtering/permutohedral/permutohedral_cuda.cu index b87a88a84f..5ec3967aa4 100644 --- a/monai/csrc/filtering/permutohedral/permutohedral_cuda.cu +++ b/monai/csrc/filtering/permutohedral/permutohedral_cuda.cu @@ -38,7 +38,7 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. */ -#define BLOCK_SIZE 64 +#define BLOCK_SIZE 32 #include #include @@ -47,6 +47,7 @@ SOFTWARE. #include #include "hash_table.cuh" +#include "permutohedral.h" #include "utils/meta_macros.h" template