From acfac384b87ba2651cfd21ff9eeb9b55e6c109e0 Mon Sep 17 00:00:00 2001
From: cliffburdick <cburdick@nvidia.com>
Date: Tue, 23 Aug 2022 12:55:37 -0700
Subject: [PATCH] Fixed sum() using wrong iterator type

---
 include/matx/transforms/cub.h       |  8 ++++----
 include/matx/transforms/reduce.h    | 30 +++++++++++++++++++++++++++++
 test/00_operators/ReductionTests.cu |  3 ++-
 3 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/include/matx/transforms/cub.h b/include/matx/transforms/cub.h
index cf038512..18301f28 100644
--- 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
index 02973c0f..8ab3afe0 100644
--- 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
index a68bc21d..ba8a4469 100644
--- 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;