1 //============================================================================
2 //  Copyright (c) Kitware, Inc.
3 //  All rights reserved.
4 //  See LICENSE.txt for details.
5 //
6 //  This software is distributed WITHOUT ANY WARRANTY; without even
7 //  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
8 //  PURPOSE.  See the above copyright notice for more information.
9 //============================================================================
10 
11 #ifndef vtk_m_worklet_connectivity_InnerJoin_h
12 #define vtk_m_worklet_connectivity_InnerJoin_h
13 
14 #include <vtkm/cont/ArrayHandleIndex.h>
15 #include <vtkm/worklet/DispatcherMapField.h>
16 #include <vtkm/worklet/ScatterCounting.h>
17 #include <vtkm/worklet/WorkletMapField.h>
18 
19 namespace vtkm
20 {
21 namespace worklet
22 {
23 namespace connectivity
24 {
25 class InnerJoin
26 {
27 public:
28   struct Merge : vtkm::worklet::WorkletMapField
29   {
30     using ControlSignature =
31       void(FieldIn, FieldIn, FieldIn, WholeArrayIn, FieldOut, FieldOut, FieldOut);
32     using ExecutionSignature = void(_1, _2, _3, VisitIndex, _4, _5, _6, _7);
33     using InputDomain = _1;
34 
35     using ScatterType = vtkm::worklet::ScatterCounting;
36 
37     // TODO: type trait for array portal?
38     template <typename KeyType, typename ValueType1, typename InPortalType, typename ValueType2>
operatorMerge39     VTKM_EXEC void operator()(KeyType key,
40                               ValueType1 value1,
41                               vtkm::Id lowerBounds,
42                               vtkm::Id visitIndex,
43                               const InPortalType& value2,
44                               vtkm::Id& keyOut,
45                               ValueType1& value1Out,
46                               ValueType2& value2Out) const
47     {
48       auto v2 = value2.Get(lowerBounds + visitIndex);
49       keyOut = key;
50       value1Out = value1;
51       value2Out = v2;
52     }
53   };
54 
55   using Algorithm = vtkm::cont::Algorithm;
56 
57   // TODO: not mutating input keys and values?
58   template <typename Key, typename Value1, typename Value2>
Run(vtkm::cont::ArrayHandle<Key> & key1,vtkm::cont::ArrayHandle<Value1> & value1,vtkm::cont::ArrayHandle<Key> & key2,vtkm::cont::ArrayHandle<Value2> & value2,vtkm::cont::ArrayHandle<Key> & keyOut,vtkm::cont::ArrayHandle<Value1> & value1Out,vtkm::cont::ArrayHandle<Value2> & value2Out)59   static void Run(vtkm::cont::ArrayHandle<Key>& key1,
60                   vtkm::cont::ArrayHandle<Value1>& value1,
61                   vtkm::cont::ArrayHandle<Key>& key2,
62                   vtkm::cont::ArrayHandle<Value2>& value2,
63                   vtkm::cont::ArrayHandle<Key>& keyOut,
64                   vtkm::cont::ArrayHandle<Value1>& value1Out,
65                   vtkm::cont::ArrayHandle<Value2>& value2Out)
66   {
67     Algorithm::SortByKey(key1, value1);
68     Algorithm::SortByKey(key2, value2);
69 
70     vtkm::cont::ArrayHandle<vtkm::Id> lbs;
71     vtkm::cont::ArrayHandle<vtkm::Id> ubs;
72     Algorithm::LowerBounds(key2, key1, lbs);
73     Algorithm::UpperBounds(key2, key1, ubs);
74 
75     vtkm::cont::ArrayHandle<vtkm::Id> counts;
76     Algorithm::Transform(ubs, lbs, counts, vtkm::Subtract());
77 
78     vtkm::worklet::ScatterCounting scatter{ counts };
79     vtkm::worklet::DispatcherMapField<Merge> mergeDisp(scatter);
80     mergeDisp.Invoke(key1, value1, lbs, value2, keyOut, value1Out, value2Out);
81   }
82 };
83 
84 class Renumber
85 {
86 public:
Run(vtkm::cont::ArrayHandle<vtkm::Id> & componentsInOut)87   static void Run(vtkm::cont::ArrayHandle<vtkm::Id>& componentsInOut)
88   {
89     using Algorithm = vtkm::cont::Algorithm;
90 
91     // FIXME: we should be able to apply findRoot to each pixel and use some kind
92     // of atomic operation to get the number of unique components without the
93     // cost of copying and sorting. This might be able to be extended to also
94     // work for the renumbering (replacing InnerJoin) through atomic increment.
95     vtkm::cont::ArrayHandle<vtkm::Id> uniqueComponents;
96     Algorithm::Copy(componentsInOut, uniqueComponents);
97     Algorithm::Sort(uniqueComponents);
98     Algorithm::Unique(uniqueComponents);
99 
100     vtkm::cont::ArrayHandle<vtkm::Id> ids;
101     Algorithm::Copy(vtkm::cont::ArrayHandleIndex(componentsInOut.GetNumberOfValues()), ids);
102 
103     vtkm::cont::ArrayHandle<vtkm::Id> uniqueColor;
104     Algorithm::Copy(vtkm::cont::ArrayHandleIndex(uniqueComponents.GetNumberOfValues()),
105                     uniqueColor);
106 
107     vtkm::cont::ArrayHandle<vtkm::Id> cellColors;
108     vtkm::cont::ArrayHandle<vtkm::Id> pixelIdsOut;
109     InnerJoin::Run(componentsInOut,
110                    ids,
111                    uniqueComponents,
112                    uniqueColor,
113                    cellColors,
114                    pixelIdsOut,
115                    componentsInOut);
116 
117     Algorithm::SortByKey(pixelIdsOut, componentsInOut);
118   }
119 };
120 }
121 }
122 } // vtkm::worklet::connectivity
123 
124 #endif //vtk_m_worklet_connectivity_InnerJoin_h
125