Added proper LHS handling for transforms

NVIDIA · Aug 4, 2023 · a9a5d8e · a9a5d8e
1 parent 6213d10
commit a9a5d8e
Show file tree

Hide file tree

Showing 10 changed files with 98 additions and 11 deletions.
diff --git a/include/matx/operators/collapse.h b/include/matx/operators/collapse.h
@@ -132,7 +132,15 @@ namespace matx
             return op_.Size(DIM + dim - 1);
         }
 
-        template<typename R> __MATX_INLINE__ auto operator=(const R &rhs) { return set(*this, rhs); }
+        template<typename R> 
+        __MATX_INLINE__ auto operator=(const R &rhs) { 
+          if constexpr (is_matx_transform_op<R>()) {
+            return mtie(*this, rhs);
+          }
+          else {          
+            return set(*this, rhs); 
+          }
+        }
     };
   }
   /**
@@ -257,7 +265,15 @@ namespace matx
             return op_.Size(dim);
         }
 
-        template<typename R> __MATX_INLINE__ auto operator=(const R &rhs) { return set(*this, rhs); }
+        template<typename R> 
+        __MATX_INLINE__ auto operator=(const R &rhs) { 
+          if constexpr (is_matx_transform_op<R>()) {
+            return mtie(*this, rhs);
+          }
+          else {          
+            return set(*this, rhs); 
+          }
+        }
     };
   }
   /**

diff --git a/include/matx/operators/concat.h b/include/matx/operators/concat.h
@@ -165,7 +165,15 @@ namespace matx
           return cuda::std::get<0>(ops_).Size(dim);
       }
 
-      template<typename R> __MATX_INLINE__ auto operator=(const R &rhs) { return set(*this, rhs); }
+      template<typename R> 
+      __MATX_INLINE__ auto operator=(const R &rhs) { 
+        if constexpr (is_matx_transform_op<R>()) {
+          return mtie(*this, rhs);
+        }
+        else {          
+          return set(*this, rhs); 
+        }
+      }
 
       private:
       cuda::std::tuple<Ts...> ops_;

diff --git a/include/matx/operators/permute.h b/include/matx/operators/permute.h
@@ -119,7 +119,15 @@ namespace matx
           return op_.Size(dims_[dim]);
         }
 
-        template<typename R> __MATX_INLINE__ auto operator=(const R &rhs) { return set(*this, rhs); }
+        template<typename R> 
+        __MATX_INLINE__ auto operator=(const R &rhs) { 
+          if constexpr (is_matx_transform_op<R>()) {
+            return mtie(*this, rhs);
+          }
+          else {          
+            return set(*this, rhs); 
+          }
+        }
     };
   }
 

diff --git a/include/matx/operators/remap.h b/include/matx/operators/remap.h
@@ -120,7 +120,15 @@ namespace matx
           }
         }
 
-        template<typename R> __MATX_INLINE__ auto operator=(const R &rhs) { return set(*this, rhs); }
+        template<typename R> 
+        __MATX_INLINE__ auto operator=(const R &rhs) { 
+          if constexpr (is_matx_transform_op<R>()) {
+            return mtie(*this, rhs);
+          }
+          else {          
+            return set(*this, rhs); 
+          }
+        }
     };
   }
 

diff --git a/include/matx/operators/reshape.h b/include/matx/operators/reshape.h
@@ -135,7 +135,15 @@ namespace matx
           return sizes_[dim];
         }
 
-        template<typename R> __MATX_INLINE__ auto operator=(const R &rhs) { return set(*this, rhs); }
+        template<typename R> 
+        __MATX_INLINE__ auto operator=(const R &rhs) { 
+          if constexpr (is_matx_transform_op<R>()) {
+            return mtie(*this, rhs);
+          }
+          else {          
+            return set(*this, rhs); 
+          }
+        }
     };
   }
 

diff --git a/include/matx/operators/reverse.h b/include/matx/operators/reverse.h
@@ -116,7 +116,15 @@ namespace matx
           return op_.Size(dim);
         }
 
-        template<typename R> __MATX_INLINE__ auto operator=(const R &rhs) { return set(*this, rhs); }
+        template<typename R> 
+        __MATX_INLINE__ auto operator=(const R &rhs) { 
+          if constexpr (is_matx_transform_op<R>()) {
+            return mtie(*this, rhs);
+          }
+          else {          
+            return set(*this, rhs); 
+          }
+        }
     };
   }
 

diff --git a/include/matx/operators/shift.h b/include/matx/operators/shift.h
@@ -115,7 +115,15 @@ namespace matx
           return detail::matx_max(size1,size2);
         }
 
-        template<typename R> __MATX_INLINE__ auto operator=(const R &rhs) { return set(*this, rhs); }
+        template<typename R> 
+        __MATX_INLINE__ auto operator=(const R &rhs) { 
+          if constexpr (is_matx_transform_op<R>()) {
+            return mtie(*this, rhs);
+          }
+          else {          
+            return set(*this, rhs); 
+          }
+        }
     };
   }
   /**

diff --git a/include/matx/operators/slice.h b/include/matx/operators/slice.h
@@ -151,7 +151,15 @@ namespace matx
           return sizes_[dim];
         }
 
-        template<typename R> __MATX_INLINE__ auto operator=(const R &rhs) { return set(*this, rhs); }
+        template<typename R> 
+        __MATX_INLINE__ auto operator=(const R &rhs) { 
+          if constexpr (is_matx_transform_op<R>()) {
+            return mtie(*this, rhs);
+          }
+          else {          
+            return set(*this, rhs); 
+          }
+        }
     };
   }
 

diff --git a/include/matx/operators/stack.h b/include/matx/operators/stack.h
@@ -181,7 +181,15 @@ namespace matx
         }
       }
 
-     template<typename R> __MATX_INLINE__ auto operator=(const R &rhs) { return set(*this, rhs); }
+      template<typename R> 
+      __MATX_INLINE__ auto operator=(const R &rhs) { 
+        if constexpr (is_matx_transform_op<R>()) {
+          return mtie(*this, rhs);
+        }
+        else {          
+          return set(*this, rhs); 
+        }
+      }
 
       private:
       cuda::std::tuple<Ts...> ops_;

diff --git a/include/matx/operators/transpose.h b/include/matx/operators/transpose.h
@@ -114,7 +114,14 @@ namespace detail {
       }
 
       template<typename R> 
-      __MATX_INLINE__ auto operator=(const R &rhs) { return set(*this, rhs); }
+      __MATX_INLINE__ auto operator=(const R &rhs) { 
+        if constexpr (is_matx_transform_op<R>()) {
+          return mtie(*this, rhs);
+        }
+        else {          
+          return set(*this, rhs); 
+        }
+      }
 
   };
 }