perf(transforms): fold BatchNorm into Conv weights

docs/design/constant-handling-design.md

@@ -214,6 +214,15 @@ Direct/textual pass invocation rejects memory-address carriers before
 dereferencing them. Keeping the handoff zero-copy avoids duplicating
 potentially multi-gigabyte weights throughout shape inference.
 
+Passes that fold weight values at compile time read a carrier's payload through
+`readConstantData` (`include/hip/Dialect/IR/HipConstantData.h`), which handles
+inline values and file-backed ranges and declines memory-address carriers.
+`hip-fold-batchnorm`, `hip-decompose-conv-transpose`, and the NhwcConv
+conversion all use it, so they fold only weights at or below the importer's
+inline threshold (raise `XLNX_ABI_2_0_CLONE_EXTERNAL_DATA_THRESHOLD` to bring
+all initializers inline). A fold emits its result as a new inline carrier, so
+only the rewritten weights are copied.
+
 For externalized constants, the pass writes raw bytes to the configured
 constants file via `fs`, and emits the existing module attributes:
 

docs/pipeline_pass_menu.md

@@ -137,6 +137,7 @@ the selected perfConfig string are printed only with `--verbose`.
 | `hip-infer-loop-body-shapes` | module | Rank-establish unranked tensors inside outlined loop bodies. |
 | `outline-onnx-to-hipdnn` | module | Outline subgraphs targeted at hipDNN graph compilation. |
 | `convert-onnx-to-hip` | module | Pattern-match ONNX ops by name → HIP dialect; lower both constant sweeps to policy-neutral `hip.constant`; legalize GatherBlockQuantized (INT4 packing, unsigned, quantize_axis). |
+| `hip-fold-batchnorm` | func.func | Fold an inference `hip.batch_norm` with constant parameters into the filter and bias of the `hip.conv` producing it, directly or through a `hip.add` of two convs (the shift goes into one bias). Every intermediate result must have the batch norm as its only consumer. Reads inline and file-backed constants only; a batch norm over memory-address-backed weights is left in place. |
 | `hip-externalize-constants` | module | Validate explicit serialization order, append unordered carriers deterministically, plan/serialize, then replace `hip.constant` using the existing inline/full/streaming/hybrid metadata contract. Direct threshold default: 0. |
 | `hip-infer-shapes` | module | Refine `?` dims on HIP DPS result types via `ReifyRankedShapedTypeOpInterface`. |
 | `hip-fusion-transform` | module | HIP-to-HIP rewriting at unrestricted arity (one op → many, many → one, or subgraph → subgraph). The patterns live in `lib/Dialect/Transforms/fusion_pattern/` — PDLL files compiled into the library, plus native C++ for rewrites whose shape depends on the matched op's rank — and share one `RewritePatternSet`. |
@@ -237,6 +238,8 @@ ONNX → HIP  (buildOnnxToHipPipeline)
   «slot: AfterOnnxLoopOutline»
   [outline-onnx-to-hipdnn + compile-hipdnn-graphs]   (handle overload only)
   convert-onnx-to-hip
+  func.func(hip-fold-batchnorm)      (shared with the rocMLIR head, so it runs
+                                      before any conv rewrite or kernel outlining)
   «slot: AfterConvertOnnxToHip»  (supported hip.constant producer boundary)
   hip-infer-shapes
   hip-fusion-transform               (must precede externalization: its patterns

include/hip/Dialect/IR/HipConstantData.h

@@ -0,0 +1,37 @@
+/*
+ * Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved.
+ * Licensed under the MIT License.
+ */
+#ifndef HIP_DIALECT_IR_HIP_CONSTANT_DATA_H
+#define HIP_DIALECT_IR_HIP_CONSTANT_DATA_H
+
+#include "hip/Dialect/IR/HipDialect.h"
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Support/LLVM.h"
+
+namespace mlir {
+namespace hip {
+
+/// Read the payload of a `hip.constant` carrier as a `DenseElementsAttr` of
+/// `type`, whatever source it holds.
+///
+/// "Compile-time constant" is not the same as "inline `value`": the importer
+/// externalizes initializers past a small byte threshold, so an ordinary weight
+/// reaches the compiler as a file-backed byte range or a memory-address range
+/// rather than an inline attribute.
+///
+/// Memory-address carriers are never resolved. That address is process-local
+/// and valid only while the producer that recorded it is live, which a pass
+/// cannot establish; textual pass invocation would dereference it blind.
+///
+/// Fails (without diagnostics) when the payload does not match `type` or the
+/// source cannot be read.
+FailureOr<DenseElementsAttr> readConstantData(ConstantOp constant,
+                                              RankedTensorType type);
+
+} // namespace hip
+} // namespace mlir
+
+#endif // HIP_DIALECT_IR_HIP_CONSTANT_DATA_H

include/hip/Dialect/Transforms/Passes.td

@@ -237,6 +237,26 @@ def DecomposeConvTransposePass
   ];
 }
 
+def FoldBatchNormPass
+    : Pass<"hip-fold-batchnorm", "mlir::func::FuncOp"> {
+  let summary = "Fold constant-parameter hip.batch_norm into hip.conv weights";
+  let description = [{
+    Inference BatchNorm with constant parameters is the per-channel affine map
+    `y = s * x + t`, `s = scale / sqrt(var + eps)`, `t = B - mean * s`. When
+    its input is a `hip.conv`, or a `hip.add` of two `hip.conv`s, the scale is
+    multiplied into each filter per output channel and the shift is added to
+    one conv's bias (any existing biases are scaled by `s`). The batch norm
+    disappears, so a following pointwise op fuses straight into the conv.
+
+    Every intermediate result must have the batch norm as its only consumer,
+    and the filters, biases, and batch-norm parameters must be readable at
+    compile time. The rewritten filters are emitted as new inline
+    `hip.constant`s; the originals are left for their other users, if any.
+    Inline and file-backed carriers are read; memory-address carriers are
+    not, so a batch norm over memory-backed weights is left in place.
+  }];
+}
+
 def PromoteStridedHipOperandsPass
     : Pass<"hip-promote-strided-operands", "mlir::func::FuncOp"> {
   let summary = "Materialize identity-layout temporaries for DPS-input memrefs "

lib/Conversion/OnnxToHip/NhwcConvConversion.cpp

@@ -5,10 +5,9 @@
 
 #include "OnnxToHipUtils.h"
 
+#include "hip/Dialect/IR/HipConstantData.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "llvm/Support/FileSystem.h"
-#include "llvm/Support/MemoryBuffer.h"
 
 #include <cstring>
 
@@ -61,47 +60,6 @@ Value convOutExtent(PatternRewriter &rewriter, Location loc, Value inExtent,
   return arith::AddIOp::create(rewriter, loc, divided, i64(1));
 }
 
-// Inline and file-backed filters are readable here. Memory-address carriers
-// are process-local and are not dereferenced from a pass that also runs on
-// textual IR, so those stay a runtime transpose.
-FailureOr<DenseElementsAttr> readWeightConstant(ConstantOp constant,
-                                                RankedTensorType type) {
-  switch (constant.getSourceKind()) {
-  case ConstantOp::SourceKind::Inline: {
-    auto value = dyn_cast<DenseElementsAttr>(constant.getValueAttr());
-    if (!value || value.getType() != type)
-      return failure();
-    return value;
-  }
-  case ConstantOp::SourceKind::Memory:
-    return failure();
-  case ConstantOp::SourceKind::File:
-    break;
-  }
-
-  StringRef path = constant.getLocationAttr().getValue();
-  const int64_t offset = constant.getOffsetAttr().getInt();
-  const int64_t size = constant.getSizeAttr().getInt();
-  const int64_t elementBytes = (type.getElementTypeBitWidth() + 7) / 8;
-  if (size != type.getNumElements() * elementBytes)
-    return failure();
-  uint64_t fileSize = 0;
-  if (llvm::sys::fs::file_size(path, fileSize) ||
-      static_cast<uint64_t>(offset) + static_cast<uint64_t>(size) > fileSize)
-    return failure();
-  llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> buffer =
-      llvm::MemoryBuffer::getFileSlice(path, static_cast<uint64_t>(size),
-                                       offset);
-  if (!buffer || (*buffer)->getBufferSize() != static_cast<uint64_t>(size))
-    return failure();
-  DenseElementsAttr value = DenseElementsAttr::getFromRawBuffer(
-      type,
-      ArrayRef<char>((*buffer)->getBufferStart(), static_cast<size_t>(size)));
-  if (!value)
-    return failure();
-  return value;
-}
-
 // [M, kH, kW, C/group] -> [M, C/group, kH, kW], preserving element bits.
 FailureOr<DenseElementsAttr> transposeOhwiToOihw(DenseElementsAttr value,
                                                  RankedTensorType srcType) {
@@ -289,8 +247,11 @@ LogicalResult NhwcConvToHip::matchAndRewrite(Operation *op,
   ConstantOp deadWeight;
   if (weightType.hasStaticShape()) {
     if (auto weightConst = weights.getDefiningOp<ConstantOp>()) {
+      // Memory-address carriers are process-local and are not dereferenced
+      // from a pattern that also runs on textual IR, so those stay a runtime
+      // transpose.
       FailureOr<DenseElementsAttr> data =
-          readWeightConstant(weightConst, weightType);
+          readConstantData(weightConst, weightType);
       FailureOr<DenseElementsAttr> transposed =
           succeeded(data) ? transposeOhwiToOihw(*data, weightType)
                           : FailureOr<DenseElementsAttr>(failure());

lib/Dialect/IR/CMakeLists.txt

@@ -5,6 +5,7 @@
 
 add_library(HipDialectIR STATIC
   HipCausalConvCanonicalize.cpp
+  HipConstantData.cpp
   HipDialect.cpp
   HipDpsOpInterface.cpp
   HipOpStateInterface.cpp

lib/Dialect/IR/HipConstantData.cpp

@@ -0,0 +1,60 @@
+/*
+ * Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved.
+ * Licensed under the MIT License.
+ */
+
+#include "hip/Dialect/IR/HipConstantData.h"
+
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/MemoryBuffer.h"
+
+#include <cstdint>
+#include <memory>
+
+using namespace mlir;
+using namespace mlir::hip;
+
+FailureOr<DenseElementsAttr>
+mlir::hip::readConstantData(ConstantOp constant, RankedTensorType type) {
+  switch (constant.getSourceKind()) {
+  case ConstantOp::SourceKind::Inline: {
+    // The verifier ties an inline value to the carrier's result type, but a
+    // failed match has to be cheaper than an assertion if it ever does not.
+    auto value = dyn_cast<DenseElementsAttr>(constant.getValueAttr());
+    if (!value || value.getType() != type)
+      return failure();
+    return value;
+  }
+  case ConstantOp::SourceKind::Memory:
+    return failure();
+  case ConstantOp::SourceKind::File:
+    break;
+  }
+
+  const int64_t size = constant.getSizeAttr().getInt();
+  const int64_t elementBytes = (type.getElementTypeBitWidth() + 7) / 8;
+  if (!type.hasStaticShape() || size != type.getNumElements() * elementBytes)
+    return failure();
+
+  StringRef path = constant.getLocationAttr().getValue();
+  const int64_t offset = constant.getOffsetAttr().getInt();
+  // The range is validated against the file before mapping: a slice running
+  // past the end can otherwise be mapped and fault on access rather than
+  // failing here.
+  uint64_t fileSize = 0;
+  if (offset < 0 || llvm::sys::fs::file_size(path, fileSize) ||
+      static_cast<uint64_t>(offset) + static_cast<uint64_t>(size) > fileSize)
+    return failure();
+  llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> buffer =
+      llvm::MemoryBuffer::getFileSlice(path, static_cast<uint64_t>(size),
+                                       offset);
+  if (!buffer || (*buffer)->getBufferSize() != static_cast<uint64_t>(size))
+    return failure();
+  ArrayRef<char> raw((*buffer)->getBufferStart(), static_cast<size_t>(size));
+  if (!DenseElementsAttr::isValidRawBuffer(type, raw))
+    return failure();
+  DenseElementsAttr value = DenseElementsAttr::getFromRawBuffer(type, raw);
+  if (!value)
+    return failure();
+  return value;
+}

lib/Dialect/Transforms/CMakeLists.txt

@@ -33,6 +33,7 @@ add_library(HipDialectTransforms STATIC
   DecomposeConvTranspose.cpp
   ExternalizeConstants.cpp
   FixLoopAccumulatorOffset.cpp
+  FoldBatchNorm.cpp
   FuseROCMlir.cpp
   GenerateInterface.cpp
   GenerateOpStateInit.cpp

lib/Dialect/Transforms/DecomposeConvTranspose.cpp

@@ -3,6 +3,7 @@
  * Licensed under the MIT License.
  */
 
+#include "hip/Dialect/IR/HipConstantData.h"
 #include "hip/Dialect/IR/HipDialect.h"
 #include "hip/Dialect/Transforms/Passes.h"
 
@@ -13,11 +14,7 @@
 #include <mlir/Support/LLVM.h>
 #include <mlir/Transforms/GreedyPatternRewriteDriver.h>
 
-#include <llvm/Support/FileSystem.h>
-#include <llvm/Support/MemoryBuffer.h>
-
 #include <algorithm>
-#include <memory>
 #include <numeric>
 
 namespace mlir::hip {
@@ -113,62 +110,6 @@ int64_t flatIndex(ArrayRef<int64_t> shape, int64_t a, int64_t b, int64_t c,
   return ((a * shape[1] + b) * shape[2] + c) * shape[3] + d;
 }
 
-/// Read a filter carrier's payload, whatever source it holds.
-///
-/// "Compile-time constant" is not the same as "inline `value`": the importer
-/// externalizes initializers past a small byte threshold, so an ordinary
-/// ConvTranspose filter reaches this pass as a file-backed byte range rather
-/// than a `DenseElementsAttr`. Requiring an inline value would therefore
-/// decline nearly every real filter, and a ConvTranspose-only module would
-/// then reach the compiler with no rocMLIR kernel in it at all.
-///
-/// Memory-address carriers are deliberately not resolved. That address is
-/// process-local and valid only while the producer that recorded it is live,
-/// which this pattern cannot establish -- it is equally reachable from textual
-/// pass invocation, where the address would be dereferenced blind.
-FailureOr<DenseElementsAttr> readConstant(ConstantOp constant,
-                                          RankedTensorType type) {
-  switch (constant.getSourceKind()) {
-  case ConstantOp::SourceKind::Inline: {
-    // The verifier ties an inline value to the carrier's result type, but a
-    // failed match has to be cheaper than an assertion if it ever does not.
-    auto value = dyn_cast<DenseElementsAttr>(constant.getValueAttr());
-    if (!value || value.getType() != type)
-      return failure();
-    return value;
-  }
-  case ConstantOp::SourceKind::Memory:
-    return failure();
-  case ConstantOp::SourceKind::File:
-    break;
-  }
-
-  StringRef path = constant.getLocationAttr().getValue();
-  const int64_t offset = constant.getOffsetAttr().getInt();
-  const int64_t size = constant.getSizeAttr().getInt();
-  const int64_t elementBytes = (type.getElementTypeBitWidth() + 7) / 8;
-  if (size != type.getNumElements() * elementBytes)
-    return failure();
-  // The range is validated against the file before mapping: a slice running
-  // past the end can otherwise be mapped and fault on access rather than
-  // failing here.
-  uint64_t fileSize = 0;
-  if (llvm::sys::fs::file_size(path, fileSize) ||
-      static_cast<uint64_t>(offset) + static_cast<uint64_t>(size) > fileSize)
-    return failure();
-  llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>> buffer =
-      llvm::MemoryBuffer::getFileSlice(path, static_cast<uint64_t>(size),
-                                       offset);
-  if (!buffer || (*buffer)->getBufferSize() != static_cast<uint64_t>(size))
-    return failure();
-  DenseElementsAttr value = DenseElementsAttr::getFromRawBuffer(
-      type,
-      ArrayRef<char>((*buffer)->getBufferStart(), static_cast<size_t>(size)));
-  if (!value)
-    return failure();
-  return value;
-}
-
 LogicalResult
 DecomposeConvTranspose::matchAndRewrite(ConvTransposeOp op,
                                         PatternRewriter &rewriter) const {
@@ -204,12 +145,16 @@ DecomposeConvTranspose::matchAndRewrite(ConvTransposeOp op,
     return rewriter.notifyMatchFailure(op, "expected matching element types");
 
   // The tap selection and flip are folded into new constants, so the filter
-  // has to be readable at compile time.
+  // has to be readable at compile time. Requiring an inline value would
+  // decline nearly every real filter, which the importer externalizes to a
+  // file-backed range. Memory-address carriers are not resolved: this pattern
+  // is equally reachable from textual pass invocation, where the address
+  // would be dereferenced blind.
   auto weightConst = op.getWeights().getDefiningOp<ConstantOp>();
   if (!weightConst)
     return rewriter.notifyMatchFailure(op, "weights are not a hip.constant");
   FailureOr<DenseElementsAttr> weightData =
-      readConstant(weightConst, weightType);
+      readConstantData(weightConst, weightType);
   if (failed(weightData))
     return rewriter.notifyMatchFailure(op, "weights are not readable here");
   DenseElementsAttr weightAttr = *weightData;