[ayanshul-amd]

Summary

Adds hip-fold-batchnorm, which folds an inference hip.batch_norm with constant parameters into the filter and bias of the hip.conv that feeds it. The conv can feed the batch norm directly or through a hip.add of two convs. With the batch norm gone, a following Relu fuses into the conv kernel.

On iri-v1, all 72 Relu(BN(Add(Conv, Conv))) sites fold, and latency improves by about 2% (215.8 ms to 211.3 ms).

Related issue or design

None.

Why

Inference BatchNorm is the per-channel affine map y = s*x + t, with s = scale / sqrt(var + eps) and t = B - mean*s. When it follows a conv, it can be absorbed into the conv weights at compile time.

Without the fold, each site runs a standalone hip_batch_norm kernel followed by a standalone Relu. That is two extra full read/write passes over the activation. It also blocks fusion: fuse-rocmlir only grows along pointwise successors, and hip.batch_norm is not pointwise.

What

• The rewrite.
  • Scales each conv filter by s per output channel, and scales any existing bias by s.
  • Adds t to exactly one conv's bias, so it isn't counted twice through the Add.
  • Leaves the Add in place as the fused epilogue.
  • Writes the folded weights as new inline hip.constants and leaves the originals for any other users.
  • Works for grouped convs and for f16, bf16, f32 and f64. The math is done in double and rounded to the weight type.
• When it applies.
  • Each intermediate result (the conv, or the Add and both of its convs) must have the batch norm as its only user.
  • Filters, biases and BN parameters must be compile-time constants.
  • The pass reads inline and file-backed hip.constants only. A batch norm over memory-address-backed weights is left in place, the same rule hip-decompose-conv-transpose and the NhwcConv conversion already follow. With the default 127-byte inline threshold, real weights arrive memory-backed, so the fold needs XLNX_ABI_2_0_CLONE_EXTERNAL_DATA_THRESHOLD (or the per-session XLNX_model_clone_external_data_threshold option) raised.
• Placement. It runs right after convert-onnx-to-hip in buildOnnxToHipPipelineHead, before the AfterConvertOnnxToHip plugin slot. The head is shared by the default and rocMLIR flows, so in both the fold runs before any conv rewrite or kernel outlining.
• Shared reader. DecomposeConvTranspose and NhwcConvConversion each had a private copy of the same constant reader. Both now use a shared readConstantData helper (include/hip/Dialect/IR/HipConstantData.h), which the new pass also uses. Their behavior is unchanged.
• Docs. Updated docs/pipeline_pass_menu.md (pass table and pipeline order) and docs/design/constant-handling-design.md.

Compiler IR

Before / after

// Before
%a = hip.conv(%ctx) ins(%x, %wa : ..., tensor<128x128x3x3xf32>) ...
%b = hip.conv(%ctx) ins(%y, %wb : ..., tensor<128x384x3x3xf32>) ...
%s = hip.add(%ctx) ins(%a, %b : ...) ...
%n = hip.batch_norm(%ctx) ins(%s, %scale, %B, %mean, %var : ...) ...
%r = hip.max(%ctx) ins(%n, %zero : ...) ...

// After
%wa2 = hip.constant {value = dense<s * wa>} : tensor<128x128x3x3xf32>
%t   = hip.constant {value = dense<t>} : tensor<128xf32>
%a = hip.conv(%ctx) ins(%x, %wa2, %t : ...) ...
%wb2 = hip.constant {value = dense<s * wb>} : tensor<128x384x3x3xf32>
%b = hip.conv(%ctx) ins(%y, %wb2 : ...) ...
%s = hip.add(%ctx) ins(%a, %b : ...) ...
%r = hip.max(%ctx) ins(%s, %zero : ...) ...

Performance

Metric (iri-v1)	Before	After
Average latency	215.8 ms	211.3 ms
P50 latency	217.7 ms	212.9 ms
hip.batch_norm ops	72	0
Standalone Relu (hip.max)	76	4

Hardware, workload, build configuration, and measurement method:

• Hardware: Strix Halo (gfx1151).
• Model: iri-v1 f32 800x800.
• Build: Release, ENABLE_ROCMLIRTRITON=ON, --hip_arch gfx1151. Before and after are both on base 787825af.
• Run settings: HIPDNN_EP_ROCMLIR=1, HIPDNN_EP_ROCMLIR_AUTOTUNE=quick, HIPDNN_EP_STRICT=1, XLNX_ABI_2_0_CLONE_EXTERNAL_DATA_THRESHOLD=1073741824.
• Measurement: onnxruntime_perf_test -I -t 30 -c 1, runs done serially.
• Op counts are from the post-head IR dump.

The gain is smaller than a DRAM-bandwidth estimate predicted. A proxy model at the same shapes timed the removed kernels at:

• 1×128×100×100: 21 µs (BN) + 26 µs (Relu) per site.
• 1×384×100×100: 45 µs + 126 µs per site.

That is about 8 ms in total across the 36 + 36 sites. The 128-channel tensors are largely served from the GPU's last-level cache. After the fold, rocMLIR conv/GEMM kernels take about 129 ms of roughly 194 ms of kernel time per inference, and layout copies, transposes and pads take about 50 ms.

Test plan

• test/lit/Transforms/test_fold_batchnorm.mlir (new):
  • 5 folding cases: conv without bias, conv with bias, Add(Conv, Conv), grouped, f16.
  • 3 cases that must be left alone: multi-use conv result, non-constant BN parameter, memory-address filter.
• LIT suite: 661 passed, 9 unsupported.
• test/numeric/tests/test_conv_batchnorm.py (new) covers f32/f16 × with or without bias × group 1/4, plus Relu(BN(Add(Conv, Conv))) in f32/f16. Run with HIPDNN_EP_STRICT=1 and --ep-option XLNX_model_clone_external_data_threshold=1073741824, together with test_conv.py:
  • default flow: 16 passed;
  • HIPDNN_EP_ROCMLIR=1: 16 passed.
• Confirmed in the IR dump after the pass:
  • With the 1 GiB threshold, the batch norm is folded.
  • With the 127-byte default, the BN is left in place (memory-backed weights) and the output still matches CPU.
• iri-v1 compared with the CPU EP at both thresholds: max abs error 1.8e-6 / 1.9e-6, cosine 1.0.
• pre-commit run on the changed files: all hooks pass.

Notes for reviewers

• hipDNN. When a hipDNN handle is active, outline-onnx-to-hipdnn claims onnx.Conv before conversion, so those convs never become hip.conv and the fold doesn't see them.
• Default threshold. At the default 127-byte inline threshold, weights are memory-backed, so neither this fold nor DecomposeConvTranspose fires. On iri-v1 this leaves the 3 decoder ConvTransposes undecomposed. That is very likely most of the gap between 625 ms at the default threshold and 211 ms at 1 GiB. This predates this PR.
• Autotune bug (separate). With HIPDNN_EP_ROCMLIR_AUTOTUNE=quick, a plain 16→64 channel 3x3 conv at 8x8 returns all zeros, even without any BatchNorm. It's correct without autotune and at 32x32. The rocMLIR numeric results above were therefore run without autotune. I'll file this separately.

Checklist

• The change is focused, or links a design/series explaining its scope.
• Relevant tests were added or updated and the results are documented.
• User-facing or design documentation was updated when needed.
• Substantial AI assistance is disclosed, and I reviewed and understand the result.

AI assistance: implementation, tests, benchmarking and this description were drafted with Cursor (Claude Opus 5.5). They were validated with the LIT, numeric and iri-v1 runs listed above.

[github-actions]

Thanks for opening a PR!

This project follows LLVM's incremental-development and AI-tool-use
guidance. See CONTRIBUTING.md
for the project workflow.

Before requesting review, please check that:

1. The change is focused. Substantial work links the relevant issue
   or design discussion.
2. The PR documents relevant test results and updates affected
   documentation.
3. If AI tools provided substantial assistance, the description
   explains what was assisted and how it was validated, and commit
   trailers identify the tool. The contributor has reviewed and
   understands the result.

Reviewers are assigned through
CODEOWNERS where ownership
is configured.

[github-actions]

L2 Accuracy Results (EP vs CPU)

Model	Combined L2	Total Elems	Skipped NaN/Inf
conv_test_hybrid	1.12256E-07	64	0
GroupQueryAttention_seq256	59.2561	2621440	0
MatMulNBits_o_seq128	259.905	368640	0
QMoE_seq128	24.2583	368640	0

Threshold: 0.01 | Run: 7922 - Commit: b90744c

[github-actions]

Binary sizes

Baseline: run 40294400486 - 7579ebf (dll-sizes)

File	This (MB)	Main (MB)	Delta (MB)
hipgpu.dll	59.32	59.31	+0.01
custom_kernels_gfx1170.dll	5.55	5.55	0.00
custom_kernels_gfx11-generic.dll	5.79	5.79	0.00

MorphiZen EP Performance Results

Model	QPS	Session (s)	1st Infer (ms)	CPU%	Mem (MB)
full_model_seq128	8.77	7.58	250	3	1158
GroupQueryAttention_seq128	11234.60	2.55325	1	5	229
matmul_down_seq128	515.99	3.25	127	3	274

EPContext Export Performance

Model	QPS	Session (s)	1st Infer (ms)	CPU%	Mem (MB)
full_model_seq128	8.83	48.30	252	3	15505

EPContext Import Performance

Model	QPS	Session (s)	1st Infer (ms)	CPU%	Mem (MB)
full_model_seq128	8.83	10.12	251	3	15638

OGA Benchmark Results

Model	Warmup	Reps	Prompt Len	Gen Tokens	TTFT (ms)	TPS	Peak Mem (GB)	GPU Mem (GB)
gpt-oss-20b-webgpu-int4-rtn-block-32	1	5	128	128	92.4	74.3	1.25	12.24
Llama-3.1-8B-awq-g128-int4-asym-fp16-onnx-dml	1	5	128	128	118.1	47.5	1.15	5.23

OGA Wheel Smoke (Python benchmark_e2e.py)

Model	TTFT (ms)	TPS
Llama-3.1-8B-awq-g128-int4-asym-fp16-onnx-dml	119	48.1

VLM Benchmark Results

Model	TTFT (ms)	TPS	Output
gemma-4-E4B-it-fp16-int4-kquant-prune-logits	4049.8	7.2	This is a high-quality, outdoor photograph featuring a beautiful, fluffy white dog, likely a West Highland White Terrier or a si
Qwen3.8-27B-fp16-ve-fp16-int4-k_quant-gs128-text-dml-onnx	14592.8	14.0	The user wants a detailed description of the image. Let me carefully observe all elements before writing. Subject: A white do

Run: 7922 - Commit: b90744c