WIP: ENH: Add naive CUDA schedule

KitwareMedical · Sep 4, 2024 · 5383a77 · 5383a77
1 parent 5f77feb
commit 5383a77
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 3 deletions.
diff --git a/include/itkHalideDiscreteGaussianImageFilter.hxx b/include/itkHalideDiscreteGaussianImageFilter.hxx
@@ -22,6 +22,7 @@
 
 #include "itkHalideDiscreteGaussianImpl.h"
 
+#include <Halide.h>
 #include <HalideBuffer.h>
 
 namespace itk
@@ -64,7 +65,9 @@ HalideDiscreteGaussianImageFilter<TInputImage, TOutputImage>::GenerateData()
   Halide::Runtime::Buffer<const InputPixelType> inputBuffer(input->GetBufferPointer(), sizes);
   Halide::Runtime::Buffer<OutputPixelType>      outputBuffer(output->GetBufferPointer(), sizes);
 
+  inputBuffer.set_host_dirty();
   itkHalideDiscreteGaussianImpl(inputBuffer, sigma_x, sigma_y, sigma_z, outputBuffer);
+  outputBuffer.copy_to_host();
 }
 
 } // end namespace itk

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -6,7 +6,8 @@ target_link_libraries(itkHalideGenerators PRIVATE Halide::Generator)
 add_halide_library(itkHalideDiscreteGaussianImpl
   FROM itkHalideGenerators
   HEADER itkHalideDiscreteGaussianImpl_h
-  #  PARAMS input.type=float32 output.type=float32
+  # PARAMS input.type=float32 output.type=float32
+  FEATURES cuda # debug
   )
 
 set(HalideFilters_SRCS

diff --git a/src/generators.cpp b/src/generators.cpp
@@ -61,7 +61,10 @@ class DiscreteGaussianGenerator : public Generator<DiscreteGaussianGenerator>
 
     output(x, y, z) = blur_z(x, y, z);
 
-    output.vectorize(x, 4);
+
+    Var ix, iy, iz, ox, oy, oz;
+
+    output.gpu_tile(x, y, z, ix, iy, iz, ox, oy, oz, 4, 4, 4);
   }
 };
 

diff --git a/test/itkHalideDiscreteGaussianImageFilterTest.cxx b/test/itkHalideDiscreteGaussianImageFilterTest.cxx
@@ -65,7 +65,7 @@ itkHalideDiscreteGaussianImageFilterTest(int argc, char * argv[])
   }
   const char * outputImageFileName = argv[1];
 
-  constexpr unsigned int Dimension = 2;
+  constexpr unsigned int Dimension = 3;
   using PixelType = float;
   using ImageType = itk::Image<PixelType, Dimension>;
 
@@ -85,6 +85,7 @@ itkHalideDiscreteGaussianImageFilterTest(int argc, char * argv[])
   ShowProgress::Pointer showProgress = ShowProgress::New();
   filter->AddObserver(itk::ProgressEvent(), showProgress);
   filter->SetInput(image);
+  filter->SetSigma(2);
 
   using WriterType = itk::ImageFileWriter<ImageType>;
   WriterType::Pointer writer = WriterType::New();