[Codegen][GlobalOpt] Detach non-splat constant inits (#24605)
A constant used directly as the `outs` (init/destination) of a
`linalg.matmul` or convolution e.g. a fully-connected/conv bias imported
from TFLite via TOSA survives into dispatch formation. It is then cloned
into the workgroup on the written (output) side and bufferizes as a
read-only global, which fails workgroup-distribution verification with
the following error:
```
'linalg.generic' op write affecting operations on global resources are restricted to workgroup distributed contexts
'func.func' op failed on workgroup distribution verification
```
This kind of IR was seen in TFLite -> TOSA conversion flows and and the
IR below is an example of that:
```mlir
util.func public @main(%a: tensor<1x1xf32>, %w: tensor<1x9xf32>) -> tensor<1x9xf32> {
%bias = arith.constant dense<[[1.0, 1.0, 2.0, 2.0, 3.0, 3.0, 4.0, 4.0, 5.0]]> : tensor<1x9xf32>
%0 = linalg.matmul ins(%a, %w : tensor<1x1xf32>, tensor<1x9xf32>)
outs(%bias : tensor<1x9xf32>) -> tensor<1x9xf32>
util.return %0 : tensor<1x9xf32>
}
```
The chosen solution is to use is the `DetachElementwiseFromNamedOps`
pattern which exists to keep non-writable inits out of dispatches. It
only covers two of our three relevant cases. The pass will
1. Take a non-const init and detach it into a zero-fill plus a trailing
add
2. Take a splat constant init and convert it to a fill.
Our use-case of a non-splat constant init would not be handled. Here we
extend the elementwise detach pattern so a non-splat constant init is
detached the same way a computed init is. A fresh zero-fill is created
and the constant is added back as a read-only input. This form of IR can
be then handled by codegen.
After the fix the reproducer compiles and the bias survives read-only:
```
%bias = arith.constant ... : tensor<1x9xf32>
%z = linalg.fill ins(%c0 : f32) outs(%empty : tensor<1x9xf32>) -> tensor<1x9xf32>
%mm = linalg.matmul ins(%a, %w) outs(%z) -> tensor<1x9xf32>
%0 = linalg.generic ins(%mm, %bias) outs(%z) { ... } // %mm + %bias
```
---
_Code Examples and Tests Assisted-by Claude_diff --git a/compiler/src/iree/compiler/GlobalOptimization/DetachElementwiseFromNamedOps.cpp b/compiler/src/iree/compiler/GlobalOptimization/DetachElementwiseFromNamedOps.cpp
index 24648f6..aa695f9 100644
--- a/compiler/src/iree/compiler/GlobalOptimization/DetachElementwiseFromNamedOps.cpp
+++ b/compiler/src/iree/compiler/GlobalOptimization/DetachElementwiseFromNamedOps.cpp
@@ -58,11 +58,27 @@
}
Value outputOperand = outputOperands.front()->get();
- auto outsDefiningOp = outputOperand.getDefiningOp<linalg::LinalgOp>();
- if (!outsDefiningOp || isa<linalg::FillOp>(outsDefiningOp.getOperation())) {
- // If not linalg op, or is a fill op, do nothing.
+ // By definition constants are read-only however some conversion paths
+ // place the bias initial value as the init/output argument of DPS
+ // operations. This later gets converted to a writable buffer. Here we
+ // detach the constant from the init argument to avoid this.
+ // Splat constant are handled later by DetachSplatConstantOutsOperands.
+ Operation *outsDefiningOp = outputOperand.getDefiningOp();
+ bool isNonSplatConstant = false;
+ if (auto constOp = dyn_cast_or_null<arith::ConstantOp>(outsDefiningOp)) {
+ auto elementsAttr = dyn_cast<ElementsAttr>(constOp.getValue());
+ isNonSplatConstant = elementsAttr && !elementsAttr.isSplat();
+ }
+
+ // outsDefiningOp is null for a block-argument init. The outsLinalgOp guard
+ // keeps the isa<FillOp> check off a null Operation*
+ auto outsLinalgOp = dyn_cast_or_null<linalg::LinalgOp>(outsDefiningOp);
+ bool needsDetach = isNonSplatConstant ||
+ (outsLinalgOp && !isa<linalg::FillOp>(outsDefiningOp));
+ if (!needsDetach) {
return failure();
}
+
auto outputType = cast<RankedTensorType>(outputOperand.getType());
if (!outputType.getElementType().isIntOrFloat()) {
return failure();
@@ -124,6 +140,9 @@
/// with allocations. Using `fill` will allow for fusing with the op just like
/// fill -> linalg ops are fused. If not as a fallback they would be converted
/// to a splat, but both without stack allocations.
+///
+/// Non-splat constant `outs` inits on contractions/convolutions are detached
+/// separately by DetachElementwisePattern (zero fill + trailing add).
template <typename InterfaceOp>
struct DetachSplatConstantOutsOperands
: OpInterfaceRewritePattern<InterfaceOp> {
diff --git a/compiler/src/iree/compiler/GlobalOptimization/test/detach_elementwise_from_named_ops.mlir b/compiler/src/iree/compiler/GlobalOptimization/test/detach_elementwise_from_named_ops.mlir
index 870e351..7075e0e 100644
--- a/compiler/src/iree/compiler/GlobalOptimization/test/detach_elementwise_from_named_ops.mlir
+++ b/compiler/src/iree/compiler/GlobalOptimization/test/detach_elementwise_from_named_ops.mlir
@@ -201,3 +201,55 @@
// CHECK: %[[GENERIC:.+]] = linalg.generic
// CHECK-SAME: outs(%[[FILL]] :
// CHECK: util.return %[[GENERIC]]
+
+// -----
+
+/// A non-splat constant init on a contraction must be detached: a constant
+/// cannot be a writable dispatch output. It is replaced with a zero fill and
+/// the constant is added back via a trailing elementwise op (as a read-only
+/// input).
+util.func public @matmul_nonsplat_cst_output(%arg0: tensor<2x3xf32>, %arg1: tensor<3x2xf32>) -> tensor<2x2xf32> {
+ %cst = arith.constant dense<[[1.0, 2.0], [3.0, 4.0]]> : tensor<2x2xf32>
+ %0 = linalg.matmul ins(%arg0, %arg1 : tensor<2x3xf32>, tensor<3x2xf32>)
+ outs(%cst : tensor<2x2xf32>) -> tensor<2x2xf32>
+ util.return %0 : tensor<2x2xf32>
+}
+// CHECK-LABEL: util.func public @matmul_nonsplat_cst_output
+// CHECK-SAME: %[[ARG0:.+]]: tensor<2x3xf32>, %[[ARG1:.+]]: tensor<3x2xf32>
+// CHECK-DAG: %[[CST:.+]] = arith.constant dense<{{.*}}> : tensor<2x2xf32>
+// CHECK-DAG: %[[F0:.+]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<2x2xf32>
+// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[F0]] : f32) outs(%[[INIT]] : tensor<2x2xf32>)
+// CHECK: %[[MM:.+]] = linalg.matmul
+// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] : tensor<2x3xf32>, tensor<3x2xf32>)
+// CHECK-SAME: outs(%[[FILL]] : tensor<2x2xf32>)
+// CHECK: %[[EW:.+]] = linalg.generic
+// CHECK-SAME: ins(%[[MM]], %[[CST]] : tensor<2x2xf32>, tensor<2x2xf32>)
+// CHECK-SAME: outs(%[[FILL]] : tensor<2x2xf32>)
+// CHECK: arith.addf
+// CHECK: util.return %[[EW]]
+
+// -----
+
+/// A non-splat constant init on a convolution is detached the same way as a
+/// contraction: zero fill + trailing add, with the constant as a read-only
+/// input.
+util.func public @conv_nonsplat_cst_output(%input: tensor<1x4x4x2xf32>, %filter: tensor<3x3x2x2xf32>) -> tensor<1x2x2x2xf32> {
+ %cst = arith.constant dense<[[[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]]> : tensor<1x2x2x2xf32>
+ %0 = linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+ ins(%input, %filter : tensor<1x4x4x2xf32>, tensor<3x3x2x2xf32>)
+ outs(%cst : tensor<1x2x2x2xf32>) -> tensor<1x2x2x2xf32>
+ util.return %0 : tensor<1x2x2x2xf32>
+}
+// CHECK-LABEL: util.func public @conv_nonsplat_cst_output
+// CHECK-DAG: %[[CST:.+]] = arith.constant dense<{{.*}}> : tensor<1x2x2x2xf32>
+// CHECK-DAG: %[[F0:.+]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1x2x2x2xf32>
+// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[F0]] : f32) outs(%[[INIT]] : tensor<1x2x2x2xf32>)
+// CHECK: %[[CONV:.+]] = linalg.conv_2d_nhwc_hwcf
+// CHECK-SAME: outs(%[[FILL]] : tensor<1x2x2x2xf32>)
+// CHECK: %[[EW:.+]] = linalg.generic
+// CHECK-SAME: ins(%[[CONV]], %[[CST]] : tensor<1x2x2x2xf32>, tensor<1x2x2x2xf32>)
+// CHECK-SAME: outs(%[[FILL]] : tensor<1x2x2x2xf32>)
+// CHECK: arith.addf
+// CHECK: util.return %[[EW]]
