Support for dynamic parameteric inputs in inference on heterogeneous architectures

check_style.sh

@@ -1,14 +1,17 @@
 #!/bin/bash
 set -e
 
-# Directories
-SRC_DIR="./include"
-TEST_DIR="./tests"
+# Directories (SOFIE's own layout: source lives under core/parsers/utils,
+# tests under test/ - unlike sofieBLAS this script was originally written for,
+# SOFIE has no top-level include/ or tests/ directory)
+SRC_DIRS="./core ./parsers ./utils"
+TEST_DIR="./test"
 
 echo "📝 Discovering source/header files..."
 
-FILES=$(find "$SRC_DIR" "$TEST_DIR" \
+FILES=$(find $SRC_DIRS "$TEST_DIR" \
     -path "$TEST_DIR/build" -prune -o \
+    -path "$TEST_DIR/input_models/references" -prune -o \
     -type f \( \
         -name '*.cpp' -o -name '*.cc' -o -name '*.cxx' -o \
         -name '*.h' -o -name '*.hpp' -o -name '*.hxx' -o -name '*.hh' \
@@ -31,7 +34,7 @@ done
 echo "🔍 Running clang-tidy..."
 for file in $FILES; do
     echo "Linting $file"
-    clang-tidy "$file" --extra-arg=-std=c++20 -- -I"$SRC_DIR" || true
+    clang-tidy "$file" --extra-arg=-std=c++20 -- -I./core/inc -I./parsers/inc -I./utils || true
 done
 
 echo "✅ Formatting and linting complete."

core/inc/SOFIE/RModel.hxx

@@ -8,6 +8,10 @@
 
 namespace SOFIE {
 
+/// Session member holding the construction-time value of a shape parameter (N -> fN);
+/// shared by the CPU and GPU session generators
+std::string GetMemberNameForDimShape(std::string name);
+
 class RModel final : public RModel_Base {
 
    friend class RModelProfiler;
@@ -60,7 +64,7 @@ private:
       std::vector<size_t> opIndices; ///< consecutive op indices forming this group
       std::string inputTensor;       ///< input tensor name of the first op
       std::string outputTensor;      ///< output tensor name of the last op
-      size_t numElements = 0;
+      std::string iterationLengthExpr;        ///< element count: literal for static tensors, runtime expression for dynamic
       bool isFused() const { return opIndices.size() > 1; }
       std::string suffix() const {
          std::string s;
@@ -211,6 +215,12 @@ public:
    // generate the infer function signature for inference on ALPAKA. If isdecl= false generate the calling infer function
    // used to infer the sub-graphs
    std::string GenerateInferSignature_GPU_ALPAKA(bool isdecl = true);
+   /// iterate the GPU infer argument sequence in order: onParam(param) for each dynamic
+   /// shape parameter at the position of the input introducing it, onInput(name) for each
+   /// input tensor; single source of the argument-order convention shared by the
+   /// signatures, the _infer_impl call sites and the Session constructor
+   void ForEachInferArg_GPU_ALPAKA(const std::function<void(const std::string &)> &onParam,
+                                   const std::function<void(const std::string &)> &onInput) const;
 
    // generate the _infer_impl signature using ViewPlainPtr types instead of Buf types
    std::string GenerateImplSignature_GPU_ALPAKA(bool isdecl = true);

core/inc/SOFIE/RModelProfilerGPU.hxx

@@ -35,7 +35,7 @@ public:
    static std::string GenerateMemoryReport(const MemoryInfo &info);
 
    static std::string GenerateBeginInferCode();
-   static std::string GenerateOperatorCode(ROperator &op, size_t op_idx);
+   static std::string GenerateOperatorCode(ROperator &op, size_t op_idx, const std::vector<std::string> &dynParamNames);
    static std::string GenerateEndInferCode();
 
    RModelProfilerGPU() = delete;

core/inc/SOFIE/ROperator.hxx

@@ -77,6 +77,10 @@ public:
    virtual void Initialize(RModel&) = 0;
    virtual std::string Generate(std::string OpName) = 0;  //expect unique opName for each operator within the same RModel
    virtual std::string Generate_GPU_ALPAKA(std::string OpName){ return "";} //expect unique opName for each operator within the same RModel
+   //dynParamNames: the model's shape parameters in infer-argument order
+   virtual std::string Generate_GPU_ALPAKA(std::string OpName, const std::vector<std::string> & /*dynParamNames*/) {
+      return Generate_GPU_ALPAKA(OpName);
+   }
    // generate initialization code for session constructor
    virtual std::string GenerateInitCode() { return "";}
    virtual std::string GenerateInitCode_GPU_ALPAKA() { return "";};
@@ -87,6 +91,9 @@ public:
    // generate session data members specific to operator
    virtual std::string GenerateSessionMembersCode(std::string /*opName*/) { return ""; }
    virtual std::string Generate_GPU_Kernel_ALPAKA(std::string /*opName*/) { return ""; }
+   virtual std::string Generate_GPU_Kernel_ALPAKA(std::string opName, const std::vector<std::string> & /*dynParamNames*/) {
+      return Generate_GPU_Kernel_ALPAKA(opName);
+   }
    virtual std::string Generate_GPU_Kernel_Definitions_ALPAKA(std::string /*opName*/) { return ""; }
    virtual std::string Header() { return "";}
    virtual std::string GetFusableOutputTensorName() { return "";}

core/inc/SOFIE/ROperator_BasicBinary.hxx

@@ -276,6 +276,9 @@ public:
                          << ConvertShapeToString(fShapeY) << std::endl;
             }
             // we convert non-dim shapes to Dim shapes
+            // (fShapeA/fShapeB were rank-padded to fShapeY by MultidirectionalBroadcastShape)
+            fDimShapeA = ConvertShapeToDim(fShapeA);
+            fDimShapeB = ConvertShapeToDim(fShapeB);
             fDimShapeY = ConvertShapeToDim(fShapeY);
          }
       } else {
@@ -451,97 +454,63 @@ public:
       return out.str();
    }
 
-   std::string Generate_GPU_Kernel_ALPAKA(std::string opName) {
+   // fDimShapeA/fDimShapeB are rank-padded to fDimShapeY by Initialize, so a broadcast
+   // dim is a literal 1. An input is either fully broadcast (scalar, all dims 1), not
+   // broadcast (same shape as the output), or partially broadcast - only the last case
+   // needs per-dim index decomposition in the kernel.
+   bool IsPartialBroadcast(const std::vector<Dim> &shape) const {
+      return ConvertDimShapeToLength(shape) != "1" && !UTILITY::AreSameShape(shape, fDimShapeY);
+   }
+
+   std::string Generate_GPU_Kernel_ALPAKA(std::string opName, const std::vector<std::string> &dynParamNames) override {
       if (fIsOutputConstant)
          return "";
+      if (fDimShapeY.empty())
+         throw std::runtime_error("SOFIE Operator Basic Binary called to Generate without being initialized first");
+
+      const std::size_t D = fDimShapeY.size();
+      auto stridesY = UTILITY::ComputeStrideFromShape(fDimShapeY);
+      auto stridesA = UTILITY::ComputeStrideFromShape(fDimShapeA);
+      auto stridesB = UTILITY::ComputeStrideFromShape(fDimShapeB);
+      bool isAScalar = ConvertDimShapeToLength(fDimShapeA) == "1";
+      bool isBScalar = ConvertDimShapeToLength(fDimShapeB) == "1";
+      bool isAPartialBroadcast = IsPartialBroadcast(fDimShapeA);
+      bool isBPartialBroadcast = IsPartialBroadcast(fDimShapeB);
 
       std::string op;
       op = "\n//------ "+opName+"_"+BinaryOperatorTrait<T, Op>::Name()+"_KERNEL_ALPAKA\n";
       op += SP + "struct Binary"+opName+BinaryOperatorTrait<T, Op>::Name()+"Kernel {\n";
       op += SP + SP + "template<typename TAcc, typename T>\n";
-      op += SP + SP + "ALPAKA_FN_ACC void operator()(TAcc const & acc, T const * A, T const * B, T * C) const {\n";
+      op += SP + SP + "ALPAKA_FN_ACC void operator()(TAcc const & acc, T const * A, T const * B, T * C";
+      for (auto &p : dynParamNames)
+         op += ", std::size_t const " + p;
+      op += ", std::size_t const totalElements) const {\n";
       op += SP + SP + SP + "auto idx = alpaka::getIdx<alpaka::Grid, alpaka::Threads>(acc)[0];\n";
-      op += SP + SP + SP + "if (idx < " + std::to_string(ConvertShapeToLength(fShapeY)) + ") {\n";
-      auto stridesA = UTILITY::ComputeStrideFromShape(fShapeA);
-      auto stridesB = UTILITY::ComputeStrideFromShape(fShapeB);
-
-      for(size_t id_s = 0; id_s < stridesA.size(); ++id_s){
-         if(fShapeA[id_s] == 1)
-            stridesA[id_s] = 0;
-      }
-
-      for(size_t id_s = 0; id_s < stridesB.size(); ++id_s){
-         if(fShapeB[id_s] == 1)
-            stridesB[id_s] = 0;
-      }
-
-      auto stridesY = UTILITY::ComputeStrideFromShape(fShapeY);
-
-      // --- Fast-path index simplifications ---
-      // Check whether A is broadcast (all strides zero → single element)
-      bool isAScalar = true;
-      for (const auto& s : stridesA) { if (s != 0) { isAScalar = false; break; } }
-      // Check whether B is broadcast (all strides zero → single element)
-      bool isBScalar = true;
-      for (const auto& s : stridesB) { if (s != 0) { isBScalar = false; break; } }
-      // Check whether A has the same contiguous layout as Y (no broadcasting)
-      bool isAContiguous = (fShapeA.size() == fShapeY.size());
-      if (isAContiguous) {
-         for (size_t i = 0; i < fShapeA.size(); ++i)
-            if (fShapeA[i] != fShapeY[i]) { isAContiguous = false; break; }
-      }
-      // Check whether B has the same contiguous layout as Y (no broadcasting)
-      bool isBContiguous = (fShapeB.size() == fShapeY.size());
-      if (isBContiguous) {
-         for (size_t i = 0; i < fShapeB.size(); ++i)
-            if (fShapeB[i] != fShapeY[i]) { isBContiguous = false; break; }
-      }
-
-      std::string flattened_index_A = "";
-      std::string flattened_index_B = "";
-
-      if (isAScalar) {
-         // A is a single broadcast value
-         flattened_index_A = "0";
-      } else if (isAContiguous) {
-         // A and Y have identical shapes → direct index
-         flattened_index_A = "idx";
-      } else {
-         // General broadcast case: decompose idx into per-dim coords
-         std::string temp = "idx";
-         for (size_t id_s = 0; id_s < fShapeA.size(); ++id_s) {
-            auto strideY = stridesY[id_s];
-            auto strideA = stridesA[id_s];
-            std::string coord = "(int)(" + temp + " / " + std::to_string(strideY) + ")";
-            flattened_index_A += coord + " * " + std::to_string(strideA) + " + ";
-            temp = temp + " - (" + coord + " * " + std::to_string(strideY) + ")";
+      op += SP + SP + SP + "if (idx < totalElements) {\n";
+
+      // decompose idx into output coords, skipping broadcast dims (stride 0)
+      if (isAPartialBroadcast || isBPartialBroadcast) {
+         op += SP + SP + SP + SP + "std::size_t remaining = idx;\n";
+         op += SP + SP + SP + SP + "std::size_t coord;\n";
+         if (isAPartialBroadcast) op += SP + SP + SP + SP + "std::size_t idxA = 0;\n";
+         if (isBPartialBroadcast) op += SP + SP + SP + SP + "std::size_t idxB = 0;\n";
+         for (std::size_t d = 0; d < D; d++) {
+            std::string sY = "(" + stridesY[d].GetVal() + ")";
+            op += SP + SP + SP + SP + "coord = remaining / " + sY + ";\n";
+            if (d + 1 < D)
+               op += SP + SP + SP + SP + "remaining -= coord * " + sY + ";\n";
+            if (isAPartialBroadcast && fDimShapeA[d].GetVal() != "1")
+               op += SP + SP + SP + SP + "idxA += coord * (" + stridesA[d].GetVal() + ");\n";
+            if (isBPartialBroadcast && fDimShapeB[d].GetVal() != "1")
+               op += SP + SP + SP + SP + "idxB += coord * (" + stridesB[d].GetVal() + ");\n";
          }
-         if (!flattened_index_A.empty())
-            flattened_index_A.erase(flattened_index_A.size() - 3);
       }
-
-      if (isBScalar) {
-         // B is a single broadcast value
-         flattened_index_B = "0";
-      } else if (isBContiguous) {
-         // B and Y have identical shapes → direct index
-         flattened_index_B = "idx";
-      } else {
-         // General broadcast case
-         std::string temp = "idx";
-         for (size_t id_s = 0; id_s < fShapeB.size(); ++id_s) {
-            auto strideY = stridesY[id_s];
-            auto strideB = stridesB[id_s];
-            std::string coord = "(int)(" + temp + " / " + std::to_string(strideY) + ")";
-            flattened_index_B += coord + " * " + std::to_string(strideB) + " + ";
-            temp = temp + " - (" + coord + " * " + std::to_string(strideY) + ")";
-         }
-         if (!flattened_index_B.empty())
-            flattened_index_B.erase(flattened_index_B.size() - 3);
-      }
-
-      op += "C[idx] = " + BinaryOperatorTrait<T, Op>::Op("A["+flattened_index_A+"]", "B["+flattened_index_B+"]") + ";\n";
-      op += "}\n}\n};\n";
+      std::string indexA = isAScalar ? "0" : (isAPartialBroadcast ? "idxA" : "idx");
+      std::string indexB = isBScalar ? "0" : (isBPartialBroadcast ? "idxB" : "idx");
+      op += SP + SP + SP + SP + "C[idx] = " + BinaryOperatorTrait<T, Op>::Op("A["+indexA+"]", "B["+indexB+"]") + ";\n";
+      op += SP + SP + SP + "}\n";
+      op += SP + SP + "}\n";
+      op += SP + "};\n";
       return op;
    }
 
@@ -552,7 +521,7 @@ public:
       return SP + "Binary"+OpName+BinaryOperatorTrait<T, Op>::Name()+"Kernel binary" + OpName + "Kernel;\n";
    }
 
-   std::string Generate_GPU_ALPAKA(std::string OpName) {
+   std::string Generate_GPU_ALPAKA(std::string OpName, const std::vector<std::string> &dynParamNames) override {
       if (fIsOutputConstant)
          return "";
 
@@ -567,7 +536,10 @@ public:
       out << SP << "auto const workDiv_" << fNY << " = sofie_workdiv(elementsPerGrid_" << fNY << ");\n";
       out << SP << "auto task_" << OpName << " = alpaka::createTaskKernel<Acc>(workDiv_" << fNY
          << ", binary" << OpName << "Kernel, alpaka::getPtrNative(deviceBuf_" << fNA
-         << "), alpaka::getPtrNative(deviceBuf_" << fNB << "), alpaka::getPtrNative(deviceBuf_" << fNY << "));\n";
+         << "), alpaka::getPtrNative(deviceBuf_" << fNB << "), alpaka::getPtrNative(deviceBuf_" << fNY << ")";
+      for (auto &p : dynParamNames)
+         out << ", static_cast<std::size_t>(" << p << ")";
+      out << ", static_cast<Idx>(" << length << "));\n";
       out << SP << "alpaka::enqueue(queue, task_" << OpName << ");\n";
       return out.str();
    }

core/inc/SOFIE/ROperator_BasicNary.hxx

@@ -145,6 +145,11 @@ public:
             fNBroadcastedInputs.emplace_back("tensor_" + fNInputs[i]);
          }
       }
+      // rank-pad each input shape to fShapeY; the GPU emitters read them directly
+      for (auto &s : fShapeInputs) {
+         auto y = fShapeY;
+         UTILITY::MultidirectionalBroadcastShape(s, y);
+      }
       fType = ConvertTypeToString(model.GetTensorType(fNInputs[0]));
    }
 
@@ -185,6 +190,109 @@ public:
       return out.str();
    }
 
+   std::string GetGPUCombine(const std::string& acc_v, const std::string& val) const {
+      if (Op == EBasicNaryOperator::Max)
+         return acc_v + " = (" + acc_v + " > " + val + ") ? " + acc_v + " : " + val + ";";
+      if (Op == EBasicNaryOperator::Min)
+         return acc_v + " = (" + acc_v + " < " + val + ") ? " + acc_v + " : " + val + ";";
+      return acc_v + " = " + acc_v + " + " + val + ";"; // Sum and Mean both accumulate
+   }
+
+   std::string Generate_GPU_Kernel_ALPAKA(std::string OpName, const std::vector<std::string> &dynParamNames) override {
+      OpName = "op_" + OpName;
+      size_t nIn = fNInputs.size();
+      const std::size_t D = fShapeY.size();
+      auto stridesY = UTILITY::ComputeStrideFromShape(fShapeY);
+      auto isOne = [](const Dim &d) { return d.GetVal() == "1"; };
+      std::vector<bool> isScalar(nIn), isContiguous(nIn), isPartial(nIn);
+      std::vector<std::vector<Dim>> strides(nIn);
+      bool anyPartial = false;
+      for (size_t i = 0; i < nIn; i++) {
+         isScalar[i]     = ConvertDimShapeToLength(fShapeInputs[i]) == "1";
+         isContiguous[i] = UTILITY::AreSameShape(fShapeInputs[i], fShapeY);
+         isPartial[i]    = !isScalar[i] && !isContiguous[i];
+         if (isPartial[i]) {
+            strides[i] = UTILITY::ComputeStrideFromShape(fShapeInputs[i]);
+            anyPartial = true;
+         }
+      }
+      std::string op;
+      op += "\n//------ BASICNARY_KERNEL_ALPAKA\n";
+      op += SP + "struct BasicNaryKernel_" + OpName + " {\n";
+      op += SP + SP + "template<typename TAcc, typename TOut";
+      for (size_t i = 0; i < nIn; i++)
+         op += ", typename Tin" + std::to_string(i);
+      op += ">\n";
+      op += SP + SP + "ALPAKA_FN_ACC void operator()(TAcc const& acc";
+      for (size_t i = 0; i < nIn; i++)
+         op += ", Tin" + std::to_string(i) + " const* in" + std::to_string(i);
+      op += ", TOut* out";
+      for (auto &p : dynParamNames)
+         op += ", std::size_t const " + p;
+      op += ", std::size_t n) const {\n";
+      op += SP + SP + SP + "auto const idx = alpaka::getIdx<alpaka::Grid, alpaka::Threads>(acc)[0];\n";
+      op += SP + SP + SP + "if (idx >= n) return;\n";
+
+      // decompose idx into output coords, skipping broadcast dims (stride 0)
+      if (anyPartial) {
+         op += SP + SP + SP + "std::size_t remaining = idx;\n";
+         op += SP + SP + SP + "std::size_t coord;\n";
+         for (size_t i = 0; i < nIn; i++)
+            if (isPartial[i])
+               op += SP + SP + SP + "std::size_t idx" + std::to_string(i) + " = 0;\n";
+         for (std::size_t d = 0; d < D; d++) {
+            std::string sY = "(" + stridesY[d].GetVal() + ")";
+            op += SP + SP + SP + "coord = remaining / " + sY + ";\n";
+            if (d + 1 < D)
+               op += SP + SP + SP + "remaining -= coord * " + sY + ";\n";
+            for (size_t i = 0; i < nIn; i++) {
+               if (!isPartial[i] || isOne(fShapeInputs[i][d])) continue;
+               op += SP + SP + SP + "idx" + std::to_string(i) + " += coord * (" + strides[i][d].GetVal() + ");\n";
+            }
+         }
+      }
+      auto index = [&](size_t i) -> std::string {
+         if (isContiguous[i]) return "idx";
+         if (isScalar[i]) return "0";
+         return "idx" + std::to_string(i);
+      };
+      op += SP + SP + SP + "TOut v = static_cast<TOut>(in0[" + index(0) + "]);\n";
+      for (size_t i = 1; i < nIn; i++)
+         op += SP + SP + SP + "{ TOut w = static_cast<TOut>(in" + std::to_string(i) + "[" + index(i) + "]); " + GetGPUCombine("v", "w") + " }\n";
+      if (Op == EBasicNaryOperator::Mean)
+         op += SP + SP + SP + "v = v / static_cast<TOut>(" + std::to_string(nIn) + ");\n";
+      op += SP + SP + SP + "out[idx] = v;\n";
+      op += SP + SP + "}\n";
+      op += SP + "};\n";
+      return op;
+   }
+
+   std::string Generate_GPU_Kernel_Definitions_ALPAKA(std::string OpName) override {
+      OpName = "op_" + OpName;
+      return SP + "BasicNaryKernel_" + OpName + " basicNaryKernel_" + OpName + ";\n";
+   }
+
+   std::string Generate_GPU_ALPAKA(std::string OpName, const std::vector<std::string> &dynParamNames) override {
+      if (fShapeY.empty())
+         throw std::runtime_error("SOFIE BasicNary Op called to Generate without being initialized first");
+      OpName = "op_" + OpName;
+      std::stringstream out;
+      std::string length = ConvertDimShapeToLength(fShapeY);
+      out << "\n//------ BASICNARY_GPU_ALPAKA\n";
+      out << SP << "auto const elementsPerGrid_" << OpName << " = Vec::all(Idx{" << length << "});\n";
+      out << SP << "auto const workDiv_" << OpName << " = sofie_workdiv(elementsPerGrid_" << OpName << ");\n";
+      out << SP << "auto task_" << OpName << " = alpaka::createTaskKernel<Acc>(workDiv_" << OpName
+          << ", basicNaryKernel_" << OpName;
+      for (auto &in : fNInputs)
+         out << ", alpaka::getPtrNative(deviceBuf_" << in << ")";
+      out << ", alpaka::getPtrNative(deviceBuf_" << fNY << ")";
+      for (auto &p : dynParamNames)
+         out << ", static_cast<std::size_t>(" << p << ")";
+      out << ", static_cast<std::size_t>(" << length << "));\n";
+      out << SP << "alpaka::enqueue(queue, task_" << OpName << ");\n";
+      return out.str();
+   }
+
    std::vector<std::string> GetStdLibs() override {return { std::string("cmath") }; }
 };