Fix stack() operator-operand composition (#1203)

* Fix stack() operator-operand composition

Three bugs that only surface when a stack() expression is materialized
via run() with an operator (non-tensor) operand:

(A) The stack() factory stored a local `first = pp_get<0>(ts...)` used
    only to call the static Rank() method, so its value was never read.
    Under -Werror=unused-but-set-variable this fails to compile when the
    first operand has a trivial destructor (e.g. any expression operator).

(B) StackOp inherited the no-op BaseOp::PreRun/PostRun and never
    forwarded them to its variadic operands. When an operand is a
    transform operator that allocates temporary storage during PreRun,
    that temporary was never allocated or computed, so operator() read
    uninitialized memory.

(C) StackOp::Size() computed Size(dim-1) on an operand whose Size()
    indexes a fixed-size cuda::std::array. GCC's -Werror=array-bounds
    analysis cannot prove dim-1 >= 0 (axis_ is a runtime int) and
    rejects it.

The existing Stack test only reads operator() elements directly and
never materializes a stack expression, so none of these were triggered.

Fixes:
  (A) Call ::Rank() from a type alias for the type of first. The
      alias includes a [[maybe_unused]] because Release builds elide
      the assertion and do not reference the local alias.
  (B) Add PreRun/PostRun to StackOp that forward to every operand via
      cuda::std::apply, guarded by is_matx_op<>().
  (C) Compute the operand dim once and clamp it to a valid non-negative
      index so the fixed-size-array access is provably in-bounds.

Adds StackOperatorInput test to verify Pre/PostRun forwarding.

Signed-off-by: Thomas Benson <tbenson@nvidia.com>

Author: tbensonatl

include/matx/operators/stack.h

@@ -332,12 +332,17 @@ namespace matx
       {
         if(dim==axis_)
           return sizeof...(Ts);
-        else if (dim < axis_) {
-          return cuda::std::get<0>(ops_).Size(dim);
-        } else {
-          // remove axis_ dim from dim.
-          return cuda::std::get<0>(ops_).Size(dim-1);
+        // Map the requested output dimension to the operand dimension by dropping
+        // the stacked axis. For any valid query op_dim is already >= 0 (a dim > axis_
+        // implies dim >= 1, and the only legitimate dim==0 query is the axis itself,
+        // handled above). The clamp is therefore unreachable for valid input and
+        // exists solely to give the compiler a provable lower bound for
+        // -Werror=array-bounds.
+        int op_dim = (dim < axis_) ? dim : dim - 1;
+        if (op_dim < 0) {
+          op_dim = 0;
         }
+        return cuda::std::get<0>(ops_).Size(op_dim);
       }
 
       __MATX_INLINE__ __MATX_HOST__ int32_t DynRank() const {
@@ -353,6 +358,35 @@ namespace matx
         else return Rank();
       }
 
+      // Forward PreRun/PostRun to every stacked operand. An operand may itself
+      // be a transform operator that allocates and fills temporary storage during
+      // PreRun, so its PreRun must be invoked before operator() reads from it.
+      template <typename ShapeType, typename Executor>
+      __MATX_INLINE__ void PreRun(ShapeType &&shape, Executor &&ex) const noexcept
+      {
+        cuda::std::apply([&](const auto&... ops) {
+          auto prerun_one = [&](const auto &op) {
+            if constexpr (is_matx_op<cuda::std::decay_t<decltype(op)>>()) {
+              op.PreRun(shape, ex);
+            }
+          };
+          (prerun_one(ops), ...);
+        }, ops_);
+      }
+
+      template <typename ShapeType, typename Executor>
+      __MATX_INLINE__ void PostRun(ShapeType &&shape, Executor &&ex) const noexcept
+      {
+        cuda::std::apply([&](const auto&... ops) {
+          auto postrun_one = [&](const auto &op) {
+            if constexpr (is_matx_op<cuda::std::decay_t<decltype(op)>>()) {
+              op.PostRun(shape, ex);
+            }
+          };
+          (postrun_one(ops), ...);
+        }, ops_);
+      }
+
       ~StackOp() = default;
       StackOp(const StackOp &rhs) = default;
 
@@ -451,9 +485,8 @@ namespace matx
   template <typename... Ts>
     __MATX_INLINE__ __MATX_HOST__  auto stack(int axis, const Ts&... ts)
     {
-      auto first = detail::pp_get<0>(ts...);
-
-      MATX_ASSERT_STR(axis <= first.Rank(),matxInvalidDim, "stack must take an axis less than or equal to the the rank of the operators");
+      using first_type [[maybe_unused]] = cuda::std::tuple_element_t<0, cuda::std::tuple<Ts...>>;
+      MATX_ASSERT_STR(axis <= first_type::Rank(),matxInvalidDim, "stack must take an axis less than or equal to the the rank of the operators");
       return detail::StackOp<Ts...>{axis, ts...};
     }  
 } // end namespace matx

test/00_operators/stack_test.cu

@@ -2,6 +2,7 @@
 #include "matx.h"
 #include "test_types.h"
 #include "utilities.h"
+#include "prerun_tester.h"
 
 using namespace matx;
 using namespace matx::test;
@@ -47,4 +48,62 @@ TYPED_TEST(OperatorTestsNumericAllExecs, Stack)
   }  
 
   MATX_EXIT_HANDLER();
-}
+}
+
+// Verifies that stack() correctly forwards PreRun()/PostRun() to its operands
+// when a stack expression is materialized via run(). Each variadic operand is
+// wrapped in a PreRunTesterOp lifecycle probe: stack()'s PreRun/PostRun fold
+// must forward to every operand exactly once (an unforwarded operand leaves
+// prerun_count == 0). The probe is a transparent pass-through, so the
+// materialized result is still the correct stacked output and is checked against
+// a reference. The cumsum() operands are real transforms whose temporaries are
+// only allocated/filled if PreRun is forwarded.
+TYPED_TEST(OperatorTestsFloatNonComplexNonHalfAllExecsWithoutJIT, StackOperatorInput)
+{
+  MATX_ENTER_HANDLER();
+  using TestType = cuda::std::tuple_element_t<0, TypeParam>;
+  using ExecType = cuda::std::tuple_element_t<1, TypeParam>;
+
+  ExecType exec{};
+
+  auto a = make_tensor<TestType>({5});
+  auto b = make_tensor<TestType>({5});
+  auto c = make_tensor<TestType>({5});
+  a.SetVals({1, 2, 3, 4, 5});       // cumsum(a) = {1, 3, 6, 10, 15}
+  b.SetVals({10, 20, 30, 40, 50});  // leaf operand
+  c.SetVals({2, 4, 6, 8, 10});      // cumsum(c) = {2, 6, 12, 20, 30}
+
+  // Reference: materialize the transform operands into tensors first.
+  auto ca = make_tensor<TestType>({5});
+  auto cc = make_tensor<TestType>({5});
+  (ca = cumsum(a)).run(exec);
+  (cc = cumsum(c)).run(exec);
+  auto out_ref = make_tensor<TestType>({3, 5});
+  (out_ref = stack(0, ca, b, cc)).run(exec);
+
+  // Under test: wrap each stacked operand in its own lifecycle probe. A mix of
+  // two transform operands and a leaf exercises the variadic fold over more than
+  // two operands.
+  PreRunLifecycle s0, s1, s2;
+  auto out_test = make_tensor<TestType>({3, 5});
+  (out_test = stack(0,
+                    make_prerun_tester(cumsum(a), s0),
+                    make_prerun_tester(b, s1),
+                    make_prerun_tester(cumsum(c), s2))).run(exec);
+
+  exec.sync();
+
+  // Correctness preserved (probe is a transparent pass-through).
+  for (int i = 0; i < 3; i++) {
+    for (int j = 0; j < 5; j++) {
+      ASSERT_EQ(out_test(i, j), out_ref(i, j)) << "mismatch at (" << i << "," << j << ")";
+    }
+  }
+
+  // Lifecycle: stack() forwarded a balanced PreRun/PostRun to every operand.
+  ExpectLifecycleClean(s0, "cumsum(a)");
+  ExpectLifecycleClean(s1, "b");
+  ExpectLifecycleClean(s2, "cumsum(c)");
+
+  MATX_EXIT_HANDLER();
+}