Initial dynamic input support for gpu codegen by harz05 · Pull Request #45 · ML4EP/SOFIE

harz05 · 2026-06-30T10:31:11Z

Implements #43

Adds dynamic shape support to the GPU (alpaka) code generator: models whose input dimensions are symbolic (N, n_pf, n_sv) now generate, compile and run, instead ofhaving their shapes baked in at codegen time.

The operator set, tests and results are driven by particle-net.onnx, the target dynamic model for this work.

The dynamic buffer fix

Generating a dynamic model on the GPU path failed to compile. The dynamic intermediate tensor buffers were emitted with a type alias built on the runtime device object devAcc where a type is needed, declared under the name bufDev_while every operator referencesdeviceBuf_(so the buffer was undefined at use) and allocated into a discardedautolocal in the constructor with a hardcoded float type, so the member was never set. After fixing those, a second error showed up becausealpaka::Buf` has no default constructor, so the bare member declaration would not compile.

The fix declares the dynamic buffers as Session members with the correct Buf type and deviceBuf_ name according to the tensor dtype, default initializes each member to a one element placeholder buffer to satisfy the no default constructor requirement and allocates the real buffer in the constructor sized by the runtime length once N is known.

Operator changes:

Extended for dynamic shapes: Tile, Transpose, Concat, Gather, Slice, Reduce, Conv, BatchNormalization, BasicBinary, Comparision, Range. Each keeps a dual size_t/Dim representation mirroring the ROOT's SOFIE cpu operator with a dynamic Initialize branch that registers the output as a dynamic tensor without materializing it and index math driven by runtime dimensions.

BatchNormalization: the scale/variance fusion produces a per-channel [C] array rather than materializing weights to the full [N,C,...] tensor, which would hardcode in the batch size and block any dynamic shape. Batch and spatial dims are handled by the kernel's index math instead.

New GPU kernels for two operators that previously had none:

TopK": one thread per slice with a K-sized insertion-sorted register buffer; k stays a compile-time constant so the buffer is a fixed register array, while the axis length is a runtime argument.
Softmax: block-per-row online softmax (a single fused max/sum pass, then a shared-memory tree reduction). Threads per row is the axis length rounded up to a power of two, clamped to [32,1024], falling back to 256 for a dynamic axis.

Other operators:

Range: output length derived symbolically from a shape tensor instead of an opaque range_size variable that collided across ops.
Gather: a Gather indexing a shape tensor now produces a shape tensor whose value is written host-side and copied to the device buffer, so a shape tensor consumed by a compute kernel has its value available on device.
BasicNary (Max/Min/Sum/Mean): had no GPU codegen at all; the base returned "" silently, so the output was never computed on device. Added a generic elementwise kernel with per-input template types (inputs can have mixed dtypes) and per-input index decomposition for multidirectional broadcast, so differently-shaped inputs are read through their own strides instead of a shared flat index.
RModel_ALPAKA shape tensor declarations, and an _infer_impl argument order fix. The definition interleaves dynamic params with inputs, but the two call sites passed all params first. Those orders coincide when a single input introduces every symbol, so the bug only appears on multi-symbol models.

Kernel argument convention

Index math kernels that reference a runtime dimension receive the model shape parameters (for example N) as size_t kernel arguments, supplied at the createTaskKernel callsite. GetGPUDynParams() computes that parameter list once per operator and is called by both the kernel signature generation and the launch, so the two cannot change

Known gaps

Skipped deliberately with a comment rather than emitting wrong code:

Range with a fully run time size: the length expression dereferences the scalar inputs host-side, which on GPU would need a device to host read.

Tests

30 new gtests, each constructed and run at two different sizes and compared against an independent host reference:

area	tests
shape / indexing	Transpose, Concat, Tile, Gather, Slice, Range, RangeMul
reduce	ReduceSumLast, ReduceMeanMid, ReduceMaxFirst, ReduceSumMulti
conv	Conv1D, Conv1DNoBias, Conv2DNoBias
normalization	BatchNorm4D, BatchNormDynSpatialRelu, BatchNorm2D
elementwise	AddBroadcast, Equal, NegRelu
linear	Linear (Gemm+Relu fusion)
new operators	TopK, Softmax 1D/2D/3D/4D
BasicNary broadcast	Max/Min/Mean/SumMultidirectionalBroadcast

Results

140/140 alpaka gtests pass on an NVIDIA H100

Verified end to end on particle-net.onnx, the dynamic target model: generates with 0 failures, compiles under nvcc and runs.

Memory reported by the built-in profiler:

Per-operator timings, N=1, n_pf=100, n_sv=10, averaged over 100 runs:

operator class	count	total
Conv	15	~594 us
BatchNormalization	17	~182 us
TopK	3	~153 us
ReduceMean	3	~149 us
ReduceSum	6	~65 us
MatMul	3	~53 us
Gemm	2	~26 us

Conv dominates. The profiler synchronizes after every operator, so its overall figure is not a throughput measurement; the per-op breakdown is the meaningful part.

Note on multi-size sessions

A Session registers its cuBLASLt layouts at construction size, so each test here constructs a fresh Session per size. With ML4EP/sofieBLAS#11 a single Session serves multiple sizes: verified by running one ParticleNet Session at n_pf/n_sv of 100/10, 50/5 and 128/16, each returning a valid softmax.

…h tests

…no bias conv

…ticle-net op set

…input-alpaka # Conflicts: # core/inc/SOFIE/ROperator_Conv.hxx # test/TestCustomModelsFromONNXForAlpakaCuda.cxx

sanjibansg · 2026-08-18T14:28:25Z

/runtest h100

github-actions · 2026-08-18T14:28:35Z

/runtest (h100): triggered - view run

sanjibansg · 2026-08-18T14:47:39Z

/runtest h100-47gb

github-actions · 2026-08-18T14:47:50Z

/runtest (h100-47gb): triggered - view run

github-actions · 2026-08-18T14:56:02Z

/runtest (h100-47gb): GPU Unit Tests ✅ passed - view run

sanjibansg

Thanks for this very useful implementation, couple of initial comments.

sanjibansg · 2026-08-19T05:06:04Z

   }

-   std::vector<std::string> GetBlasRoutines() override { return { std::string("Copy"), std::string("Axpy") }; }
+   std::vector<std::string> GetBlasRoutines() override { return {}; }


we should keep the blas routines here anyway, since they might be needed for the cpu inference.

actually the older implementation expanded the channels weights to full size at generation time so that scopy/saxpy could operate at same vector length, however the new implementation now uses fused per channel scale and a single fused loop (this approach is the same as what we have for root/sofie); so to sum it up we would not need to have blas routines now, i kept the function here since it was as is in root/sofie's implementation but maybe its better if we don't keep this (?)

sanjibansg · 2026-08-19T18:39:00Z

      return out.str();
   }

+   // Computed once so the kernel signature and the call site cannot drift apart.


add a bit more descriptive comment here, i.e. signalling where is this computed and for what purpose.

sanjibansg · 2026-08-19T18:48:44Z

      }
      // find shape of Y and add it in the list of intermediate tensors
-      fShapeY = ShapeInference({fShapeX})[0];
+      fShapeY = DoShapeInference(fShapeX);


Instead of a differently named method, can we have an overloaded one?

sanjibansg · 2026-08-19T18:51:09Z

+      op += SP + SP + SP + "for (std::size_t l = tid; l < axis_size; l += " + bs + "u) {\n";
+      op += SP + SP + SP + SP + "T x = X[row_base + l * inner_stride];\n";
+      op += SP + SP + SP + SP + "T m_new = (x > m) ? x : m;\n";
+      op += SP + SP + SP + SP + "d = d * alpaka::math::exp(acc, m - m_new) + alpaka::math::exp(acc, x - m_new);\n";


Suggested change

op += SP + SP + SP + SP + "d = d * alpaka::math::exp(acc, m - m_new) + alpaka::math::exp(acc, x - m_new);\n";

op += SP + SP + SP + SP + "d = d * exp(acc, m - m_new) + exp(acc, x - m_new);\n";

gpu/alpaka now has SOFIE_DEVICE_exp macro, so would it be better to use that here instead of simple exp??

sanjibansg · 2026-08-19T18:51:34Z

+      op += SP + SP + SP + SP + "T e = alpaka::math::exp(acc, X[idx] - vmax) * inv;\n";
+      op += SP + SP + SP + SP + "Y[idx] = e;\n";
+      if (fLogSoftmax)
+         op += SP + SP + SP + SP + "Y[idx] = alpaka::math::log(acc, e);\n";


Suggested change

op += SP + SP + SP + SP + "Y[idx] = alpaka::math::log(acc, e);\n";

op += SP + SP + SP + SP + "Y[idx] = log(acc, e);\n";

sanjibansg · 2026-08-19T18:51:45Z

+      op += SP + SP + SP + "T const inv = static_cast<T>(1) / sum;\n";
+      op += SP + SP + SP + "for (std::size_t l = tid; l < axis_size; l += " + bs + "u) {\n";
+      op += SP + SP + SP + SP + "std::size_t const idx = row_base + l * inner_stride;\n";
+      op += SP + SP + SP + SP + "T e = alpaka::math::exp(acc, X[idx] - vmax) * inv;\n";


Suggested change

op += SP + SP + SP + SP + "T e = alpaka::math::exp(acc, X[idx] - vmax) * inv;\n";

op += SP + SP + SP + SP + "T e = exp(acc, X[idx] - vmax) * inv;\n";

sanjibansg · 2026-08-19T18:52:11Z

+      op += SP + SP + SP + SP + SP + "T m_a = smax[tid];\n";
+      op += SP + SP + SP + SP + SP + "T m_b = smax[tid + s];\n";
+      op += SP + SP + SP + SP + SP + "T m_r = (m_b > m_a) ? m_b : m_a;\n";
+      op += SP + SP + SP + SP + SP + "ssum[tid] = ssum[tid] * alpaka::math::exp(acc, m_a - m_r) + ssum[tid + s] * alpaka::math::exp(acc, m_b - m_r);\n";


Suggested change

op += SP + SP + SP + SP + SP + "ssum[tid] = ssum[tid] * alpaka::math::exp(acc, m_a - m_r) + ssum[tid + s] * alpaka::math::exp(acc, m_b - m_r);\n";

op += SP + SP + SP + SP + SP + "ssum[tid] = ssum[tid] * exp(acc, m_a - m_r) + ssum[tid + s] * exp(acc, m_b - m_r);\n";

sanjibansg · 2026-08-19T18:54:28Z

+      out << SP << "auto const elementsPerThread_" << fNVal << " = Vec::all(static_cast<Idx>(1));\n";
+      out << SP << "auto const elementsPerGrid_"   << fNVal << " = Vec::all(static_cast<Idx>(" << numSlices << "));\n";
+      out << SP << "auto const workDiv_" << fNVal << " = sofie_workdiv(elementsPerGrid_" << fNVal << ");\n";
+      out << SP << "alpaka::exec<Acc>(queue, workDiv_" << fNVal << ", topKernel_" << fNVal


can we instead create a task kernel and enqueue it?

sanjibansg · 2026-08-19T18:56:08Z

-      // for(int i=0;i<fShapeX.size();i++)
-      // std::cout<<fShapeX[i]<<" ";
-      // std::cout<<"\ny size -> "<<fShapeY.size()<<std::endl;
+      fK = fShapeX[fAttrAxis].isParam ? kval : std::min(kval, fShapeX[fAttrAxis].dim);


fK seems like an important variable, be a bit more descriptive in its naming.

sanjibansg · 2026-08-19T18:56:50Z

+
+      op += SP + SP + SP + "for (int64_t l = " + K + "; l < (int64_t)nElAxis; ++l) {\n";
+      op += SP + SP + SP + SP + "T v = x[xbase + strideXAxis * (std::size_t)l];\n";
+      op += SP + SP + SP + SP + "if (v " + CMP + " bestV[" + K + "-1]) {\n";


if cases are critical in GPUs, can we have a workaround here to avoid branching?

…input-alpaka # Conflicts: # core/src/RModel_ALPAKA.cxx # test/alpaka/TestAlpakaShapeIndexing.cxx

…xing codegen bugs

…on GPU

harz05 · 2026-08-25T10:13:02Z

-      // Extend chain: only if CURRENT op is elementwise and its single output can be fused
      size_t current = i;
-      while (fOperators[current]->IsElementwise()) {
+      while (fOperators[current]->IsElementwise() && !fOperators[current]->IsOutputConstant()) {


remove 2nd condition. If IsElementwise() is true already then we won't need to check if the output is constant as it is not constant for current ops

harz05 · 2026-08-25T10:20:17Z

         if (nextIdx != current + 1) break;
         if (opAssigned[nextIdx]) break;
-         if (!fOperators[nextIdx]->IsElementwise()) break;
+         if (!fOperators[nextIdx]->IsElementwise() || fOperators[nextIdx]->IsOutputConstant()) break;


same as line 93

harz05 · 2026-08-25T11:14:04Z

+
+void EmitOutputCoords(std::string &op, const std::string &indent,


better descp for its use case in Tile, Comparison, etc

harz05 · 2026-08-25T11:14:34Z


+UTILITY::SliceInfo UTILITY::ComputeSliceInfo(const std::vector<Dim> & shape, size_t axis) {


topK, softmax use case descp as a comment

harz05 · 2026-08-25T11:33:45Z

-      op += SP + SP + SP + SP + "std::size_t const col_row = elem_idx / " + std::to_string(colCols) + "u;\n";
-      op += SP + SP + SP + SP + "std::size_t const col_col = elem_idx % " + std::to_string(colCols) + "u;\n\n";
+      op += SP + SP + SP + SP + "std::size_t const col_row = elem_idx / (oDepth * oHeight * oWidth);\n";
+      op += SP + SP + SP + SP + "std::size_t const col_col = elem_idx % (oDepth * oHeight * oWidth);\n\n";


restore to unsigned int

harz05 · 2026-08-25T11:35:19Z

-         op += SP + SP + SP + SP + "std::size_t const od = col_col / " + std::to_string(oHeight * oWidth) + "u;\n";
-         op += SP + SP + SP + SP + "std::size_t const oh = (col_col / " + std::to_string(oWidth) + "u) % " + std::to_string(oHeight) + "u;\n";
-         op += SP + SP + SP + SP + "std::size_t const ow = col_col % " + std::to_string(oWidth) + "u;\n\n";


same as L713

harz05 · 2026-08-25T12:59:34Z

   }

   // shape of output tensors given input tensors
-   std::vector<std::vector<size_t>> ShapeInference(std::vector<std::vector<size_t>> input) override {


base ShapeInference change to Dim based, replacing the DoShapeInference func

harz05 · 2026-08-25T13:07:26Z

+   std::string rGC;
+   ForEachInferArg_GPU_ALPAKA(
+      [&](const std::string &p) {
+         if (isdecl) rGC += "size_t ";


size_t needs to be reviewed

harz05 · 2026-08-25T13:21:57Z

      SOFIE::OperatorKind::SIGMOID,
      SOFIE::OperatorKind::TANH,
-      SOFIE::OperatorKind::SOFTMAX,
      SOFIE::OperatorKind::LEAKYRELU,


harz05 · 2026-08-25T13:44:33Z

      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 lengthExpr;        ///< element count: literal for static tensors, runtime expression for dynamic


review once

…e helpers

harz05 added 7 commits June 29, 2026 13:10

dynamic tensor allocation

59a3568

testing: dynamic support for Tile

1c083ca

test added for dynamic input

3b731e2

dynamic shape support for comparision op

ed22f6f

onnx models added and regression tests for dynamic Tile and Equal

07bafea

dynamic input for batchnorm op and test cases

dba8052

dynamic input for conv op and tests

78eb28e

harz05 force-pushed the feat/dynamic-input-alpaka branch from 2946253 to 78eb28e Compare July 5, 2026 17:38

harz05 added 4 commits July 5, 2026 23:41

emit dynamic session ctor params in declaration order to match infer

772ab95

dynamic input for binary, gather, transpose, concat and slice ops wit…

164feda

…h tests

fix for conv output shape expression for stride > 1

f7023d9

dynamic support for fused kernels and tests for gemm relu fusion and …

8ad9749

…no bias conv

harz05 mentioned this pull request Jul 7, 2026

Support runtime GEMM dimensions in the cuBLASLt layout store for dynamic shapes ML4EP/sofieBLAS#10

Open

dynamic input for reduce op and test cases

e8623a1

harz05 force-pushed the feat/dynamic-input-alpaka branch from 01ddd4c to e8623a1 Compare July 8, 2026 16:15

add diagnostic logging

e56dbf8

harz05 marked this pull request as draft July 10, 2026 18:10

harz05 added 3 commits July 18, 2026 14:29

chore: remove init-fail diagnostic logging

0efee60

feat: dynamic-input GPU codegen for Shape/Range/BasicNary and the par…

ac7c311

…ticle-net op set

test: add dynamic-input GPU operator tests

c5adb72

harz05 marked this pull request as ready for review July 18, 2026 09:30

harz05 added 4 commits August 18, 2026 03:18

Merge remote-tracking branch 'upstream/gpu/alpaka' into feat/dynamic-…

7aacb17

…input-alpaka # Conflicts: # core/inc/SOFIE/ROperator_Conv.hxx # test/TestCustomModelsFromONNXForAlpakaCuda.cxx

chore: in line comments cleanup in the code

9cb225f

fix: add missing reference headers for softmax and TopK tests

f285d59

feat: broadcast support in BasicNary GPU codegen

0cf26ae

sanjibansg requested changes Aug 19, 2026

View reviewed changes

harz05 added 5 commits August 21, 2026 03:51

Merge remote-tracking branch 'upstream/gpu/alpaka' into feat/dynamic-…

0768333

…input-alpaka # Conflicts: # core/src/RModel_ALPAKA.cxx # test/alpaka/TestAlpakaShapeIndexing.cxx

refactor: address review comments on naming, launch style and comments

fc4dd50

refactor: modifying dynamic input codegen, share shape helpers and fi…

dd0600b

…xing codegen bugs

fix: guard dynamic sizes at infer and finish the run-time Range path …

a55fccb

…on GPU

style: drop alpaka::math from exp and log in the Softmax kernel

625d58c

harz05 commented Aug 25, 2026

View reviewed changes

refactor: pass dyn params to the gpu emitters and rename the Dim shap…

390814a

…e helpers

	op += SP + SP + SP + SP + "d = d * alpaka::math::exp(acc, m - m_new) + alpaka::math::exp(acc, x - m_new);\n";
	op += SP + SP + SP + SP + "d = d * exp(acc, m - m_new) + exp(acc, x - m_new);\n";

	op += SP + SP + SP + SP + "Y[idx] = alpaka::math::log(acc, e);\n";
	op += SP + SP + SP + SP + "Y[idx] = log(acc, e);\n";

	op += SP + SP + SP + SP + "T e = alpaka::math::exp(acc, X[idx] - vmax) * inv;\n";
	op += SP + SP + SP + SP + "T e = exp(acc, X[idx] - vmax) * inv;\n";

	op += SP + SP + SP + SP + SP + "ssum[tid] = ssum[tid] * alpaka::math::exp(acc, m_a - m_r) + ssum[tid + s] * alpaka::math::exp(acc, m_b - m_r);\n";
	op += SP + SP + SP + SP + SP + "ssum[tid] = ssum[tid] * exp(acc, m_a - m_r) + ssum[tid + s] * exp(acc, m_b - m_r);\n";


		void EmitOutputCoords(std::string &op, const std::string &indent,


		UTILITY::SliceInfo UTILITY::ComputeSliceInfo(const std::vector<Dim> & shape, size_t axis) {

Uh oh!

Conversation

harz05 commented Jun 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

The dynamic buffer fix

Operator changes:

Kernel argument convention

Known gaps

Tests

Results

Note on multi-size sessions

Uh oh!

sanjibansg commented Aug 18, 2026

Uh oh!

github-actions Bot commented Aug 18, 2026

Uh oh!

sanjibansg commented Aug 18, 2026

Uh oh!

github-actions Bot commented Aug 18, 2026

Uh oh!

github-actions Bot commented Aug 18, 2026

Uh oh!

sanjibansg left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

harz05 commented Jun 30, 2026 •

edited

Loading