diff --git a/docs/design/constant-handling-design.md b/docs/design/constant-handling-design.md index 7db27a8589..a9bcf163fd 100644 --- a/docs/design/constant-handling-design.md +++ b/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: diff --git a/docs/pipeline_pass_menu.md b/docs/pipeline_pass_menu.md index cc6b2ec307..0bd2dba8c7 100644 --- a/docs/pipeline_pass_menu.md +++ b/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 diff --git a/include/hip/Dialect/IR/HipConstantData.h b/include/hip/Dialect/IR/HipConstantData.h new file mode 100644 index 0000000000..4e030f5c7e --- /dev/null +++ b/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 readConstantData(ConstantOp constant, + RankedTensorType type); + +} // namespace hip +} // namespace mlir + +#endif // HIP_DIALECT_IR_HIP_CONSTANT_DATA_H diff --git a/include/hip/Dialect/Transforms/Passes.td b/include/hip/Dialect/Transforms/Passes.td index f3246793db..104814c5c4 100644 --- a/include/hip/Dialect/Transforms/Passes.td +++ b/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 " diff --git a/lib/Conversion/OnnxToHip/NhwcConvConversion.cpp b/lib/Conversion/OnnxToHip/NhwcConvConversion.cpp index b0bead8885..adf5a78043 100644 --- a/lib/Conversion/OnnxToHip/NhwcConvConversion.cpp +++ b/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 @@ -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 readWeightConstant(ConstantOp constant, - RankedTensorType type) { - switch (constant.getSourceKind()) { - case ConstantOp::SourceKind::Inline: { - auto value = dyn_cast(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(offset) + static_cast(size) > fileSize) - return failure(); - llvm::ErrorOr> buffer = - llvm::MemoryBuffer::getFileSlice(path, static_cast(size), - offset); - if (!buffer || (*buffer)->getBufferSize() != static_cast(size)) - return failure(); - DenseElementsAttr value = DenseElementsAttr::getFromRawBuffer( - type, - ArrayRef((*buffer)->getBufferStart(), static_cast(size))); - if (!value) - return failure(); - return value; -} - // [M, kH, kW, C/group] -> [M, C/group, kH, kW], preserving element bits. FailureOr transposeOhwiToOihw(DenseElementsAttr value, RankedTensorType srcType) { @@ -289,8 +247,11 @@ LogicalResult NhwcConvToHip::matchAndRewrite(Operation *op, ConstantOp deadWeight; if (weightType.hasStaticShape()) { if (auto weightConst = weights.getDefiningOp()) { + // 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 data = - readWeightConstant(weightConst, weightType); + readConstantData(weightConst, weightType); FailureOr transposed = succeeded(data) ? transposeOhwiToOihw(*data, weightType) : FailureOr(failure()); diff --git a/lib/Dialect/IR/CMakeLists.txt b/lib/Dialect/IR/CMakeLists.txt index e7e3e1dec9..667f2ddce8 100644 --- a/lib/Dialect/IR/CMakeLists.txt +++ b/lib/Dialect/IR/CMakeLists.txt @@ -5,6 +5,7 @@ add_library(HipDialectIR STATIC HipCausalConvCanonicalize.cpp + HipConstantData.cpp HipDialect.cpp HipDpsOpInterface.cpp HipOpStateInterface.cpp diff --git a/lib/Dialect/IR/HipConstantData.cpp b/lib/Dialect/IR/HipConstantData.cpp new file mode 100644 index 0000000000..221c567fe8 --- /dev/null +++ b/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 +#include + +using namespace mlir; +using namespace mlir::hip; + +FailureOr +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(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(offset) + static_cast(size) > fileSize) + return failure(); + llvm::ErrorOr> buffer = + llvm::MemoryBuffer::getFileSlice(path, static_cast(size), + offset); + if (!buffer || (*buffer)->getBufferSize() != static_cast(size)) + return failure(); + ArrayRef raw((*buffer)->getBufferStart(), static_cast(size)); + if (!DenseElementsAttr::isValidRawBuffer(type, raw)) + return failure(); + DenseElementsAttr value = DenseElementsAttr::getFromRawBuffer(type, raw); + if (!value) + return failure(); + return value; +} diff --git a/lib/Dialect/Transforms/CMakeLists.txt b/lib/Dialect/Transforms/CMakeLists.txt index fe7e4edc74..82b5a51387 100644 --- a/lib/Dialect/Transforms/CMakeLists.txt +++ b/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 diff --git a/lib/Dialect/Transforms/DecomposeConvTranspose.cpp b/lib/Dialect/Transforms/DecomposeConvTranspose.cpp index e1827dbc45..dcd69ff267 100644 --- a/lib/Dialect/Transforms/DecomposeConvTranspose.cpp +++ b/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 #include -#include -#include - #include -#include #include namespace mlir::hip { @@ -113,62 +110,6 @@ int64_t flatIndex(ArrayRef 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 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(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(offset) + static_cast(size) > fileSize) - return failure(); - llvm::ErrorOr> buffer = - llvm::MemoryBuffer::getFileSlice(path, static_cast(size), - offset); - if (!buffer || (*buffer)->getBufferSize() != static_cast(size)) - return failure(); - DenseElementsAttr value = DenseElementsAttr::getFromRawBuffer( - type, - ArrayRef((*buffer)->getBufferStart(), static_cast(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(); if (!weightConst) return rewriter.notifyMatchFailure(op, "weights are not a hip.constant"); FailureOr weightData = - readConstant(weightConst, weightType); + readConstantData(weightConst, weightType); if (failed(weightData)) return rewriter.notifyMatchFailure(op, "weights are not readable here"); DenseElementsAttr weightAttr = *weightData; diff --git a/lib/Dialect/Transforms/FoldBatchNorm.cpp b/lib/Dialect/Transforms/FoldBatchNorm.cpp new file mode 100644 index 0000000000..c9fabbd2b8 --- /dev/null +++ b/lib/Dialect/Transforms/FoldBatchNorm.cpp @@ -0,0 +1,254 @@ +/* + * Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. + * Licensed under the MIT License. + */ + +#include "hip/Dialect/IR/HipConstantData.h" +#include "hip/Dialect/IR/HipDialect.h" +#include "hip/Dialect/Transforms/Passes.h" + +#include +#include +#include +#include +#include + +#include + +#include +#include + +namespace mlir::hip { +#define GEN_PASS_DEF_FOLDBATCHNORMPASS +#include "hip/Dialect/Transforms/Passes.h.inc" + +namespace { + +/// Inference BatchNorm with constant parameters is the per-channel affine map +/// y = s * x + t, with s = scale / sqrt(var + eps) and t = B - mean * s. A +/// convolution is linear in its weights per output channel, so s moves into +/// the filter and t into the bias. Because the map distributes over a sum, the +/// same holds through an add of two convolutions: both filters are scaled and +/// t is added once, to one of the two biases. +/// +/// 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} : tensor<128x128x3x3xf32> +/// %t = hip.constant {value = dense} : tensor<128xf32> +/// %a = hip.conv(%ctx) ins(%x, %wa2, %t : ...) ... +/// %wb2 = hip.constant {value = dense} : tensor<128x384x3x3xf32> +/// %b = hip.conv(%ctx) ins(%y, %wb2 : ...) ... +/// %s = hip.add(%ctx) ins(%a, %b : ...) ... +/// %r = hip.max(%ctx) ins(%s, %zero : ...) ... +/// +/// Every intermediate result must have the batch norm as its only consumer, +/// since the rescaled convolutions no longer produce the original values. +struct FoldBatchNormIntoConv : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(BatchNormOp op, + PatternRewriter &rewriter) const override; +}; + +bool isFoldableFloat(Type type) { + return type.isF16() || type.isBF16() || type.isF32() || type.isF64(); +} + +/// Per-channel values of a 1-D constant of length `channels`. +std::optional> readVector(Value value, int64_t channels) { + auto constant = value.getDefiningOp(); + auto type = dyn_cast(value.getType()); + if (!constant || !type || !type.hasStaticShape() || type.getRank() != 1 || + type.getDimSize(0) != channels || !isFoldableFloat(type.getElementType())) + return std::nullopt; + FailureOr data = readConstantData(constant, type); + if (failed(data)) + return std::nullopt; + SmallVector values; + values.reserve(channels); + for (const APFloat &v : data->getValues()) + values.push_back(v.convertToDouble()); + return values; +} + +/// `type` filled with `values`, rounded to its element type. +DenseElementsAttr buildFloatAttr(RankedTensorType type, + ArrayRef values) { + Type elementType = type.getElementType(); + if (elementType.isF32()) { + SmallVector narrowed(values.begin(), values.end()); + return DenseElementsAttr::get(type, ArrayRef(narrowed)); + } + if (elementType.isF64()) + return DenseElementsAttr::get(type, values); + const llvm::fltSemantics &semantics = + cast(elementType).getFloatSemantics(); + SmallVector rounded; + rounded.reserve(values.size()); + for (double v : values) { + APFloat f(v); + bool losesInfo = false; + f.convert(semantics, APFloat::rmNearestTiesToEven, &losesInfo); + rounded.push_back(f); + } + return DenseElementsAttr::get(type, rounded); +} + +/// A convolution that the batch norm can be folded into: its only consumer is +/// on the path to the batch norm, and its filter and bias (if any) are +/// readable here. +struct FoldableConv { + ConvOp op; + DenseElementsAttr weights; + std::optional> bias; +}; + +LogicalResult +FoldBatchNormIntoConv::matchAndRewrite(BatchNormOp op, + PatternRewriter &rewriter) const { + if (op->getNumResults() != 1) + return rewriter.notifyMatchFailure(op, "expected tensor mode"); + auto resultType = dyn_cast(op->getResult(0).getType()); + if (!resultType || !resultType.hasStaticShape() || resultType.getRank() < 3 || + op.getInput().getType() != resultType || + !isFoldableFloat(resultType.getElementType())) + return rewriter.notifyMatchFailure(op, "expected static float NC... input"); + const int64_t channels = resultType.getDimSize(1); + + // Collect the convolutions feeding the batch norm, directly or through one + // add whose operands are both convolutions of the batch norm's shape. + SmallVector convs; + AddOp add; + Value input = op.getInput(); + if (!input.hasOneUse()) + return rewriter.notifyMatchFailure(op, "input has other consumers"); + if (auto conv = input.getDefiningOp()) { + convs.push_back(conv); + } else if ((add = input.getDefiningOp())) { + if (add->getNumResults() != 1) + return rewriter.notifyMatchFailure(op, "expected tensor-mode add"); + for (Value operand : {add.getLhs(), add.getRhs()}) { + auto conv = operand.getDefiningOp(); + if (!conv || operand.getType() != resultType || !operand.hasOneUse()) + return rewriter.notifyMatchFailure( + op, "add operand is not a single-use conv of the same shape"); + convs.push_back(conv); + } + } else { + return rewriter.notifyMatchFailure(op, "input is not a conv or an add"); + } + + SmallVector foldable; + for (ConvOp conv : convs) { + if (conv->getNumResults() != 1 || + conv->getResult(0).getType() != resultType) + return rewriter.notifyMatchFailure(op, "conv result does not match"); + auto weightType = dyn_cast(conv.getWeights().getType()); + auto weightConst = conv.getWeights().getDefiningOp(); + if (!weightType || !weightConst || !weightType.hasStaticShape() || + weightType.getDimSize(0) != channels || + weightType.getElementType() != resultType.getElementType()) + return rewriter.notifyMatchFailure(op, "conv filter is not a constant"); + FailureOr weights = + readConstantData(weightConst, weightType); + if (failed(weights)) + return rewriter.notifyMatchFailure(op, "conv filter is not readable"); + FoldableConv entry{conv, *weights, std::nullopt}; + if (Value bias = conv.getBias()) { + entry.bias = readVector(bias, channels); + if (!entry.bias) + return rewriter.notifyMatchFailure(op, "conv bias is not readable"); + } + foldable.push_back(std::move(entry)); + } + + std::optional> gamma = + readVector(op.getScale(), channels); + std::optional> beta = readVector(op.getBias(), channels); + std::optional> mean = readVector(op.getMean(), channels); + std::optional> var = + readVector(op.getVariance(), channels); + if (!gamma || !beta || !mean || !var) + return rewriter.notifyMatchFailure(op, + "batch norm parameters not readable"); + + const double epsilon = op.getEpsilon().convertToDouble(); + SmallVector scale(channels), shift(channels); + for (int64_t c : llvm::seq(channels)) { + scale[c] = (*gamma)[c] / std::sqrt((*var)[c] + epsilon); + shift[c] = (*beta)[c] - (*mean)[c] * scale[c]; + } + + auto biasType = + RankedTensorType::get({channels}, resultType.getElementType()); + SmallVector replacements; + for (auto [index, entry] : llvm::enumerate(foldable)) { + ConvOp conv = entry.op; + auto weightType = cast(conv.getWeights().getType()); + const int64_t perChannel = weightType.getNumElements() / channels; + + SmallVector weights; + weights.reserve(weightType.getNumElements()); + int64_t flat = 0; + for (const APFloat &w : entry.weights.getValues()) + weights.push_back(w.convertToDouble() * scale[flat++ / perChannel]); + + // The shift is added once: through an add, the second conv's output is + // summed with the first's, so a shift on both would count it twice. + std::optional> bias; + if (entry.bias || index == 0) { + bias.emplace(channels, 0.0); + for (int64_t c : llvm::seq(channels)) + (*bias)[c] = (entry.bias ? (*entry.bias)[c] * scale[c] : 0.0) + + (index == 0 ? shift[c] : 0.0); + } + + rewriter.setInsertionPoint(conv); + Location loc = conv.getLoc(); + Value newWeights = ConstantOp::create(rewriter, loc, weightType, + buildFloatAttr(weightType, weights)) + .getResult(); + SmallVector operands{conv.getCtx(), conv.getInput(), newWeights}; + if (bias) + operands.push_back(ConstantOp::create(rewriter, loc, biasType, + buildFloatAttr(biasType, *bias)) + .getResult()); + operands.push_back(conv.getOutput()); + auto folded = ConvOp::create(rewriter, loc, TypeRange{resultType}, operands, + conv->getAttrs()); + replacements.push_back(folded->getResult(0)); + } + + Value result; + if (add) { + rewriter.modifyOpInPlace(add, [&] { + add.getLhsMutable().assign(replacements[0]); + add.getRhsMutable().assign(replacements[1]); + }); + result = add->getResult(0); + } else { + result = replacements[0]; + } + rewriter.replaceOp(op, result); + for (ConvOp conv : convs) + rewriter.eraseOp(conv); + return success(); +} + +struct FoldBatchNormPass + : public impl::FoldBatchNormPassBase { + void runOnOperation() override { + RewritePatternSet patterns(&getContext()); + patterns.add(&getContext()); + if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) + signalPassFailure(); + } +}; + +} // namespace +} // namespace mlir::hip diff --git a/lib/Dialect/Transforms/Pipelines.cpp b/lib/Dialect/Transforms/Pipelines.cpp index afa7e09e3a..fe1c50cde4 100644 --- a/lib/Dialect/Transforms/Pipelines.cpp +++ b/lib/Dialect/Transforms/Pipelines.cpp @@ -416,6 +416,12 @@ void mlir::hip::buildOnnxToHipPipelineHead(OpPassManager &pm, pm.addPass(createConvertOnnxToHipPass()); + // Fold constant-parameter BatchNorm into the producing convolutions while + // the conv -> add -> batch_norm chain is still as converted: this head is + // shared with the rocMLIR flow, so it runs before any conv rewrite or kernel + // outlining there. + pm.addNestedPass(createFoldBatchNormPass()); + // Plugin slot: AfterConvertOnnxToHip. The most common slot for // vendor lowerings of `onnx.Custom` ops or vendor-specific hip.* // canonicalisations. This is the supported plugin boundary for emitting diff --git a/test/lit/Transforms/test_fold_batchnorm.mlir b/test/lit/Transforms/test_fold_batchnorm.mlir new file mode 100644 index 0000000000..15ea90af57 --- /dev/null +++ b/test/lit/Transforms/test_fold_batchnorm.mlir @@ -0,0 +1,271 @@ +// Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +// Licensed under the MIT License. + +// ============================================================================ +// TEST PURPOSE: +// Verify an inference hip.batch_norm with constant parameters is folded into +// the filter and bias of the convolution(s) producing its input, and that the +// cases where the fold would be wrong or unreadable are left alone. +// +// Every positive case uses scale = [1, 2], B = [0.5, 1], mean = [1, 0], +// var = [3, 0], epsilon = 1, so s = scale / sqrt(var + eps) = [0.5, 2] and +// t = B - mean * s = [0, 1]. +// +// Test cases: +// 1. conv_no_bias - filter scaled per output channel, t becomes the bias +// 2. conv_with_bias - existing bias b becomes s * b + t +// 3. add_conv_conv - both filters scaled, t added to the first conv only +// 4. grouped - grouped conv folds the same way (per output channel) +// 5. f16 - folded values are rounded to the filter type +// 6. multi_use - conv result has a second consumer, BN survives +// 7. dynamic_param - non-constant BN parameter, BN survives +// 8. memory_filter - memory-address filter is never read, BN survives +// ============================================================================ + +// RUN: hip-mlir-opt %s --hip-fold-batchnorm --canonicalize | FileCheck %s + +module { + // -------------------------------------------------------------------------- + // 1. BN(conv), no conv bias. + // -------------------------------------------------------------------------- + func.func @conv_no_bias(%ctx: !hip.context, %x: tensor<1x1x4x4xf32>) + -> tensor<1x2x4x4xf32> { + %w = hip.constant {value = dense<[[[[2.0]]], [[[3.0]]]]> : tensor<2x1x1x1xf32>} : tensor<2x1x1x1xf32> + %scale = hip.constant {value = dense<[1.0, 2.0]> : tensor<2xf32>} : tensor<2xf32> + %b = hip.constant {value = dense<[0.5, 1.0]> : tensor<2xf32>} : tensor<2xf32> + %mean = hip.constant {value = dense<[1.0, 0.0]> : tensor<2xf32>} : tensor<2xf32> + %var = hip.constant {value = dense<[3.0, 0.0]> : tensor<2xf32>} : tensor<2xf32> + %i0 = tensor.empty() : tensor<1x2x4x4xf32> + %c = hip.conv(%ctx) ins(%x, %w : tensor<1x1x4x4xf32>, tensor<2x1x1x1xf32>) + outs(%i0 : tensor<1x2x4x4xf32>) + {kernel_shape = [1, 1], strides = [1, 1], pads = [0, 0, 0, 0], + dilations = [1, 1], group = 1 : i64} : tensor<1x2x4x4xf32> + %i1 = tensor.empty() : tensor<1x2x4x4xf32> + %y = hip.batch_norm(%ctx) ins(%c, %scale, %b, %mean, %var : + tensor<1x2x4x4xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, + tensor<2xf32>) outs(%i1 : tensor<1x2x4x4xf32>) + {epsilon = 1.0 : f32} : tensor<1x2x4x4xf32> + return %y : tensor<1x2x4x4xf32> + } + + // CHECK-LABEL: func.func @conv_no_bias + // CHECK-NOT: hip.batch_norm + // CHECK-DAG: %[[W:.*]] = hip.constant {value = dense<{{\[}}{{\[}}{{\[}}[1.000000e+00]]], {{\[}}{{\[}}[6.000000e+00]]]]> : tensor<2x1x1x1xf32>} + // CHECK-DAG: %[[B:.*]] = hip.constant {value = dense<[0.000000e+00, 1.000000e+00]> : tensor<2xf32>} + // CHECK: %[[Y:.*]] = hip.conv(%{{.*}}) ins(%{{.*}}, %[[W]], %[[B]] : + // CHECK-NOT: hip.batch_norm + // CHECK: return %[[Y]] + + // -------------------------------------------------------------------------- + // 2. BN(conv) with an existing bias b = [1, 1] -> s * b + t = [0.5, 3]. + // -------------------------------------------------------------------------- + func.func @conv_with_bias(%ctx: !hip.context, %x: tensor<1x1x4x4xf32>) + -> tensor<1x2x4x4xf32> { + %w = hip.constant {value = dense<[[[[2.0]]], [[[3.0]]]]> : tensor<2x1x1x1xf32>} : tensor<2x1x1x1xf32> + %cb = hip.constant {value = dense<[1.0, 1.0]> : tensor<2xf32>} : tensor<2xf32> + %scale = hip.constant {value = dense<[1.0, 2.0]> : tensor<2xf32>} : tensor<2xf32> + %b = hip.constant {value = dense<[0.5, 1.0]> : tensor<2xf32>} : tensor<2xf32> + %mean = hip.constant {value = dense<[1.0, 0.0]> : tensor<2xf32>} : tensor<2xf32> + %var = hip.constant {value = dense<[3.0, 0.0]> : tensor<2xf32>} : tensor<2xf32> + %i0 = tensor.empty() : tensor<1x2x4x4xf32> + %c = hip.conv(%ctx) ins(%x, %w, %cb : tensor<1x1x4x4xf32>, tensor<2x1x1x1xf32>, tensor<2xf32>) + outs(%i0 : tensor<1x2x4x4xf32>) + {kernel_shape = [1, 1], strides = [1, 1], pads = [0, 0, 0, 0], + dilations = [1, 1], group = 1 : i64} : tensor<1x2x4x4xf32> + %i1 = tensor.empty() : tensor<1x2x4x4xf32> + %y = hip.batch_norm(%ctx) ins(%c, %scale, %b, %mean, %var : + tensor<1x2x4x4xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, + tensor<2xf32>) outs(%i1 : tensor<1x2x4x4xf32>) + {epsilon = 1.0 : f32} : tensor<1x2x4x4xf32> + return %y : tensor<1x2x4x4xf32> + } + + // CHECK-LABEL: func.func @conv_with_bias + // CHECK-NOT: hip.batch_norm + // CHECK-DAG: %[[B:.*]] = hip.constant {value = dense<[5.000000e-01, 3.000000e+00]> : tensor<2xf32>} + // CHECK: hip.conv(%{{.*}}) ins(%{{.*}}, %{{.*}}, %[[B]] : + // CHECK-NOT: hip.batch_norm + + // -------------------------------------------------------------------------- + // 3. BN(add(conv, conv)): both filters scaled by s; t lands on one bias only. + // -------------------------------------------------------------------------- + func.func @add_conv_conv(%ctx: !hip.context, %x: tensor<1x1x4x4xf32>, + %z: tensor<1x1x4x4xf32>) -> tensor<1x2x4x4xf32> { + %wa = hip.constant {value = dense<[[[[2.0]]], [[[3.0]]]]> : tensor<2x1x1x1xf32>} : tensor<2x1x1x1xf32> + %wb = hip.constant {value = dense<[[[[4.0]]], [[[5.0]]]]> : tensor<2x1x1x1xf32>} : tensor<2x1x1x1xf32> + %scale = hip.constant {value = dense<[1.0, 2.0]> : tensor<2xf32>} : tensor<2xf32> + %b = hip.constant {value = dense<[0.5, 1.0]> : tensor<2xf32>} : tensor<2xf32> + %mean = hip.constant {value = dense<[1.0, 0.0]> : tensor<2xf32>} : tensor<2xf32> + %var = hip.constant {value = dense<[3.0, 0.0]> : tensor<2xf32>} : tensor<2xf32> + %i0 = tensor.empty() : tensor<1x2x4x4xf32> + %ca = hip.conv(%ctx) ins(%x, %wa : tensor<1x1x4x4xf32>, tensor<2x1x1x1xf32>) + outs(%i0 : tensor<1x2x4x4xf32>) + {kernel_shape = [1, 1], strides = [1, 1], pads = [0, 0, 0, 0], + dilations = [1, 1], group = 1 : i64} : tensor<1x2x4x4xf32> + %i1 = tensor.empty() : tensor<1x2x4x4xf32> + %cb = hip.conv(%ctx) ins(%z, %wb : tensor<1x1x4x4xf32>, tensor<2x1x1x1xf32>) + outs(%i1 : tensor<1x2x4x4xf32>) + {kernel_shape = [1, 1], strides = [1, 1], pads = [0, 0, 0, 0], + dilations = [1, 1], group = 1 : i64} : tensor<1x2x4x4xf32> + %i2 = tensor.empty() : tensor<1x2x4x4xf32> + %s = hip.add(%ctx) ins(%ca, %cb : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) + outs(%i2 : tensor<1x2x4x4xf32>) -> tensor<1x2x4x4xf32> + %i3 = tensor.empty() : tensor<1x2x4x4xf32> + %y = hip.batch_norm(%ctx) ins(%s, %scale, %b, %mean, %var : + tensor<1x2x4x4xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, + tensor<2xf32>) outs(%i3 : tensor<1x2x4x4xf32>) + {epsilon = 1.0 : f32} : tensor<1x2x4x4xf32> + return %y : tensor<1x2x4x4xf32> + } + + // CHECK-LABEL: func.func @add_conv_conv + // CHECK-SAME: (%{{.*}}: !hip.context, %[[X:.*]]: tensor<1x1x4x4xf32>, %[[Z:.*]]: tensor<1x1x4x4xf32>) + // CHECK-NOT: hip.batch_norm + // CHECK-DAG: %[[WA:.*]] = hip.constant {value = dense<{{\[}}{{\[}}{{\[}}[1.000000e+00]]], {{\[}}{{\[}}[6.000000e+00]]]]> : tensor<2x1x1x1xf32>} + // CHECK-DAG: %[[WB:.*]] = hip.constant {value = dense<{{\[}}{{\[}}{{\[}}[2.000000e+00]]], {{\[}}{{\[}}[1.000000e+01]]]]> : tensor<2x1x1x1xf32>} + // CHECK-DAG: %[[T:.*]] = hip.constant {value = dense<[0.000000e+00, 1.000000e+00]> : tensor<2xf32>} + // CHECK-DAG: %[[CA:.*]] = hip.conv(%{{.*}}) ins(%[[X]], %[[WA]], %[[T]] : + // CHECK-DAG: %[[CB:.*]] = hip.conv(%{{.*}}) ins(%[[Z]], %[[WB]] : tensor<1x1x4x4xf32>, tensor<2x1x1x1xf32>) + // CHECK: %[[S:.*]] = hip.add(%{{.*}}) ins(%[[CA]], %[[CB]] : + // CHECK-NOT: hip.batch_norm + // CHECK: return %[[S]] + + // -------------------------------------------------------------------------- + // 4. Grouped conv: the filter is still [M, C/group, ...], scaled per M. + // -------------------------------------------------------------------------- + func.func @grouped(%ctx: !hip.context, %x: tensor<1x2x4x4xf32>) + -> tensor<1x2x4x4xf32> { + %w = hip.constant {value = dense<[[[[2.0]]], [[[3.0]]]]> : tensor<2x1x1x1xf32>} : tensor<2x1x1x1xf32> + %scale = hip.constant {value = dense<[1.0, 2.0]> : tensor<2xf32>} : tensor<2xf32> + %b = hip.constant {value = dense<[0.5, 1.0]> : tensor<2xf32>} : tensor<2xf32> + %mean = hip.constant {value = dense<[1.0, 0.0]> : tensor<2xf32>} : tensor<2xf32> + %var = hip.constant {value = dense<[3.0, 0.0]> : tensor<2xf32>} : tensor<2xf32> + %i0 = tensor.empty() : tensor<1x2x4x4xf32> + %c = hip.conv(%ctx) ins(%x, %w : tensor<1x2x4x4xf32>, tensor<2x1x1x1xf32>) + outs(%i0 : tensor<1x2x4x4xf32>) + {kernel_shape = [1, 1], strides = [1, 1], pads = [0, 0, 0, 0], + dilations = [1, 1], group = 2 : i64} : tensor<1x2x4x4xf32> + %i1 = tensor.empty() : tensor<1x2x4x4xf32> + %y = hip.batch_norm(%ctx) ins(%c, %scale, %b, %mean, %var : + tensor<1x2x4x4xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, + tensor<2xf32>) outs(%i1 : tensor<1x2x4x4xf32>) + {epsilon = 1.0 : f32} : tensor<1x2x4x4xf32> + return %y : tensor<1x2x4x4xf32> + } + + // CHECK-LABEL: func.func @grouped + // CHECK-NOT: hip.batch_norm + // CHECK-DAG: %[[W:.*]] = hip.constant {value = dense<{{\[}}{{\[}}{{\[}}[1.000000e+00]]], {{\[}}{{\[}}[6.000000e+00]]]]> : tensor<2x1x1x1xf32>} + // CHECK: hip.conv(%{{.*}}) ins(%{{.*}}, %[[W]], %{{.*}} :{{.*}}group = 2 : i64 + // CHECK-NOT: hip.batch_norm + + // -------------------------------------------------------------------------- + // 5. f16: the folded filter and bias are f16. + // -------------------------------------------------------------------------- + func.func @f16(%ctx: !hip.context, %x: tensor<1x1x4x4xf16>) + -> tensor<1x2x4x4xf16> { + %w = hip.constant {value = dense<[[[[2.0]]], [[[3.0]]]]> : tensor<2x1x1x1xf16>} : tensor<2x1x1x1xf16> + %scale = hip.constant {value = dense<[1.0, 2.0]> : tensor<2xf16>} : tensor<2xf16> + %b = hip.constant {value = dense<[0.5, 1.0]> : tensor<2xf16>} : tensor<2xf16> + %mean = hip.constant {value = dense<[1.0, 0.0]> : tensor<2xf16>} : tensor<2xf16> + %var = hip.constant {value = dense<[3.0, 0.0]> : tensor<2xf16>} : tensor<2xf16> + %i0 = tensor.empty() : tensor<1x2x4x4xf16> + %c = hip.conv(%ctx) ins(%x, %w : tensor<1x1x4x4xf16>, tensor<2x1x1x1xf16>) + outs(%i0 : tensor<1x2x4x4xf16>) + {kernel_shape = [1, 1], strides = [1, 1], pads = [0, 0, 0, 0], + dilations = [1, 1], group = 1 : i64} : tensor<1x2x4x4xf16> + %i1 = tensor.empty() : tensor<1x2x4x4xf16> + %y = hip.batch_norm(%ctx) ins(%c, %scale, %b, %mean, %var : + tensor<1x2x4x4xf16>, tensor<2xf16>, tensor<2xf16>, tensor<2xf16>, + tensor<2xf16>) outs(%i1 : tensor<1x2x4x4xf16>) + {epsilon = 1.0 : f32} : tensor<1x2x4x4xf16> + return %y : tensor<1x2x4x4xf16> + } + + // CHECK-LABEL: func.func @f16 + // CHECK-NOT: hip.batch_norm + // CHECK-DAG: hip.constant {value = dense<{{\[}}{{\[}}{{\[}}[1.000000e+00]]], {{\[}}{{\[}}[6.000000e+00]]]]> : tensor<2x1x1x1xf16>} + // CHECK-DAG: hip.constant {value = dense<[0.000000e+00, 1.000000e+00]> : tensor<2xf16>} + // CHECK: hip.conv + // CHECK-NOT: hip.batch_norm + + // -------------------------------------------------------------------------- + // 6. The conv result is also returned: rescaling it would change that + // consumer, so the batch norm stays. + // -------------------------------------------------------------------------- + func.func @multi_use(%ctx: !hip.context, %x: tensor<1x1x4x4xf32>) + -> (tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32>) { + %w = hip.constant {value = dense<[[[[2.0]]], [[[3.0]]]]> : tensor<2x1x1x1xf32>} : tensor<2x1x1x1xf32> + %scale = hip.constant {value = dense<[1.0, 2.0]> : tensor<2xf32>} : tensor<2xf32> + %b = hip.constant {value = dense<[0.5, 1.0]> : tensor<2xf32>} : tensor<2xf32> + %mean = hip.constant {value = dense<[1.0, 0.0]> : tensor<2xf32>} : tensor<2xf32> + %var = hip.constant {value = dense<[3.0, 0.0]> : tensor<2xf32>} : tensor<2xf32> + %i0 = tensor.empty() : tensor<1x2x4x4xf32> + %c = hip.conv(%ctx) ins(%x, %w : tensor<1x1x4x4xf32>, tensor<2x1x1x1xf32>) + outs(%i0 : tensor<1x2x4x4xf32>) + {kernel_shape = [1, 1], strides = [1, 1], pads = [0, 0, 0, 0], + dilations = [1, 1], group = 1 : i64} : tensor<1x2x4x4xf32> + %i1 = tensor.empty() : tensor<1x2x4x4xf32> + %y = hip.batch_norm(%ctx) ins(%c, %scale, %b, %mean, %var : + tensor<1x2x4x4xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, + tensor<2xf32>) outs(%i1 : tensor<1x2x4x4xf32>) + {epsilon = 1.0 : f32} : tensor<1x2x4x4xf32> + return %y, %c : tensor<1x2x4x4xf32>, tensor<1x2x4x4xf32> + } + + // CHECK-LABEL: func.func @multi_use + // CHECK: hip.batch_norm + + // -------------------------------------------------------------------------- + // 7. A runtime BN parameter cannot be folded at compile time. + // -------------------------------------------------------------------------- + func.func @dynamic_param(%ctx: !hip.context, %x: tensor<1x1x4x4xf32>, + %scale: tensor<2xf32>) -> tensor<1x2x4x4xf32> { + %w = hip.constant {value = dense<[[[[2.0]]], [[[3.0]]]]> : tensor<2x1x1x1xf32>} : tensor<2x1x1x1xf32> + %b = hip.constant {value = dense<[0.5, 1.0]> : tensor<2xf32>} : tensor<2xf32> + %mean = hip.constant {value = dense<[1.0, 0.0]> : tensor<2xf32>} : tensor<2xf32> + %var = hip.constant {value = dense<[3.0, 0.0]> : tensor<2xf32>} : tensor<2xf32> + %i0 = tensor.empty() : tensor<1x2x4x4xf32> + %c = hip.conv(%ctx) ins(%x, %w : tensor<1x1x4x4xf32>, tensor<2x1x1x1xf32>) + outs(%i0 : tensor<1x2x4x4xf32>) + {kernel_shape = [1, 1], strides = [1, 1], pads = [0, 0, 0, 0], + dilations = [1, 1], group = 1 : i64} : tensor<1x2x4x4xf32> + %i1 = tensor.empty() : tensor<1x2x4x4xf32> + %y = hip.batch_norm(%ctx) ins(%c, %scale, %b, %mean, %var : + tensor<1x2x4x4xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, + tensor<2xf32>) outs(%i1 : tensor<1x2x4x4xf32>) + {epsilon = 1.0 : f32} : tensor<1x2x4x4xf32> + return %y : tensor<1x2x4x4xf32> + } + + // CHECK-LABEL: func.func @dynamic_param + // CHECK: hip.batch_norm + + // -------------------------------------------------------------------------- + // 8. A memory-address filter is process-local; the pass must not + // dereference it. + // -------------------------------------------------------------------------- + func.func @memory_filter(%ctx: !hip.context, %x: tensor<1x1x4x4xf32>) + -> tensor<1x2x4x4xf32> { + %w = hip.constant {memory_address = 4096 : i64, size = 8 : i64} : tensor<2x1x1x1xf32> + %scale = hip.constant {value = dense<[1.0, 2.0]> : tensor<2xf32>} : tensor<2xf32> + %b = hip.constant {value = dense<[0.5, 1.0]> : tensor<2xf32>} : tensor<2xf32> + %mean = hip.constant {value = dense<[1.0, 0.0]> : tensor<2xf32>} : tensor<2xf32> + %var = hip.constant {value = dense<[3.0, 0.0]> : tensor<2xf32>} : tensor<2xf32> + %i0 = tensor.empty() : tensor<1x2x4x4xf32> + %c = hip.conv(%ctx) ins(%x, %w : tensor<1x1x4x4xf32>, tensor<2x1x1x1xf32>) + outs(%i0 : tensor<1x2x4x4xf32>) + {kernel_shape = [1, 1], strides = [1, 1], pads = [0, 0, 0, 0], + dilations = [1, 1], group = 1 : i64} : tensor<1x2x4x4xf32> + %i1 = tensor.empty() : tensor<1x2x4x4xf32> + %y = hip.batch_norm(%ctx) ins(%c, %scale, %b, %mean, %var : + tensor<1x2x4x4xf32>, tensor<2xf32>, tensor<2xf32>, tensor<2xf32>, + tensor<2xf32>) outs(%i1 : tensor<1x2x4x4xf32>) + {epsilon = 1.0 : f32} : tensor<1x2x4x4xf32> + return %y : tensor<1x2x4x4xf32> + } + + // CHECK-LABEL: func.func @memory_filter + // CHECK: hip.constant {memory_address = 4096 : i64 + // CHECK: hip.batch_norm +} diff --git a/test/numeric/tests/test_conv_batchnorm.py b/test/numeric/tests/test_conv_batchnorm.py new file mode 100644 index 0000000000..53602c53e6 --- /dev/null +++ b/test/numeric/tests/test_conv_batchnorm.py @@ -0,0 +1,143 @@ +# +# Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. +# Licensed under the MIT License. +# + +"""Tests for BatchNormalization folded into the producing Conv weights. + +hip-fold-batchnorm rewrites an inference BatchNormalization with constant +parameters into the filter and bias of the Conv(s) feeding it, directly or +through an Add of two Convs. A wrong per-channel scale, a shift added twice, +or a shift dropped all give outputs that stay correlated with the reference, +so these compare tightly against the ORT CPU reference. + +The fold reads only inline constants, and every filter and BN parameter vector +here (64 channels, 256 B in f32) is above the importer's default 127-byte +inline threshold. Run with +`--ep-option XLNX_model_clone_external_data_threshold=1073741824` (or the +`XLNX_ABI_2_0_CLONE_EXTERNAL_DATA_THRESHOLD` environment variable) to exercise +the fold; without it the same graphs check the unfolded path. +""" + +import numpy as np +import pytest +from onnx import helper, numpy_helper + +from framework.comparator import compare_outputs +from framework.onnx_utils import make_model_from_nodes, np_to_onnx_type + +CHANNELS = 64 +SPATIAL = 8 + + +def _bn_initializers(rng, dtype, prefix="bn"): + """BatchNorm scale, B, mean, var with a spread wide enough to matter.""" + return [ + numpy_helper.from_array( + rng.uniform(0.5, 2.0, [CHANNELS]).astype(dtype), f"{prefix}_scale" + ), + numpy_helper.from_array( + rng.uniform(-1.0, 1.0, [CHANNELS]).astype(dtype), f"{prefix}_b" + ), + numpy_helper.from_array( + rng.uniform(-0.5, 0.5, [CHANNELS]).astype(dtype), f"{prefix}_mean" + ), + numpy_helper.from_array( + rng.uniform(0.1, 2.0, [CHANNELS]).astype(dtype), f"{prefix}_var" + ), + ] + + +def _bn_node(x, y, prefix="bn"): + return helper.make_node( + "BatchNormalization", + [x, f"{prefix}_scale", f"{prefix}_b", f"{prefix}_mean", f"{prefix}_var"], + [y], + epsilon=1e-5, + ) + + +def _conv_node(x, w, y, bias=None, group=1): + return helper.make_node( + "Conv", + [x, w] + ([bias] if bias else []), + [y], + kernel_shape=[3, 3], + strides=[1, 1], + pads=[1, 1, 1, 1], + dilations=[1, 1], + group=group, + ) + + +def _weight(rng, dtype, name, in_c, group=1): + w = rng.uniform(-0.2, 0.2, [CHANNELS, in_c // group, 3, 3]).astype(dtype) + return numpy_helper.from_array(w, name) + + +_TOLERANCE = { + np.float32: dict(atol=1e-3, rtol=1e-3, cos_threshold=0.99999), + np.float16: dict(atol=3e-2, rtol=2e-2, cos_threshold=0.999), +} + + +class TestConvBatchNorm: + @pytest.mark.parametrize("dtype", [np.float32, np.float16]) + @pytest.mark.parametrize("with_bias", [False, True]) + @pytest.mark.parametrize("group", [1, 4]) + def test_conv_bn(self, model_runner, dtype, with_bias, group): + """Conv -> BatchNormalization folds into one conv with a bias.""" + in_c = 16 + tp = np_to_onnx_type(dtype) + rng = np.random.default_rng(11) + inits = [_weight(rng, dtype, "w", in_c, group)] + _bn_initializers(rng, dtype) + if with_bias: + inits.append( + numpy_helper.from_array( + rng.uniform(-0.5, 0.5, [CHANNELS]).astype(dtype), "cb" + ) + ) + nodes = [ + _conv_node("x", "w", "c", bias="cb" if with_bias else None, group=group), + _bn_node("c", "y"), + ] + model = make_model_from_nodes( + nodes, + [helper.make_tensor_value_info("x", tp, [1, in_c, SPATIAL, SPATIAL])], + [helper.make_tensor_value_info("y", tp, [1, CHANNELS, SPATIAL, SPATIAL])], + initializers=inits, + ) + x = rng.uniform(-1.0, 1.0, [1, in_c, SPATIAL, SPATIAL]).astype(dtype) + actual, expected = model_runner.run_sample(model, [x]) + compare_outputs(actual, expected, **_TOLERANCE[dtype]) + + @pytest.mark.parametrize("dtype", [np.float32, np.float16]) + def test_add_conv_conv_bn_relu(self, model_runner, dtype): + """Relu(BN(Conv(x) + Conv(z))): both filters scaled, shift added once.""" + in_x, in_z = 16, 48 + tp = np_to_onnx_type(dtype) + rng = np.random.default_rng(23) + inits = [ + _weight(rng, dtype, "wa", in_x), + _weight(rng, dtype, "wb", in_z), + ] + _bn_initializers(rng, dtype) + nodes = [ + _conv_node("x", "wa", "ca"), + _conv_node("z", "wb", "cb"), + helper.make_node("Add", ["ca", "cb"], ["s"]), + _bn_node("s", "n"), + helper.make_node("Relu", ["n"], ["y"]), + ] + model = make_model_from_nodes( + nodes, + [ + helper.make_tensor_value_info("x", tp, [1, in_x, SPATIAL, SPATIAL]), + helper.make_tensor_value_info("z", tp, [1, in_z, SPATIAL, SPATIAL]), + ], + [helper.make_tensor_value_info("y", tp, [1, CHANNELS, SPATIAL, SPATIAL])], + initializers=inits, + ) + x = rng.uniform(-1.0, 1.0, [1, in_x, SPATIAL, SPATIAL]).astype(dtype) + z = rng.uniform(-1.0, 1.0, [1, in_z, SPATIAL, SPATIAL]).astype(dtype) + actual, expected = model_runner.run_sample(model, [x, z]) + compare_outputs(actual, expected, **_TOLERANCE[dtype])