feat: on-the-fly broadcasting

sanjibansg · lmoneta · commit 07223cdd6399 · 2025-08-11T18:29:43.000+02:00
diff --git a/tmva/sofie/inc/TMVA/ROperator_BasicBinary.hxx b/tmva/sofie/inc/TMVA/ROperator_BasicBinary.hxx
@@ -144,7 +144,7 @@ public:
                   fDimShapeA = ConvertShapeToDim(fShapeA);
                } else {
                   // Add an intermediate tensor for broadcasting A
-                  model.AddIntermediateTensor(fNBroadcastedA, model.GetTensorType(fNA), fShapeY);
+                  // model.AddIntermediateTensor(fNBroadcastedA, model.GetTensorType(fNA), fShapeY);
                }
             }
             // Broadcast B to Y
@@ -170,7 +170,7 @@ public:
                   fDimShapeB = ConvertShapeToDim(fShapeB);
                } else {
                   // Add an intermediate tensor for broadcasting B
-                  model.AddIntermediateTensor(fNBroadcastedB, model.GetTensorType(fNB), fShapeY);
+                  // model.AddIntermediateTensor(fNBroadcastedB, model.GetTensorType(fNB), fShapeY);
                }
             }
          } else {
@@ -281,6 +281,52 @@ public:
       out << SP << "\n//------ " << BinaryOperatorTrait<T,Op>::Name() << "\n";
       auto length = ConvertDimShapeToLength(fDimShapeY);
       std::string typeName = TensorType<T>::Name();
+
+
+      // non dynamic case
+      if (!fShapeA.empty() && !fShapeB.empty() && !fShapeY.empty() ) {
+
+      auto stridesA = UTILITY::ComputeStrideFromShape(fShapeA);
+      auto stridesB = UTILITY::ComputeStrideFromShape(fShapeB);
+      auto stridesY = UTILITY::ComputeStrideFromShape(fShapeY);
+
+      std::string compute_idx_A, compute_idx_B, compute_idx_Y;
+      if (std::all_of(fShapeA.begin(), fShapeA.end(), [](size_t x) { return x == 1; })){
+         compute_idx_A = "0";
+      } else {
+         for(size_t i = 0; i<fShapeA.size(); ++i){
+            if(fShapeA[i]==1) continue;
+            compute_idx_A += " idx_"+fNY+std::to_string(i+(fShapeY.size()-fShapeA.size()))+" * "+stridesA[i]+" +";
+         }
+         compute_idx_A.pop_back();
+      }
+      if (std::all_of(fShapeB.begin(), fShapeB.end(), [](size_t x) { return x == 1; })){
+         compute_idx_B = "0";
+      } else {
+         for(size_t i = 0; i<fShapeB.size(); ++i){
+            if(fShapeB[i]==1) continue;
+            compute_idx_B += " idx_"+fNY+std::to_string(i+(fShapeY.size()-fShapeB.size()))+" * "+stridesB[i]+" +";
+         }
+         compute_idx_B.pop_back();
+      }
+
+      for(size_t j = 0; j<fShapeY.size(); ++j){
+               out << SP << "size_t "<<"idx_"<<fNY<<j<<";\n";
+      }
+      out << SP << "for(size_t idx = 0; idx < " << length << "; ++idx){\n";
+      out<< SP << SP << "idx_"<<fNY<<"0 = idx / " << stridesY[0]<<";\n";
+      compute_idx_Y += "idx_"+fNY+"0 * " + std::to_string(stridesY[0]);
+      std::string modulo_op = "idx % " + std::to_string(stridesY[0]);
+      for(size_t j = 1; j<fShapeY.size(); ++j){
+
+               out << SP << SP << "idx_"<<fNY<<j<<" = ("<<modulo_op<<") / "<<stridesY[j]<<";\n";
+               modulo_op += "% " + std::to_string(stridesY[j]);
+               compute_idx_Y = "idx_"+fNY+std::to_string(j)+" * "+std::to_string(stridesY[j])+" + "+compute_idx_Y;
+      }
+      out << SP << SP << "tensor_" << fNY <<"["<<compute_idx_Y<<"] = "<<BinaryOperatorTrait<T,Op>::Op("tensor_"+ fNA + "["+compute_idx_A+"]", "tensor_"+ fNB + "["+compute_idx_B+"]")<<" ;\n";
+
+   }  else {
+      // dynamic case with broadcasting
       // we need to check if we can broadcast (case flag has bit 4 set)
       if (fBroadcastFlag & 4) {
          // need to check if shapes are the same
@@ -337,6 +383,8 @@ public:
       out << SP << SP << "tensor_" << fNY << "[id] = "
                << BinaryOperatorTrait<T,Op>::Op( "tensor_" + nameA + "[id]" , "tensor_" + nameB + "[id]")
                <<  " ;\n";
+   }
+
       out << SP << "}\n";
       return out.str();
    }
diff --git a/tmva/sofie/inc/TMVA/ROperator_Comparision.hxx b/tmva/sofie/inc/TMVA/ROperator_Comparision.hxx
@@ -131,10 +131,6 @@ public:
                // Update the data and the shape of A
                model.UpdateInitializedTensor(fNX1, model.GetTensorType(fNX1), fShapeY, broadcastedData);
                fShapeX1 = fShapeY;
-            } else {
-               // Add an intermediate tensor for broadcasting A
-               fNBroadcastedX1 = "Broadcasted" + fNX1;
-               model.AddIntermediateTensor(fNBroadcastedX1, model.GetTensorType(fNX1), fShapeY);
             }
          }
          // Broadcast B to Y
@@ -147,10 +143,6 @@ public:
                // Update the data and the shape of B
                model.UpdateInitializedTensor(fNX2, model.GetTensorType(fNX2), fShapeY, broadcastedData);
                fShapeX2 = fShapeY;
-            } else {
-               // Add an intermediate tensor for broadcasting B
-               fNBroadcastedX2 = "Broadcasted" + fNX2;
-               model.AddIntermediateTensor(fNBroadcastedX2, model.GetTensorType(fNX2), fShapeY);
             }
          }
       } else {
@@ -264,32 +256,47 @@ public:
       out << SP << "\n//------ " << ComparisionTrait<T,Op>::Name() << "  " << opName
                                  << " --> " << ConvertShapeToString(fShapeY) << "\n";
       size_t length = ConvertShapeToLength(fShapeY);
-      // Broadcast A if it's uninitialized
-      if (!fNBroadcastedX1.empty()) {
-         std::string type1 = ConvertTypeToString(fTensorType1);
-         out << SP << "// Broadcasting uninitialized tensor " << fNX1 << "\n";
-         out << SP << "{\n";
-         out << SP << SP << type1 << "* data = TMVA::Experimental::SOFIE::UTILITY::UnidirectionalBroadcast<" << type1 << ">(tensor_" << fNX1 << ", " << ConvertShapeToString(fShapeX1) << ", " << ConvertShapeToString(fShapeY) << ");\n";
-         out << SP << SP << "std::copy(data, data + " << length << ", tensor_" << fNBroadcastedX1 << ");\n";
-         out << SP << SP << "delete[] data;\n";
-         out << SP << "}\n";
+
+      auto stridesX1 = UTILITY::ComputeStrideFromShape(fShapeX1);
+      auto stridesX2 = UTILITY::ComputeStrideFromShape(fShapeX2);
+      auto stridesY = UTILITY::ComputeStrideFromShape(fShapeY);
+
+      std::string compute_idx_X1, compute_idx_X2, compute_idx_Y;
+      if (std::all_of(fShapeX1.begin(), fShapeX1.end(), [](size_t x) { return x == 1; })){
+         compute_idx_X1 = "0";
+      } else {
+         for(size_t i = 0; i<fShapeX1.size(); ++i){
+            if(fShapeX1[i]==1) continue;
+            compute_idx_X1 += " idx_"+fNY+std::to_string(i+(fShapeY.size()-fShapeX1.size()))+" * "+stridesX1[i]+" +";
+         }
+         compute_idx_X1.pop_back();
+      }
+      if (std::all_of(fShapeX2.begin(), fShapeX2.end(), [](size_t x) { return x == 1; })){
+         compute_idx_X2 = "0";
+      } else {
+         for(size_t i = 0; i<fShapeX2.size(); ++i){
+            if(fShapeX2[i]==1) continue;
+            compute_idx_X2 += " idx_"+fNY+std::to_string(i+(fShapeY.size()-fShapeX2.size()))+" * "+stridesX2[i]+" +";
+         }
+         compute_idx_X2.pop_back();
       }
-      // Broadcast B if it's uninitialized
-      if (!fNBroadcastedX2.empty()) {
-         std::string type2 = ConvertTypeToString(fTensorType2);
-         out << SP << "// Broadcasting uninitialized tensor " << fNX2 << "\n";
-         out << SP << "{\n";
-         out << SP << SP << type2 << "* data = TMVA::Experimental::SOFIE::UTILITY::UnidirectionalBroadcast<" << type2 << ">(tensor_" << fNX2 << ", " << ConvertShapeToString(fShapeX2) << ", " << ConvertShapeToString(fShapeY) << ");\n";
-         out << SP << SP << "std::copy(data, data + " << length << ", tensor_" << fNBroadcastedX2 << ");\n";
-         out << SP << SP << "delete[] data;\n";
-         out << SP << "}\n";
+      
+      for(size_t j = 0; j<fShapeY.size(); ++j){
+               out << SP << "size_t "<<"idx_"<<fNY<<j<<";\n";
       }
-      const std::string& nameX1 = fNBroadcastedX1.empty()? fNX1 : fNBroadcastedX1;
-      const std::string& nameX2 = fNBroadcastedX2.empty()? fNX2 : fNBroadcastedX2;
-
-      out << SP << "for (size_t id = 0; id < " << length << " ; id++){\n";
-      out << SP << SP << "fTensor_" << fNY << "[id] = " << ComparisionTrait<T,Op>::Op( "tensor_" + nameX1 + "[id]" , "tensor_" + nameX2 + "[id]") <<  " ;\n";
+      out << SP << "for(size_t idx = 0; idx < " << length << "; ++idx){\n";
+      out<< SP << SP << "idx_"<<fNY<<"0 = idx / " << stridesY[0]<<";\n";
+      compute_idx_Y += "idx_"+fNY+"0 * " + std::to_string(stridesY[0]);
+      std::string modulo_op = "idx % " + std::to_string(stridesY[0]);
+      for(size_t j = 1; j<fShapeY.size(); ++j){
+               
+               out << SP << SP << "idx_"<<fNY<<j<<" = ("<<modulo_op<<") / "<<stridesY[j]<<";\n";
+               modulo_op += "% " + std::to_string(stridesY[j]);
+               compute_idx_Y = "idx_"+fNY+std::to_string(j)+" * "+std::to_string(stridesY[j])+" + "+compute_idx_Y;
+      }
+      out << SP << SP << "tensor_" << fNY <<"["<<compute_idx_Y<<"] = "<<ComparisionTrait<T,Op>::Op("tensor_"+ fNX1 + "["+compute_idx_X1+"]", "tensor_"+ fNX2 + "["+compute_idx_X2+"]")<<" ;\n";
       out << SP << "}\n";
+
       // since output is a boolean need to add the tensor_xxx variable since it is not defined as a pointer to a boolean std::vector
       if (!fIsModelOutput)
          out << SP << "const std::vector<std::uint8_t> & tensor_" << fNY << " = fTensor_" << fNY << ";\n";
diff --git a/tmva/sofie/inc/TMVA/ROperator_Conv.hxx b/tmva/sofie/inc/TMVA/ROperator_Conv.hxx
@@ -294,6 +294,7 @@ public:
                // we need to add a new intermediate tensor for broadcasted bias tensor
                fNB2 = fNB + "bcast";
                model.AddIntermediateTensor(fNB2, model.GetTensorType(fNB), targetShape);
+               fOutputTensorNames.emplace_back(fNB2);
             }
          }
       }
diff --git a/tmva/sofie/inc/TMVA/ROperator_Gemm.hxx b/tmva/sofie/inc/TMVA/ROperator_Gemm.hxx
@@ -251,6 +251,7 @@ namespace SOFIE{
                   fNC2 = fNC + "bcast";
                   if (!fIsDynamic) {
                      model.AddIntermediateTensor(fNC2, model.GetTensorType(fNC), shapeY);
+                     fOutputTensorNames.emplace_back(fNC2);
                   }
                   else
                      model.AddDynamicTensor(fNC2,model.GetTensorType(fNC), fShapeY);
diff --git a/tmva/sofie/inc/TMVA/ROperator_Where.hxx b/tmva/sofie/inc/TMVA/ROperator_Where.hxx
@@ -110,9 +110,6 @@ public:
                // Update the data and the shape of A
                model.AddConstantTensor(fNBroadcastedA, model.GetTensorType(fNA), fShapeY, broadcastedData);
                fShapeA = fShapeY;
-            } else {
-               // Add an intermediate tensor for broadcasting A
-               model.AddIntermediateTensor(fNBroadcastedA, model.GetTensorType(fNA), fShapeY);
             }
          }
          // Broadcast B to Y
@@ -126,9 +123,6 @@ public:
                // do not update tensor B but add broadcasted one (since it can be input to some other operators)
                model.AddConstantTensor(fNBroadcastedB, model.GetTensorType(fNB), fShapeY, broadcastedData);
                fShapeB = fShapeY;
-            } else {
-               // Add an intermediate tensor for broadcasting B
-               model.AddIntermediateTensor(fNBroadcastedB, model.GetTensorType(fNB), fShapeY);
             }
          }
          // Broadcast C to Y
@@ -142,9 +136,6 @@ public:
                // do not update tensor C but add broadcasted one (since it can be input to some other operators)
                model.AddConstantTensor(fNBroadcastedC, model.GetTensorType(fNC), fShapeY, broadcastedData);
                fShapeC = fShapeY;
-            } else {
-               // Add an intermediate tensor for broadcasting B
-               model.AddIntermediateTensor(fNBroadcastedC, model.GetTensorType(fNC), fShapeY);
             }
          }
       } else {
@@ -252,39 +243,64 @@ public:
       out << SP << "\n//-------- Where " << opName << " --> " << ConvertShapeToString(fShapeY) << "\n";
       size_t length = ConvertShapeToLength(fShapeY);
       std::string typeName = TensorType<T>::Name();
-      // Broadcast A if it's uninitialized
-      if (fShapeA != fShapeY) {
-         out << SP << "// Broadcasting uninitialized tensor " << fNA << "\n";
-         //out << SP << "{\n";
-         out << SP  << "TMVA::Experimental::SOFIE::UTILITY::UnidirectionalBroadcast<" << typeName << ">(tensor_" << fNA << ", " << ConvertShapeToString(fShapeA) << ", " << ConvertShapeToString(fShapeY)
-                         << ", fTensor_" << fNBroadcastedA << ");\n";
+
+      auto stridesA = UTILITY::ComputeStrideFromShape(fShapeA);
+      auto stridesB = UTILITY::ComputeStrideFromShape(fShapeB);
+      auto stridesC = UTILITY::ComputeStrideFromShape(fShapeC);
+      auto stridesY = UTILITY::ComputeStrideFromShape(fShapeY);
+
+      std::string compute_idx_A, compute_idx_B, compute_idx_C, compute_idx_Y;
+      
+      if (std::all_of(fShapeA.begin(), fShapeA.end(), [](size_t x) { return x == 1; })){
+         compute_idx_A = "0";
+      } else {
+         for(size_t i = 0; i<fShapeA.size(); ++i){
+            if(fShapeA[i]==1) continue;
+            compute_idx_A += " idx_"+fNY+std::to_string(i+(fShapeY.size()-fShapeA.size()))+" * "+stridesA[i]+" +";
+         }
+         compute_idx_A.pop_back();
       }
-      // Broadcast B if it's uninitialized
-      if (fShapeB != fShapeY) {
-         out << SP << "// Broadcasting uninitialized tensor " << fNB << "\n";
-         //out << SP << "{\n";
-         out << SP << "TMVA::Experimental::SOFIE::UTILITY::UnidirectionalBroadcast<" << typeName << ">(tensor_" << fNB << ", " << ConvertShapeToString(fShapeB) << ", " << ConvertShapeToString(fShapeY)
-                   << ", fTensor_" << fNBroadcastedB << ");\n";
+      if (std::all_of(fShapeB.begin(), fShapeB.end(), [](size_t x) { return x == 1; })){
+         compute_idx_B = "0";
+      } else {
+         for(size_t i = 0; i<fShapeB.size(); ++i){
+            if(fShapeB[i]==1) continue;
+            compute_idx_B += " idx_"+fNY+std::to_string(i+(fShapeY.size()-fShapeB.size()))+" * "+stridesB[i]+" +";
+         }
+         compute_idx_B.pop_back();
       }
-       // Broadcast C if it's uninitialized
-      if (fShapeC != fShapeY) {
-         // special case if C is an input tensor
-         if (fIsInputBoolTensor) {
-            size_t inputLength = ConvertShapeToLength(fShapeC);
-            out << SP << "std::vector<std::uint8_t> fTensor_" << fNC << "(tensor_" << fNC <<  ", tensor_" << fNC << " + " << inputLength << ");\n";
+      if (std::all_of(fShapeC.begin(), fShapeC.end(), [](size_t x) { return x == 1; })){
+         compute_idx_C = "0";
+      } else {
+         for(size_t i = 0; i<fShapeC.size(); ++i){
+            if(fShapeC[i]==1) continue;
+            compute_idx_C += " idx_"+fNY+std::to_string(i+(fShapeY.size()-fShapeC.size()))+" * "+stridesC[i]+" +";
          }
-         out << SP << "// Broadcasting uninitialized tensor " << fNC << "\n";
-         //out << SP << "{\n";
-         out << SP << "TMVA::Experimental::SOFIE::UTILITY::UnidirectionalBroadcast<std::uint8_t>(fTensor_" << fNC << ".data(), " << ConvertShapeToString(fShapeC) << ", " << ConvertShapeToString(fShapeY)
-                   << ", fTensor_" << fNBroadcastedC << ");\n";
+         compute_idx_C.pop_back();
+      }
+      
+      if (fIsInputBoolTensor) {
+            size_t inputLength = ConvertShapeToLength(fShapeC);
+            out << SP << "std::vector<bool> fTensor_" << fNC << "(tensor_" << fNC <<  ", tensor_" << fNC << " + " << inputLength << ");\n";
       }
-      std::string nameA = fNBroadcastedA.empty()? fNA : fNBroadcastedA;
-      std::string nameB = fNBroadcastedB.empty()? fNB : fNBroadcastedB;
-      std::string nameC = fNBroadcastedC.empty()? fNC : fNBroadcastedC;
-      out << SP << "for (size_t id = 0; id < " << length << " ; id++){\n";
+      
+      for(size_t j = 0; j<fShapeY.size(); ++j){
+               out << SP << "size_t "<< fNY << "idx_"<<j<<";\n";
+      }
+      out << SP << "for(size_t idx = 0; idx < " << length << "; ++idx){\n";
+      out<< SP << SP << "idx_"<<fNY<<"0 = idx / " << stridesY[0]<<";\n";
+      compute_idx_Y += "idx_"+fNY+"0 * " + std::to_string(stridesY[0]);
+      std::string modulo_op = "idx % " + std::to_string(stridesY[0]);
+      for(size_t j = 1; j<fShapeY.size(); ++j){
+               
+               out << SP << SP << "idx_"<<fNY<<j<<" = ("<<modulo_op<<") / "<<stridesY[j]<<";\n";
+               modulo_op += "% " + std::to_string(stridesY[j]);
+               compute_idx_Y = "idx_"+fNY+std::to_string(j)+" * "+std::to_string(stridesY[j])+" + "+compute_idx_Y;
+      }
+
       // get output tensor applying condition
-      out << SP << SP << "tensor_" << fNY << "[id] = "  << "(fTensor_" << nameC << "[id]) ? tensor_"
-                               << nameA << "[id] : tensor_" + nameB + "[id];\n";
+      out << SP << SP << "tensor_" << fNY << "["<<compute_idx_Y<<"] = "  << "(fTensor_" << fNC << "["<<compute_idx_C<<"]) ? tensor_"
+                               << fNA << "["<<compute_idx_A<<"] : tensor_" + fNB + "["<<compute_idx_B<<"];\n";
       out << SP << "}\n";
       return out.str();
    }

Original file line number	Diff line number	Diff line change
`@@ -294,6 +294,7 @@ public:`
`294`	`294`	`// we need to add a new intermediate tensor for broadcasted bias tensor`
`295`	`295`	`fNB2 = fNB + "bcast";`
`296`	`296`	`model.AddIntermediateTensor(fNB2, model.GetTensorType(fNB), targetShape);`
	`297`	`+ fOutputTensorNames.emplace_back(fNB2);`
`297`	`298`	`}`
`298`	`299`	`}`
`299`	`300`	`}`
Original file line number	Diff line number	Diff line change
`@@ -251,6 +251,7 @@ namespace SOFIE{`
`251`	`251`	`fNC2 = fNC + "bcast";`
`252`	`252`	`if (!fIsDynamic) {`
`253`	`253`	`model.AddIntermediateTensor(fNC2, model.GetTensorType(fNC), shapeY);`
	`254`	`+ fOutputTensorNames.emplace_back(fNC2);`
`254`	`255`	`}`
`255`	`256`	`else`
`256`	`257`	`model.AddDynamicTensor(fNC2,model.GetTensorType(fNC), fShapeY);`