Change Probe filter to use CellLocatorChooser

There was some compile errors with CUDA and the CellLocatorGeneral.
Apparently it added enough to the CUDA kernel to cause issues with
constant memory.

vtkm/cont/CellLocatorChooser.h

@@ -10,10 +10,12 @@
 #ifndef vtk_m_cont_CellLocatorChooser_h
 #define vtk_m_cont_CellLocatorChooser_h
 
+#include <vtkm/cont/CastAndCall.h>
 #include <vtkm/cont/CellLocatorRectilinearGrid.h>
 #include <vtkm/cont/CellLocatorTwoLevel.h>
 #include <vtkm/cont/CellLocatorUniformGrid.h>
 #include <vtkm/cont/CellSetStructured.h>
+#include <vtkm/cont/DataSet.h>
 
 namespace vtkm
 {
@@ -29,19 +31,21 @@ struct CellLocatorChooserImpl
   using type = vtkm::cont::CellLocatorTwoLevel;
 };
 
+using UniformArray = vtkm::cont::ArrayHandleUniformPointCoordinates;
+
 template <>
-struct CellLocatorChooserImpl<vtkm::cont::CellSetStructured<3>,
+struct CellLocatorChooserImpl<vtkm::cont::CellSetStructured<3>, UniformArray>
-                              vtkm::cont::ArrayHandleUniformPointCoordinates>
 {
   using type = vtkm::cont::CellLocatorUniformGrid;
 };
 
-template <>
+using RectilinearArray =
-struct CellLocatorChooserImpl<
-  vtkm::cont::CellSetStructured<3>,
   vtkm::cont::ArrayHandleCartesianProduct<vtkm::cont::ArrayHandle<vtkm::FloatDefault>,
                                           vtkm::cont::ArrayHandle<vtkm::FloatDefault>,
-                                          vtkm::cont::ArrayHandle<vtkm::FloatDefault>>>
+                                          vtkm::cont::ArrayHandle<vtkm::FloatDefault>>;
+
+template <>
+struct CellLocatorChooserImpl<vtkm::cont::CellSetStructured<3>, RectilinearArray>
 {
   using type = vtkm::cont::CellLocatorRectilinearGrid;
 };
@@ -59,6 +63,101 @@ template <typename CellSetType, typename CoordinateSystemArrayType>
 using CellLocatorChooser =
   typename detail::CellLocatorChooserImpl<CellSetType, CoordinateSystemArrayType>::type;
 
+namespace detail
+{
+
+struct CastAndCallCellLocatorChooserFunctor
+{
+  template <typename CellLocatorType, typename Functor, typename... Args>
+  void CallFunctorWithLocator(const vtkm::cont::DynamicCellSet& cellSet,
+                              const vtkm::cont::CoordinateSystem& coordinateSystem,
+                              Functor&& functor,
+                              Args&&... args) const
+  {
+    CellLocatorType locator;
+    locator.SetCellSet(cellSet);
+    locator.SetCoordinates(coordinateSystem);
+
+    functor(locator, std::forward<Args>(args)...);
+  }
+
+  template <typename CellSetType, typename Functor, typename... Args>
+  void operator()(const CellSetType& cellSet,
+                  const vtkm::cont::CoordinateSystem& coordinateSystem,
+                  Functor&& functor,
+                  Args&&... args) const
+  {
+    this->CallFunctorWithLocator<vtkm::cont::CellLocatorTwoLevel>(
+      cellSet, coordinateSystem, std::forward<Functor>(functor), std::forward<Args>(args)...);
+  }
+
+  template <typename Functor, typename... Args>
+  void operator()(const vtkm::cont::CellSetStructured<3>& cellSet,
+                  const vtkm::cont::CoordinateSystem& coordinateSystem,
+                  Functor&& functor,
+                  Args&&... args) const
+  {
+    auto coordArray = coordinateSystem.GetData();
+    if (coordArray.IsType<detail::UniformArray>())
+    {
+      this->CallFunctorWithLocator<vtkm::cont::CellLocatorUniformGrid>(
+        cellSet, coordinateSystem, std::forward<Functor>(functor), std::forward<Args>(args)...);
+    }
+    else if (coordArray.IsType<detail::RectilinearArray>())
+    {
+      this->CallFunctorWithLocator<vtkm::cont::CellLocatorRectilinearGrid>(
+        cellSet, coordinateSystem, std::forward<Functor>(functor), std::forward<Args>(args)...);
+    }
+    else
+    {
+      this->CallFunctorWithLocator<vtkm::cont::CellLocatorTwoLevel>(
+        cellSet, coordinateSystem, std::forward<Functor>(functor), std::forward<Args>(args)...);
+    }
+  }
+};
+
+} // namespace detail
+
+/// \brief Calls a functor with the appropriate type of `CellLocator`.
+///
+/// Given a cell set and a coordinate system of unknown types, calls a functor with an appropriate
+/// CellLocator of the given type. The CellLocator is populated with the provided cell set and
+/// coordinate system.
+///
+/// Any additional args are passed to the functor.
+///
+template <typename CellSetList, typename Functor, typename... Args>
+VTKM_CONT void CastAndCallCellLocatorChooser(
+  const vtkm::cont::DynamicCellSetBase<CellSetList>& cellSet,
+  const vtkm::cont::CoordinateSystem& coordinateSystem,
+  Functor&& functor,
+  Args&&... args)
+{
+  vtkm::cont::CastAndCall(cellSet,
+                          detail::CastAndCallCellLocatorChooserFunctor{},
+                          coordinateSystem,
+                          std::forward<Functor>(functor),
+                          std::forward<Args>(args)...);
+}
+
+/// \brief Calls a functor with the appropriate type of `CellLocator`.
+///
+/// Given a `DataSet`, calls a functor with an appropriate CellLocator of the given type.
+/// The CellLocator is populated with the provided cell set and coordinate system.
+///
+/// Any additional args are passed to the functor.
+///
+template <typename Functor, typename... Args>
+VTKM_CONT void CastAndCallCellLocatorChooser(const vtkm::cont::DataSet& dataSet,
+                                             Functor&& functor,
+                                             Args&&... args)
+{
+  CastAndCallCellLocatorChooser(dataSet.GetCellSet(),
+                                dataSet.GetCoordinateSystem(),
+                                std::forward<Functor>(functor),
+                                std::forward<Args>(args)...);
+}
+
 }
 } // namespace vtkm::cont
 

vtkm/cont/CellLocatorGeneral.cxx

@@ -83,6 +83,7 @@ CellLocatorGeneral::ExecObjType CellLocatorGeneral::PrepareForExecution(
   vtkm::cont::DeviceAdapterId device,
   vtkm::cont::Token& token) const
 {
+  this->Update();
   return this->LocatorImpl.CastAndCall(PrepareFunctor{}, device, token);
 }
 

vtkm/cont/CellLocatorRectilinearGrid.cxx

@@ -54,6 +54,8 @@ vtkm::exec::CellLocatorRectilinearGrid CellLocatorRectilinearGrid::PrepareForExe
   vtkm::cont::DeviceAdapterId device,
   vtkm::cont::Token& token) const
 {
+  this->Update();
+
   using ExecObjType = vtkm::exec::CellLocatorRectilinearGrid;
 
   if (this->Is3D)

vtkm/cont/CellLocatorTwoLevel.cxx

@@ -482,6 +482,7 @@ CellLocatorTwoLevel::ExecObjType CellLocatorTwoLevel::PrepareForExecution(
   vtkm::cont::DeviceAdapterId device,
   vtkm::cont::Token& token) const
 {
+  this->Update();
   ExecObjType execObject;
   this->GetCellSet().CastAndCall(MakeExecObject{}, device, token, *this, execObject);
   return execObject;

vtkm/cont/CellLocatorUniformGrid.cxx

@@ -72,6 +72,7 @@ vtkm::exec::CellLocatorUniformGrid CellLocatorUniformGrid::PrepareForExecution(
   vtkm::cont::DeviceAdapterId vtkmNotUsed(device),
   vtkm::cont::Token& vtkmNotUsed(token)) const
 {
+  this->Update();
   return vtkm::exec::CellLocatorUniformGrid(
     this->CellDims, this->Origin, this->InvSpacing, this->MaxPoint);
 }

vtkm/cont/internal/CellLocatorBase.h

@@ -138,7 +138,7 @@ class VTKM_ALWAYS_EXPORT CellLocatorBase : public vtkm::cont::ExecutionObjectBas
 {
   vtkm::cont::DynamicCellSet CellSet;
   vtkm::cont::CoordinateSystem Coords;
-  bool Modified = true;
+  mutable bool Modified = true;
 
 public:
 #ifndef VTKM_NO_DEPRECATED_VIRTUAL
@@ -167,11 +167,11 @@ public:
     this->SetModified();
   }
 
-  void Update()
+  void Update() const
   {
     if (this->Modified)
     {
-      static_cast<Derived*>(this)->Build();
+      static_cast<Derived*>(const_cast<CellLocatorBase*>(this))->Build();
       this->Modified = false;
     }
   }

vtkm/worklet/Probe.h

@@ -12,13 +12,12 @@
 
 #include <vtkm/cont/ArrayCopy.h>
 #include <vtkm/cont/ArrayHandle.h>
-#include <vtkm/cont/CellLocatorGeneral.h>
+#include <vtkm/cont/CellLocatorChooser.h>
+#include <vtkm/cont/Invoker.h>
 #include <vtkm/exec/CellInside.h>
 #include <vtkm/exec/CellInterpolate.h>
 #include <vtkm/exec/ParametricCoordinates.h>
 
-#include <vtkm/worklet/DispatcherMapField.h>
-#include <vtkm/worklet/DispatcherMapTopology.h>
 #include <vtkm/worklet/WorkletMapField.h>
 #include <vtkm/worklet/WorkletMapTopology.h>
 
@@ -53,6 +52,16 @@ public:
   };
 
 private:
+  struct RunSelectLocator
+  {
+    template <typename LocatorType, typename PointsType>
+    void operator()(const LocatorType& locator, Probe& worklet, const PointsType& points) const
+    {
+      worklet.Invoke(
+        FindCellWorklet{}, points, locator, worklet.CellIds, worklet.ParametricCoordinates);
+    }
+  };
+
   template <typename CellSetType, typename PointsType, typename PointsStorage>
   void RunImpl(const CellSetType& cells,
                const vtkm::cont::CoordinateSystem& coords,
@@ -60,14 +69,7 @@ private:
   {
     this->InputCellSet = vtkm::cont::DynamicCellSet(cells);
 
-    vtkm::cont::CellLocatorGeneral locator;
+    vtkm::cont::CastAndCallCellLocatorChooser(cells, coords, RunSelectLocator{}, *this, points);
-    locator.SetCellSet(this->InputCellSet);
-    locator.SetCoordinates(coords);
-    locator.Update();
-
-    vtkm::worklet::DispatcherMapField<FindCellWorklet> dispatcher;
-    // CellLocatorGeneral is non-copyable. Pass it via a pointer.
-    dispatcher.Invoke(points, &locator, this->CellIds, this->ParametricCoordinates);
   }
 
   //============================================================================
@@ -151,8 +153,8 @@ private:
       this->CellIds);
     this->ParametricCoordinates.Allocate(points.GetNumberOfValues());
 
-    vtkm::worklet::DispatcherMapTopology<ProbeUniformPoints> dispatcher;
+    this->Invoke(
-    dispatcher.Invoke(cells, coords, points, this->CellIds, this->ParametricCoordinates);
+      ProbeUniformPoints{}, cells, coords, points, this->CellIds, this->ParametricCoordinates);
   }
 
   //============================================================================
@@ -298,8 +300,7 @@ public:
   vtkm::cont::ArrayHandle<vtkm::UInt8> GetHiddenPointsField() const
   {
     vtkm::cont::ArrayHandle<vtkm::UInt8> field;
-    vtkm::worklet::DispatcherMapField<HiddenPointsWorklet> dispatcher;
+    this->Invoke(HiddenPointsWorklet{}, this->CellIds, field);
-    dispatcher.Invoke(this->CellIds, field);
     return field;
   }
 
@@ -331,8 +332,7 @@ public:
   vtkm::cont::ArrayHandle<vtkm::UInt8> GetHiddenCellsField(CellSetType cellset) const
   {
     vtkm::cont::ArrayHandle<vtkm::UInt8> field;
-    vtkm::worklet::DispatcherMapTopology<HiddenCellsWorklet> dispatcher;
+    this->Invoke(HiddenCellsWorklet{}, cellset, this->CellIds, field);
-    dispatcher.Invoke(cellset, this->CellIds, field);
     return field;
   }
 
@@ -343,6 +343,8 @@ private:
   vtkm::cont::ArrayHandle<vtkm::Id> CellIds;
   vtkm::cont::ArrayHandle<vtkm::Vec3f> ParametricCoordinates;
   vtkm::cont::DynamicCellSet InputCellSet;
+
+  vtkm::cont::Invoker Invoke;
 };
 }
 } // vtkm::worklet