Fixed sum() using wrong iterator type (#253)

NVIDIA · Aug 23, 2022 · 33e5898 · 33e5898
1 parent 54904fd
commit 33e5898
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 5 deletions.
diff --git a/include/matx/transforms/cub.h b/include/matx/transforms/cub.h
@@ -128,7 +128,7 @@ struct BeginOffset {
    * 
    * @return Value at offset 
    */
-  [[nodiscard]] __MATX_INLINE__ __MATX_HOST__ __MATX_DEVICE__ reference operator*() const
+  [[nodiscard]] __MATX_INLINE__ __MATX_HOST__ __MATX_DEVICE__ stride_type operator*() const
   {
     return offset_ * size_;
   }  
@@ -138,7 +138,7 @@ struct BeginOffset {
     return self_type{size_, offset_ + offset};
   }
 
-  [[nodiscard]] __MATX_INLINE__ __MATX_HOST__ __MATX_DEVICE__ reference operator[](difference_type offset) const
+  [[nodiscard]] __MATX_INLINE__ __MATX_HOST__ __MATX_DEVICE__ stride_type operator[](difference_type offset) const
   {
     return *self_type{size_, offset_ + offset};
   }  
@@ -187,7 +187,7 @@ struct EndOffset {
    * 
    * @return Value at offset 
    */
-  [[nodiscard]] __MATX_INLINE__ __MATX_HOST__ __MATX_DEVICE__ reference operator*() const
+  [[nodiscard]] __MATX_INLINE__ __MATX_HOST__ __MATX_DEVICE__ stride_type operator*() const
   {
     return (offset_ + 1) * size_;
   }  
@@ -197,7 +197,7 @@ struct EndOffset {
     return self_type{size_, offset_ + offset};
   }
 
-  [[nodiscard]] __MATX_INLINE__ __MATX_HOST__ __MATX_DEVICE__ reference operator[](difference_type offset) const
+  [[nodiscard]] __MATX_INLINE__ __MATX_HOST__ __MATX_DEVICE__ stride_type operator[](difference_type offset) const
   {
     return ( offset + 1) * size_;
   }  

diff --git a/include/matx/transforms/reduce.h b/include/matx/transforms/reduce.h
@@ -470,6 +470,36 @@ __MATX_DEVICE__ __MATX_INLINE__ void atomicAll(uint64_t *addr, uint64_t val)
   }
 }
 
+/**
+ * Atomic add for int64_t
+ *
+ * @param addr
+ *   Source and destination for result
+ * @param val
+ *   Value to add
+ */
+__MATX_DEVICE__ __MATX_INLINE__ void atomicAdd(int64_t *addr,
+                                 int64_t val)
+{
+  unsigned long long int *addri = reinterpret_cast<unsigned long long int *>(addr);
+  atomicAdd(addri, (unsigned long long int)val);
+}
+
+/**
+ * Atomic add for uint64_t
+ *
+ * @param addr
+ *   Source and destination for result
+ * @param val
+ *   Value to add
+ */
+__MATX_DEVICE__ __MATX_INLINE__ void atomicAdd(uint64_t *addr,
+                                 uint64_t val)
+{
+  unsigned long long int *addri = reinterpret_cast<unsigned long long int *>(addr);
+  atomicAdd(addri, (unsigned long long int)val);
+}
+
 /**
  * Atomic add for complex floats
  *

diff --git a/test/00_operators/ReductionTests.cu b/test/00_operators/ReductionTests.cu
@@ -134,7 +134,7 @@ TYPED_TEST(ReductionTestsComplexNonHalfTypes, VarianceStdComplex)
   MATX_EXIT_HANDLER();
 }
 
-TYPED_TEST(ReductionTestsFloatNonComplexNonHalf, Sum)
+TYPED_TEST(ReductionTestsNumericNoHalf, Sum)
 {
   MATX_ENTER_HANDLER();
   {
@@ -273,6 +273,7 @@ TEST(ReductionTests, Any)
 TEST(ReductionTests, All)
 {
   MATX_ENTER_HANDLER();
+
   using TypeParam = int;
   {
     tensor_t<int, 0> t0;