Add WorkletReduceByKey and dispatcher

And the basic type for a reduce by key worklet and its associated
adapter. Right now the worklet only supports passing in keys. Values
come next.

vtkm/cont/arg/CMakeLists.txt

@@ -27,6 +27,7 @@ set(headers
   TransportTagAtomicArray.h
   TransportTagCellSetIn.h
   TransportTagExecObject.h
+  TransportTagKeysIn.h
   TransportTagTopologyFieldIn.h
   TransportTagWholeArrayIn.h
   TransportTagWholeArrayInOut.h
@@ -36,6 +37,7 @@ set(headers
   TypeCheckTagAtomicArray.h
   TypeCheckTagCellSet.h
   TypeCheckTagExecObject.h
+  TypeCheckTagKeys.h
   )
 
 vtkm_declare_headers(${headers})

vtkm/cont/arg/TransportTagKeysIn.h

@@ -0,0 +1,45 @@
+//============================================================================
+//  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 2016 Sandia Corporation.
+//  Copyright 2016 UT-Battelle, LLC.
+//  Copyright 2016 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_cont_arg_TransportTagKeysIn_h
+#define vtk_m_cont_arg_TransportTagKeysIn_h
+
+#include <vtkm/cont/arg/Transport.h>
+
+namespace vtkm {
+namespace cont {
+namespace arg {
+
+/// \brief \c Transport tag for keys in a reduce by key.
+///
+/// \c TransportTagKeysIn is a tag used with the \c Transport class to
+/// transport vtkm::worklet::Keys objects for the input domain of a
+/// reduce by keys worklet.
+///
+struct TransportTagKeysIn {  };
+
+// Specialization of Transport class for TransportTagKeysIn is implemented in
+// vtkm/worklet/Keys.h. That class is not accessible from here due to VTK-m
+// package dependencies.
+
+}
+}
+} // namespace vtkm::cont::arg
+
+#endif //vtk_m_cont_arg_TransportTagKeysIn_h

vtkm/cont/arg/TypeCheckTagKeys.h

@@ -0,0 +1,48 @@
+//============================================================================
+//  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 2016 Sandia Corporation.
+//  Copyright 2016 UT-Battelle, LLC.
+//  Copyright 2016 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_cont_arg_TypeCheckTagKeys_h
+#define vtk_m_cont_arg_TypeCheckTagKeys_h
+
+#include <vtkm/cont/arg/TypeCheck.h>
+
+namespace vtkm {
+namespace cont {
+namespace arg {
+
+/// Check for a Keys object.
+///
+struct TypeCheckTagKeys
+{
+};
+
+// A more specific specialization that actually checks for Keys types is
+// implemented in vtkm/worklet/Keys.h. That class is not accessible from here
+// due to VTK-m package dependencies.
+template<typename Type>
+struct TypeCheck<TypeCheckTagKeys, Type>
+{
+  static const bool value = false;
+};
+
+}
+}
+} // namespace vtkm::cont::arg
+
+#endif //vtk_m_cont_arg_TypeCheckTagKeys_h

vtkm/cont/arg/testing/CMakeLists.txt

@@ -29,6 +29,7 @@ set(unit_tests
   UnitTestTypeCheckArray.cxx
   UnitTestTypeCheckCellSet.cxx
   UnitTestTypeCheckExecObject.cxx
+  UnitTestTypeCheckKeys.cxx
   )
 
 vtkm_unit_tests(SOURCES ${unit_tests})

vtkm/cont/arg/testing/UnitTestTypeCheckCellSet.cxx

@@ -32,7 +32,7 @@ struct TestNotCellSet {  };
 
 void TestCheckCellSet()
 {
-  std::cout << "Checking reporting of type checking exec object." << std::endl;
+  std::cout << "Checking reporting of type checking cell set." << std::endl;
 
   using vtkm::cont::arg::TypeCheck;
   using vtkm::cont::arg::TypeCheckTagCellSet;

vtkm/cont/arg/testing/UnitTestTypeCheckKeys.cxx

@@ -0,0 +1,66 @@
+//============================================================================
+//  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 2016 Sandia Corporation.
+//  Copyright 2016 UT-Battelle, LLC.
+//  Copyright 2016 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.
+//============================================================================
+
+#include <vtkm/cont/arg/TypeCheckTagKeys.h>
+
+#include <vtkm/worklet/Keys.h>
+
+#include <vtkm/cont/ArrayHandle.h>
+
+#include <vtkm/cont/testing/Testing.h>
+
+namespace {
+
+struct TestNotKeys {  };
+
+void TestCheckKeys()
+{
+  std::cout << "Checking reporting of type checking keys." << std::endl;
+
+  using vtkm::cont::arg::TypeCheck;
+  using vtkm::cont::arg::TypeCheckTagKeys;
+
+  VTKM_TEST_ASSERT(
+        (TypeCheck<TypeCheckTagKeys, vtkm::worklet::Keys<vtkm::Id> >::value),
+        "Type check failed.");
+  VTKM_TEST_ASSERT(
+        (TypeCheck<TypeCheckTagKeys, vtkm::worklet::Keys<vtkm::Float32> >::value),
+        "Type check failed.");
+  VTKM_TEST_ASSERT(
+        (TypeCheck<TypeCheckTagKeys, vtkm::worklet::Keys<vtkm::Id3> >::value),
+        "Type check failed.");
+
+  VTKM_TEST_ASSERT(
+        !(TypeCheck<TypeCheckTagKeys, TestNotKeys>::value),
+        "Type check failed.");
+  VTKM_TEST_ASSERT(
+        !(TypeCheck<TypeCheckTagKeys, vtkm::Id>::value),
+        "Type check failed.");
+  VTKM_TEST_ASSERT(
+        !(TypeCheck<TypeCheckTagKeys, vtkm::cont::ArrayHandle<vtkm::Id> >::value),
+        "Type check failed.");
+}
+
+} // anonymous namespace
+
+int UnitTestTypeCheckKeys(int, char*[])
+{
+  return vtkm::cont::testing::Testing::Run(TestCheckKeys);
+}

vtkm/exec/arg/CMakeLists.txt

@@ -30,11 +30,13 @@ set(headers
   FetchTagArrayTopologyMapIn.h
   FetchTagExecObject.h
   FetchTagCellSetIn.h
+  FetchTagKeysIn.h
   FetchTagWholeCellSetIn.h
   FromCount.h
   FromIndices.h
   ThreadIndices.h
   ThreadIndicesBasic.h
+  ThreadIndicesReduceByKey.h
   ThreadIndicesTopologyMap.h
   VisitIndex.h
   WorkIndex.h

vtkm/exec/arg/FetchTagKeysIn.h

@@ -0,0 +1,78 @@
+//============================================================================
+//  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 2016 Sandia Corporation.
+//  Copyright 2016 UT-Battelle, LLC.
+//  Copyright 2016 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_FetchTagKeysIn_h
+#define vtk_m_exec_arg_FetchTagKeysIn_h
+
+#include <vtkm/exec/arg/AspectTagDefault.h>
+#include <vtkm/exec/arg/Fetch.h>
+#include <vtkm/exec/arg/ThreadIndicesReduceByKey.h>
+
+#include <vtkm/exec/internal/ReduceByKeyLookup.h>
+
+namespace vtkm {
+namespace exec {
+namespace arg {
+
+/// \brief \c Fetch tag for getting key values in a reduce by key.
+///
+/// \c FetchTagKeysIn is a tag used with the \c Fetch class to retrieve keys
+/// from the input domain of a reduce by keys worklet.
+///
+struct FetchTagKeysIn {  };
+
+template<typename KeyPortalType,
+         typename IdPortalType,
+         typename IdComponentPortalType>
+struct Fetch<
+    vtkm::exec::arg::FetchTagKeysIn,
+    vtkm::exec::arg::AspectTagDefault,
+    vtkm::exec::arg::ThreadIndicesReduceByKey,
+    vtkm::exec::internal::ReduceByKeyLookup<
+      KeyPortalType,IdPortalType,IdComponentPortalType> >
+{
+  using ThreadIndicesType = vtkm::exec::arg::ThreadIndicesReduceByKey;
+  using ExecObjectType =
+    vtkm::exec::internal::ReduceByKeyLookup<
+      KeyPortalType,IdPortalType,IdComponentPortalType>;
+
+  using ValueType = typename ExecObjectType::KeyType;
+
+  VTKM_SUPPRESS_EXEC_WARNINGS
+  VTKM_EXEC
+  ValueType Load(const ThreadIndicesType &indices,
+                 const ExecObjectType &keys) const
+  {
+    return keys.UniqueKeys.Get(indices.GetInputIndex());
+  }
+
+  VTKM_EXEC
+  void Store(const ThreadIndicesType &,
+             const ExecObjectType &,
+             const ValueType &) const
+  {
+    // Store is a no-op for this fetch.
+  }
+};
+
+}
+}
+} // namespace vtkm::exec::arg
+
+#endif //vtk_m_exec_arg_FetchTagKeysIn_h

vtkm/exec/arg/ThreadIndicesReduceByKey.h

@@ -0,0 +1,77 @@
+//============================================================================
+//  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 2016 Sandia Corporation.
+//  Copyright 2016 UT-Battelle, LLC.
+//  Copyright 2016 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_ThreadIndicesReduceByKey_h
+#define vtk_m_exec_arg_ThreadIndicesReduceByKey_h
+
+#include <vtkm/exec/arg/ThreadIndicesBasic.h>
+
+#include <vtkm/exec/internal/ReduceByKeyLookup.h>
+
+namespace vtkm {
+namespace exec {
+namespace arg {
+
+/// \brief Container for thread indices in a reduce by key invocation.
+///
+/// This specialization of \c ThreadIndices adds extra indices that deal with a
+/// reduce by key. In particular, it save the indices used to map from a unique
+/// key index to the group of input values that has that key associated with
+/// it.
+///
+class ThreadIndicesReduceByKey : public vtkm::exec::arg::ThreadIndicesBasic
+{
+  using Superclass = vtkm::exec::arg::ThreadIndicesBasic;
+
+public:
+  template<typename P1, typename P2, typename P3>
+  VTKM_EXEC
+  ThreadIndicesReduceByKey(
+      vtkm::Id threadIndex,
+      vtkm::Id inIndex,
+      vtkm::IdComponent visitIndex,
+      const vtkm::exec::internal::ReduceByKeyLookup<P1,P2,P3> &keyLookup,
+      vtkm::Id globalThreadIndexOffset=0)
+    : Superclass(threadIndex, inIndex, visitIndex, globalThreadIndexOffset),
+      ValueOffset(keyLookup.Offsets.Get(inIndex)),
+      NumberOfValues(keyLookup.Counts.Get(inIndex))
+  {  }
+
+  VTKM_EXEC
+  vtkm::Id GetValueOffset() const
+  {
+    return this->ValueOffset;
+  }
+
+  VTKM_EXEC
+  vtkm::IdComponent GetNumberOfValues() const
+  {
+    return this->NumberOfValues;
+  }
+
+private:
+  vtkm::Id ValueOffset;
+  vtkm::IdComponent NumberOfValues;
+};
+
+}
+}
+} // namespace vtkm::exec::arg
+
+#endif //vtk_m_exec_arg_ThreadIndicesReduceByKey_h

vtkm/exec/internal/CMakeLists.txt

@@ -20,6 +20,7 @@
 
 set(headers
   ErrorMessageBuffer.h
+  ReduceByKeyLookup.h
   WorkletInvokeFunctor.h
   WorkletInvokeFunctorDetail.h
   )

vtkm/exec/internal/ReduceByKeyLookup.h

@@ -0,0 +1,77 @@
+//============================================================================
+//  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 2016 Sandia Corporation.
+//  Copyright 2016 UT-Battelle, LLC.
+//  Copyright 2016 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_internal_ReduceByKeyLookup_h
+#define vtk_m_exec_internal_ReduceByKeyLookup_h
+
+#include <vtkm/exec/ExecutionObjectBase.h>
+
+#include <vtkm/StaticAssert.h>
+#include <vtkm/Types.h>
+
+#include <type_traits>
+
+namespace vtkm {
+namespace exec {
+namespace internal {
+
+/// \breif Execution object holding lookup info for reduce by key.
+///
+/// A WorkletReduceByKey needs several arrays to map the current output object
+/// to the respective key and group of values. This execution object holds that
+/// state.
+///
+template<typename KeyPortalType,
+         typename IdPortalType,
+         typename IdComponentPortalType>
+struct ReduceByKeyLookup : vtkm::exec::ExecutionObjectBase
+{
+  using KeyType = typename KeyPortalType::ValueType;
+
+  VTKM_STATIC_ASSERT(
+      (std::is_same<typename IdPortalType::ValueType, vtkm::Id>::value));
+  VTKM_STATIC_ASSERT(
+      (std::is_same<typename IdComponentPortalType::ValueType, vtkm::IdComponent>::value));
+
+  KeyPortalType UniqueKeys;
+  IdPortalType SortedValuesMap;
+  IdPortalType Offsets;
+  IdComponentPortalType Counts;
+
+  VTKM_EXEC_CONT
+  ReduceByKeyLookup(const KeyPortalType &uniqueKeys,
+                    const IdPortalType &sortedValuesMap,
+                    const IdPortalType &offsets,
+                    const IdComponentPortalType &counts)
+    : UniqueKeys(uniqueKeys),
+      SortedValuesMap(sortedValuesMap),
+      Offsets(offsets),
+      Counts(counts)
+  {  }
+
+  VTKM_EXEC_CONT
+  ReduceByKeyLookup()
+  {  }
+};
+
+}
+}
+} // namespace vtkm::exec::internal
+
+#endif //vtk_m_exec_internal_ReduceByKeyLookup_h

vtkm/worklet/CMakeLists.txt

@@ -25,6 +25,7 @@ set(headers
   Clip.h
   DispatcherMapField.h
   DispatcherMapTopology.h
+  DispatcherReduceByKey.h
   DispatcherStreamingMapField.h
   ExternalFaces.h
   FieldHistogram.h
@@ -51,6 +52,7 @@ set(headers
   WaveletCompressor.h
   WorkletMapField.h
   WorkletMapTopology.h
+  WorkletReduceByKey.h
   )
 
 #-----------------------------------------------------------------------------

vtkm/worklet/DispatcherReduceByKey.h

@@ -0,0 +1,80 @@
+//============================================================================
+//  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_worklet_DispatcherReduceByKey_h
+#define vtk_m_worklet_DispatcherReduceByKey_h
+
+#include <vtkm/cont/DeviceAdapter.h>
+
+#include <vtkm/worklet/WorkletReduceByKey.h>
+
+#include <vtkm/worklet/internal/DispatcherBase.h>
+
+namespace vtkm {
+namespace worklet {
+
+/// \brief Dispatcher for worklets that inherit from \c WorkletReduceByKey.
+///
+template<typename WorkletType,
+         typename Device = VTKM_DEFAULT_DEVICE_ADAPTER_TAG>
+class DispatcherReduceByKey
+    : public vtkm::worklet::internal::DispatcherBase<
+        DispatcherReduceByKey<WorkletType, Device>,
+        WorkletType,
+        vtkm::worklet::WorkletReduceByKey>
+{
+  using Superclass =
+    vtkm::worklet::internal::DispatcherBase<
+      DispatcherReduceByKey<WorkletType, Device>,
+      WorkletType,
+      vtkm::worklet::WorkletReduceByKey>;
+
+public:
+  VTKM_CONT
+  DispatcherReduceByKey(const WorkletType &worklet = WorkletType())
+    : Superclass(worklet)
+  {  }
+
+  template<typename Invocation>
+  void DoInvoke(const Invocation &invocation) const
+  {
+    // This is the type for the input domain
+    using InputDomainType = typename Invocation::InputDomainType;
+
+    // If you get a compile error on this line, then you have tried to use
+    // something other than vtkm::worklet::Keys as the input domain, which
+    // is illegal.
+    VTKM_STATIC_ASSERT_MSG(
+          (vtkm::cont::arg::TypeCheck<vtkm::cont::arg::TypeCheckTagKeys,InputDomainType>::value),
+          "Invalid input domain for WorkletReduceByKey.");
+
+    // We can pull the input domain parameter (the data specifying the input
+    // domain) from the invocation object.
+    const InputDomainType &inputDomain = invocation.GetInputDomain();
+
+    // Now that we have the input domain, we can extract the range of the
+    // scheduling and call BadicInvoke.
+    this->BasicInvoke(invocation, inputDomain.GetInputRange(), Device());
+  }
+};
+
+}
+} // namespace vtkm::worklet
+
+#endif //vtk_m_worklet_DispatcherReduceByKey_h

vtkm/worklet/Keys.h

@@ -6,9 +6,9 @@
 //  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.
+//  Copyright 2016 Sandia Corporation.
+//  Copyright 2016 UT-Battelle, LLC.
+//  Copyright 2016 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.
@@ -26,6 +26,11 @@
 #include <vtkm/cont/ArrayHandleIndex.h>
 #include <vtkm/cont/DeviceAdapterAlgorithm.h>
 
+#include <vtkm/exec/internal/ReduceByKeyLookup.h>
+
+#include <vtkm/cont/arg/TransportTagKeysIn.h>
+#include <vtkm/cont/arg/TypeCheckTagKeys.h>
+
 #include <vtkm/BinaryOperators.h>
 
 namespace vtkm {
@@ -48,16 +53,20 @@ namespace worklet {
 /// Keys structure is reused for all the \c Invoke. This is more efficient than
 /// creating a different \c Keys structure for each \c Invoke.
 ///
-template<typename _ValueType>
+template<typename _KeyType>
 class Keys
 {
 public:
-  using ValueType = _ValueType;
-  using KeyArrayHandleType = vtkm::cont::ArrayHandle<ValueType>;
+  using KeyType = _KeyType;
+  using KeyArrayHandleType = vtkm::cont::ArrayHandle<KeyType>;
+
+  VTKM_CONT
+  Keys()
+  {  }
 
   template<typename OriginalKeyStorage, typename Device>
   VTKM_CONT
-  Keys(const vtkm::cont::ArrayHandle<ValueType,OriginalKeyStorage> &keys,
+  Keys(const vtkm::cont::ArrayHandle<KeyType,OriginalKeyStorage> &keys,
        Device)
   {
     this->BuildArrays(keys, Device());
@@ -93,6 +102,36 @@ public:
     return this->Counts;
   }
 
+  template<typename Device>
+  struct ExecutionTypes
+  {
+    using KeyPortal = typename KeyArrayHandleType::template ExecutionTypes<Device>::PortalConst;
+    using IdPortal = typename vtkm::cont::ArrayHandle<vtkm::Id>::template ExecutionTypes<Device>::PortalConst;
+    using IdComponentPortal = typename vtkm::cont::ArrayHandle<vtkm::IdComponent>::template ExecutionTypes<Device>::PortalConst;
+
+    using Lookup = vtkm::exec::internal::ReduceByKeyLookup<KeyPortal,IdPortal,IdComponentPortal>;
+  };
+
+  template<typename Device>
+  VTKM_CONT
+  typename ExecutionTypes<Device>::Lookup PrepareForInput(Device) const
+  {
+    return typename ExecutionTypes<Device>::Lookup(
+          this->UniqueKeys.PrepareForInput(Device()),
+          this->SortedValuesMap.PrepareForInput(Device()),
+          this->Offsets.PrepareForInput(Device()),
+          this->Counts.PrepareForInput(Device()));
+  }
+
+  VTKM_CONT
+  bool operator==(const vtkm::worklet::Keys<KeyType> &other) const
+  {
+    return ((this->UniqueKeys == other.UniqueKeys) &&
+            (this->SortedValuesMap == other.SortedValuesMap) &&
+            (this->Offsets == other.Offsets) &&
+            (this->Counts == other.Counts));
+  }
+
 private:
   KeyArrayHandleType UniqueKeys;
   vtkm::cont::ArrayHandle<vtkm::Id> SortedValuesMap;
@@ -138,4 +177,57 @@ private:
 }
 } // namespace vtkm::worklet
 
+// Here we implement the type checks, transports, and fetches the rely on the
+// Keys class. We implement them here because the Keys class is not accessible
+// to the arg classes. (The worklet package depends on the cont and exec
+// packages, not the other way around.)
+
+namespace vtkm {
+namespace cont {
+namespace arg {
+
+template<typename KeyType>
+struct TypeCheck<vtkm::cont::arg::TypeCheckTagKeys,
+                 vtkm::worklet::Keys<KeyType> >
+{
+  static const bool value = true;
+};
+
+template<typename KeyType, typename Device>
+struct Transport<vtkm::cont::arg::TransportTagKeysIn,
+                 vtkm::worklet::Keys<KeyType>,
+                 Device>
+{
+  using ContObjectType = vtkm::worklet::Keys<KeyType>;
+  using ExecObjectType =
+      typename ContObjectType::template ExecutionTypes<Device>::Lookup;
+
+  VTKM_CONT
+  ExecObjectType operator()(const ContObjectType &object,
+                            const ContObjectType &inputDomain,
+                            vtkm::Id) const
+  {
+    VTKM_ASSERT(object == inputDomain);
+
+    return object.PrepareForInput(Device());
+  }
+
+  template<typename InputDomainType>
+  VTKM_CONT
+  ExecObjectType operator()(const ContObjectType &,
+                            const InputDomainType &,
+                            vtkm::Id) const
+  {
+    // If you get a compile error here, it means that you have used a KeysIn
+    // tag in your ControlSignature that was not marked as the InputDomain.
+    VTKM_STATIC_ASSERT_MSG(
+          false, "A Keys class was used in a position that is not the input domain.");
+    return ExecObjectType();
+  }
+};
+
+}
+}
+} // namespace vtkm::cont::arg
+
 #endif //vtk_m_worklet_Keys_h

vtkm/worklet/WorkletReduceByKey.h

@@ -0,0 +1,82 @@
+//============================================================================
+//  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 2016 Sandia Corporation.
+//  Copyright 2016 UT-Battelle, LLC.
+//  Copyright 2016 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_worklet_WorkletReduceByKey_h
+#define vtk_m_worklet_WorkletReduceByKey_h
+
+#include <vtkm/worklet/internal/WorkletBase.h>
+
+#include <vtkm/cont/arg/TransportTagKeysIn.h>
+#include <vtkm/cont/arg/TypeCheckTagKeys.h>
+
+#include <vtkm/exec/internal/ReduceByKeyLookup.h>
+
+#include <vtkm/exec/arg/FetchTagKeysIn.h>
+#include <vtkm/exec/arg/ThreadIndicesReduceByKey.h>
+
+namespace vtkm {
+namespace worklet {
+
+class WorkletReduceByKey : public vtkm::worklet::internal::WorkletBase
+{
+public:
+  /// \brief A control signature tag for input keys.
+  ///
+  /// A \c WorkletReduceByKey operates by collected all identical keys and
+  /// then executing the worklet on each unique key. This tag specifies a
+  /// \c Keys object that defines and manages these keys.
+  ///
+  /// A \c WorkletReduceByKey should have exactly one \c KeysIn tag in its \c
+  /// ControlSignature, and the \c InputDomain should point to it.
+  ///
+  struct KeysIn : vtkm::cont::arg::ControlSignatureTagBase
+  {
+    using TypeCheckTag = vtkm::cont::arg::TypeCheckTagKeys;
+    using TransportTag = vtkm::cont::arg::TransportTagKeysIn;
+    using FetchTag = vtkm::exec::arg::FetchTagKeysIn;
+  };
+
+  /// Reduce by key worklets use the related thread indices class.
+  ///
+  VTKM_SUPPRESS_EXEC_WARNINGS
+  template<typename T,
+           typename OutToInArrayType,
+           typename VisitArrayType,
+           typename InputDomainType>
+  VTKM_EXEC
+  vtkm::exec::arg::ThreadIndicesReduceByKey
+  GetThreadIndices(const T& threadIndex,
+                   const OutToInArrayType& outToIn,
+                   const VisitArrayType& visit,
+                   const InputDomainType &inputDomain,
+                   const T& globalThreadIndexOffset=0) const
+  {
+    return vtkm::exec::arg::ThreadIndicesReduceByKey(
+          threadIndex,
+          outToIn.Get(threadIndex),
+          visit.Get(threadIndex),
+          inputDomain,
+          globalThreadIndexOffset);
+  }
+};
+
+}
+} // namespace vtkm::worklet
+
+#endif //vtk_m_worklet_WorkletReduceByKey_h

vtkm/worklet/testing/CMakeLists.txt

@@ -43,6 +43,7 @@ set(unit_tests
   UnitTestWorkletMapFieldWholeArray.cxx
   UnitTestWorkletMapTopologyExplicit.cxx
   UnitTestWorkletMapTopologyUniform.cxx
+  UnitTestWorkletReduceByKey.cxx
   UnitTestVertexClustering.cxx
   UnitTestWaveletCompressor.cxx
  )

vtkm/worklet/testing/UnitTestWorkletReduceByKey.cxx

@@ -0,0 +1,99 @@
+//=============================================================================
+//
+//  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 2016 Sandia Corporation.
+//  Copyright 2016 UT-Battelle, LLC.
+//  Copyright 2016 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.
+//
+//=============================================================================
+
+#include <vtkm/worklet/DispatcherReduceByKey.h>
+#include <vtkm/worklet/WorkletReduceByKey.h>
+
+#include <vtkm/worklet/Keys.h>
+
+#include <vtkm/cont/testing/Testing.h>
+
+namespace {
+
+static const vtkm::Id ARRAY_SIZE = 1033;
+static const vtkm::Id NUM_UNIQUE = ARRAY_SIZE/10;
+
+struct CheckReduceByKeyWorklet : vtkm::worklet::WorkletReduceByKey
+{
+  typedef void ControlSignature(KeysIn);
+  typedef void ExecutionSignature(_1, WorkIndex);
+  typedef _1 InputDomain;
+
+  template<typename T>
+  VTKM_EXEC
+  void operator()(const T &key, vtkm::Id workIndex) const
+  {
+    // These tests only work if keys are in sorted order, which is how we group
+    // them.
+
+    if (key != TestValue(workIndex, T()))
+    {
+      this->RaiseError("Unexpected key");
+    }
+  }
+};
+
+template<typename KeyType>
+void TryKeyType(KeyType)
+{
+  KeyType keyBuffer[ARRAY_SIZE];
+  for (vtkm::Id index = 0; index < ARRAY_SIZE; index++)
+  {
+    keyBuffer[index] = TestValue(index%NUM_UNIQUE, KeyType());
+  }
+
+  vtkm::cont::ArrayHandle<KeyType> keyArray =
+      vtkm::cont::make_ArrayHandle(keyBuffer, ARRAY_SIZE);
+
+  vtkm::worklet::Keys<KeyType> keys(keyArray,
+                                    VTKM_DEFAULT_DEVICE_ADAPTER_TAG());
+
+  vtkm::worklet::DispatcherReduceByKey<CheckReduceByKeyWorklet> dispatcher;
+  dispatcher.Invoke(keys);
+}
+
+void TestReduceByKey()
+{
+  typedef vtkm::cont::DeviceAdapterTraits<
+                    VTKM_DEFAULT_DEVICE_ADAPTER_TAG> DeviceAdapterTraits;
+  std::cout << "Testing Map Field on device adapter: "
+            << DeviceAdapterTraits::GetName() << std::endl;
+
+  std::cout << "Testing vtkm::Id keys." << std::endl;
+  TryKeyType(vtkm::Id());
+
+  std::cout << "Testing vtkm::IdComponent keys." << std::endl;
+  TryKeyType(vtkm::IdComponent());
+
+  std::cout << "Testing vtkm::UInt8 keys." << std::endl;
+  TryKeyType(vtkm::UInt8());
+
+  std::cout << "Testing vtkm::Id3 keys." << std::endl;
+  TryKeyType(vtkm::Id3());
+}
+
+} // anonymous namespace
+
+int UnitTestWorkletReduceByKey(int, char*[])
+{
+  return vtkm::cont::testing::Testing::Run(TestReduceByKey);
+}