Add ValueCount tag to WorkletReduceByKey

vtkm/exec/arg/CMakeLists.txt

@@ -38,6 +38,7 @@ set(headers
   ThreadIndicesBasic.h
   ThreadIndicesReduceByKey.h
   ThreadIndicesTopologyMap.h
+  ValueCount.h
   VisitIndex.h
   WorkIndex.h
   )

vtkm/exec/arg/ValueCount.h

+//============================================================================
+//  Copyright (c) Kitware, Inc.
+//  All rights reserved.
+//  See LICENSE.txt for details.
+//  This software is distributed WITHOUT ANY WARRANTY; without even
+//  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
+//  PURPOSE.  See the above copyright notice for more information.
+//
+//  Copyright 2014 Sandia Corporation.
+//  Copyright 2014 UT-Battelle, LLC.
+//  Copyright 2014 Los Alamos National Security.
+//
+//  Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
+//  the U.S. Government retains certain rights in this software.
+//
+//  Under the terms of Contract DE-AC52-06NA25396 with Los Alamos National
+//  Laboratory (LANL), the U.S. Government retains certain rights in
+//  this software.
+//============================================================================
+#ifndef vtk_m_exec_arg_ValueCount_h
+#define vtk_m_exec_arg_ValueCount_h
+
+#include <vtkm/exec/arg/Fetch.h>
+#include <vtkm/exec/arg/ExecutionSignatureTagBase.h>
+#include <vtkm/exec/arg/ThreadIndicesReduceByKey.h>
+
+namespace vtkm {
+namespace exec {
+namespace arg {
+
+/// \brief Aspect tag to use for getting the value count.
+///
+/// The \c AspectTagValueCount aspect tag causes the \c Fetch class to obtain
+/// the number of values that map to the key.
+///
+struct AspectTagValueCount {  };
+
+/// \brief The \c ExecutionSignature tag to get the number of values.
+///
+/// A \c WorkletReduceByKey operates by collecting all values associated with
+/// identical keys and then giving the worklet a Vec-like object containing all
+/// values with a matching key. This \c ExecutionSignature tag provides the
+/// number of values associated with the key and given in the Vec-like objects.
+///
+struct ValueCount : vtkm::exec::arg::ExecutionSignatureTagBase
+{
+  static const vtkm::IdComponent INDEX = 1;
+  typedef vtkm::exec::arg::AspectTagValueCount AspectTag;
+};
+
+template<typename FetchTag,
+         typename ExecObjectType>
+struct Fetch<FetchTag,
+             vtkm::exec::arg::AspectTagValueCount,
+             vtkm::exec::arg::ThreadIndicesReduceByKey,
+             ExecObjectType>
+{
+  using ThreadIndicesType = vtkm::exec::arg::ThreadIndicesReduceByKey;
+
+  using ValueType = vtkm::IdComponent;
+
+  VTKM_SUPPRESS_EXEC_WARNINGS
+  VTKM_EXEC
+  ValueType Load(const ThreadIndicesType &indices, const ExecObjectType &) const
+  {
+    return indices.GetNumberOfValues();
+  }
+
+  VTKM_EXEC
+  void Store(const ThreadIndicesType &,
+             const ExecObjectType &,
+             const ValueType &) const
+  {
+    // Store is a no-op.
+  }
+};
+
+}
+}
+} // namespace vtkm::exec::arg
+
+#endif //vtk_m_exec_arg_ValueCount_h

vtkm/worklet/WorkletReduceByKey.h

@@ -36,6 +36,7 @@
 #include <vtkm/exec/arg/FetchTagArrayDirectOut.h>
 #include <vtkm/exec/arg/FetchTagKeysIn.h>
 #include <vtkm/exec/arg/ThreadIndicesReduceByKey.h>
+#include <vtkm/exec/arg/ValueCount.h>
 
 namespace vtkm {
 namespace worklet {
@@ -81,6 +82,8 @@ public:
   /// all values with a matching key. This tag specifies an \c ArrayHandle
   /// object that holds the values.
   ///
+  /// This tag might not work with scatter operations.
+  ///
   template<typename TypeList = AllTypes>
   struct ValuesInOut : vtkm::cont::arg::ControlSignatureTagBase
   {
@@ -96,6 +99,8 @@ public:
   /// all values with a matching key. This tag specifies an \c ArrayHandle
   /// object that holds the values.
   ///
+  /// This tag might not work with scatter operations.
+  ///
   template<typename TypeList = AllTypes>
   struct ValuesOut : vtkm::cont::arg::ControlSignatureTagBase
   {
@@ -104,6 +109,15 @@ public:
     using FetchTag = vtkm::exec::arg::FetchTagArrayDirectIn;
   };
 
+  /// \brief The \c ExecutionSignature tag to get the number of values.
+  ///
+  /// A \c WorkletReduceByKey operates by collecting all values associated with
+  /// identical keys and then giving the worklet a Vec-like object containing all
+  /// values with a matching key. This \c ExecutionSignature tag provides the
+  /// number of values associated with the key and given in the Vec-like objects.
+  ///
+  struct ValueCount : vtkm::exec::arg::ValueCount {  };
+
   /// Reduce by key worklets use the related thread indices class.
   ///
   VTKM_SUPPRESS_EXEC_WARNINGS

vtkm/worklet/testing/UnitTestWorkletReduceByKey.cxx

@@ -56,7 +56,7 @@ struct CheckReduceByKeyWorklet : vtkm::worklet::WorkletReduceByKey
                                 ValuesIn<> indexValues,
                                 ValuesInOut<> valuesToModify,
                                 ValuesOut<> writeKey);
-  typedef void ExecutionSignature(_1, _2, _3, _4, _5, WorkIndex);
+  typedef void ExecutionSignature(_1, _2, _3, _4, _5, WorkIndex, ValueCount);
   typedef _1 InputDomain;
 
   template<typename T,
@@ -70,17 +70,19 @@ struct CheckReduceByKeyWorklet : vtkm::worklet::WorkletReduceByKey
                   const IndexValuesVecType &valueIndices,
                   ValuesToModifyVecType &valuesToModify,
                   WriteKeysVecType &writeKey,
-                  vtkm::Id workIndex) const
+                  vtkm::Id workIndex,
+                  vtkm::IdComponent numValues) const
   {
     // These tests only work if keys are in sorted order, which is how we group
     // them.
 
     TEST_ASSERT_WORKLET(key == TestValue(workIndex, T()));
 
-    vtkm::IdComponent numValues = keyMirror.GetNumberOfComponents();
     TEST_ASSERT_WORKLET(numValues >= GROUP_SIZE);
-    TEST_ASSERT_WORKLET(keyMirror.GetNumberOfComponents() ==
-                        valueIndices.GetNumberOfComponents());
+    TEST_ASSERT_WORKLET(keyMirror.GetNumberOfComponents() == numValues);
+    TEST_ASSERT_WORKLET(valueIndices.GetNumberOfComponents() == numValues);
+    TEST_ASSERT_WORKLET(valuesToModify.GetNumberOfComponents() == numValues);
+    TEST_ASSERT_WORKLET(writeKey.GetNumberOfComponents() == numValues);
 
 
     for (vtkm::IdComponent iComponent = 0; iComponent < numValues; iComponent++)