[StableHLO] Legalize composite ops to calls in the input pipeline (#24777)
The StableHLO input conversion pipeline had no handling for
`stablehlo.composite`, so any composite reaching the conversion to
Linalg failed with:
```
error: failed to legalize operation 'stablehlo.composite' that was explicitly marked illegal
```
JAX emits composites for some CHLO operations, e.g.
```
%0 = stablehlo.composite "chlo.acos" %arg0
{decomposition = @chlo.acos.impl, version = 1 : i32}
: (tensor<8x32xf32>) -> tensor<8x32xf32>
```
which made such models uncompilable.
Run `stablehlo-legalize-composite-to-call` early in the pipeline, right
after deserialization. It replaces each composite with a call to the
function holding its decomposition, which the inliner later in the
pipeline folds away, so the remaining conversions only ever see the
decomposed form. The pass is given an explicit empty exception list to
document where composites that should be preserved as a unit can be
opted out.
Signed-off-by: Paul Stark <paul.stark@cdprojektred.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>diff --git a/compiler/plugins/input/StableHLO/Conversion/Passes.cpp b/compiler/plugins/input/StableHLO/Conversion/Passes.cpp
index 7887e2d..3a0dd8e 100644
--- a/compiler/plugins/input/StableHLO/Conversion/Passes.cpp
+++ b/compiler/plugins/input/StableHLO/Conversion/Passes.cpp
@@ -46,6 +46,18 @@
// If the input is StableHLO, this pass is considered a NOP.
passManager.addPass(stablehlo::createCheckVHLOStableHloMixUsage());
::mlir::stablehlo::createStablehloDeserializePipeline(passManager);
+ // Composite ops carry a symbol reference to a function implementing their
+ // decomposition. Replace them with a call to that function so that the
+ // conversions below only have to handle the decomposed form; the inliner
+ // later in this pipeline folds the call away.
+ //
+ // A composite that should be handled as a unit instead of being decomposed
+ // (for example one that maps onto a dedicated lowering) can be excluded by
+ // adding its name to the exception list below, e.g. {"my_namespace.my_op"}.
+ passManager.addNestedPass<func::FuncOp>(
+ ::mlir::stablehlo::createStablehloLegalizeCompositeToCallPass(
+ ::mlir::stablehlo::StablehloLegalizeCompositeToCallPassOptions{
+ /*exceptListOption=*/{}}));
passManager.addNestedPass<func::FuncOp>(mlir::createCanonicalizerPass());
passManager.addNestedPass<func::FuncOp>(createStableHLOCanonicalize());
passManager.addNestedPass<func::FuncOp>(mlir::createCSEPass());
diff --git a/compiler/plugins/input/StableHLO/Conversion/test/BUILD.bazel b/compiler/plugins/input/StableHLO/Conversion/test/BUILD.bazel
index ca0da7f..21e06dd 100644
--- a/compiler/plugins/input/StableHLO/Conversion/test/BUILD.bazel
+++ b/compiler/plugins/input/StableHLO/Conversion/test/BUILD.bazel
@@ -22,6 +22,7 @@
"legalize_chlo_decomposition.mlir",
"legalize_chlo_no_broadcast.mlir",
"legalize_chlo_with_broadcast.mlir",
+ "legalize_composite.mlir",
"legalize_control_flow.mlir",
"legalize_shape_computations.mlir",
"stablehlo_custom_calls.mlir",
diff --git a/compiler/plugins/input/StableHLO/Conversion/test/CMakeLists.txt b/compiler/plugins/input/StableHLO/Conversion/test/CMakeLists.txt
index 5f72021..43755f7 100644
--- a/compiler/plugins/input/StableHLO/Conversion/test/CMakeLists.txt
+++ b/compiler/plugins/input/StableHLO/Conversion/test/CMakeLists.txt
@@ -19,6 +19,7 @@
"legalize_chlo_decomposition.mlir"
"legalize_chlo_no_broadcast.mlir"
"legalize_chlo_with_broadcast.mlir"
+ "legalize_composite.mlir"
"legalize_control_flow.mlir"
"legalize_shape_computations.mlir"
"stablehlo_custom_calls.mlir"
diff --git a/compiler/plugins/input/StableHLO/Conversion/test/legalize_composite.mlir b/compiler/plugins/input/StableHLO/Conversion/test/legalize_composite.mlir
new file mode 100644
index 0000000..e8dfbdb
--- /dev/null
+++ b/compiler/plugins/input/StableHLO/Conversion/test/legalize_composite.mlir
@@ -0,0 +1,48 @@
+// RUN: iree-opt --split-input-file --iree-stablehlo-input-transformation-pipeline %s | FileCheck %s
+
+// Composite ops must be replaced by their decomposition, which is then inlined
+// and lowered like any other StableHLO computation.
+
+// CHECK-LABEL: func.func @acos_composite
+// CHECK-NOT: stablehlo.composite
+// CHECK-NOT: func.call
+// CHECK: linalg.generic
+// CHECK: math.atan2
+func.func @acos_composite(%arg0: tensor<8x32xf32>) -> tensor<8x32xf32> {
+ %0 = stablehlo.composite "chlo.acos" %arg0 {decomposition = @chlo.acos.impl, version = 1 : i32} : (tensor<8x32xf32>) -> tensor<8x32xf32>
+ return %0 : tensor<8x32xf32>
+}
+func.func private @chlo.acos.impl(%arg0: tensor<8x32xf32>) -> tensor<8x32xf32> {
+ %cst = stablehlo.constant dense<1.000000e+00> : tensor<8x32xf32>
+ %0 = stablehlo.subtract %cst, %arg0 : tensor<8x32xf32>
+ %1 = stablehlo.add %cst, %arg0 : tensor<8x32xf32>
+ %2 = stablehlo.multiply %0, %1 : tensor<8x32xf32>
+ %3 = stablehlo.sqrt %2 : tensor<8x32xf32>
+ %4 = stablehlo.atan2 %3, %arg0 : tensor<8x32xf32>
+ return %4 : tensor<8x32xf32>
+}
+
+// -----
+
+// A composite taking multiple operands and returning multiple results.
+
+// CHECK-LABEL: func.func @multi_result_composite
+// CHECK-NOT: stablehlo.composite
+// CHECK-NOT: func.call
+// CHECK: arith.addf
+// CHECK: arith.mulf
+func.func @multi_result_composite(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>)
+ -> (tensor<4xf32>, tensor<4xf32>) {
+ %0:2 = stablehlo.composite "my_namespace.add_mul" %arg0, %arg1 {
+ decomposition = @add_mul.impl,
+ composite_attributes = {my_attribute = "my_value"},
+ version = 1 : i32
+ } : (tensor<4xf32>, tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>)
+ return %0#0, %0#1 : tensor<4xf32>, tensor<4xf32>
+}
+func.func private @add_mul.impl(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>)
+ -> (tensor<4xf32>, tensor<4xf32>) {
+ %0 = stablehlo.add %arg0, %arg1 : tensor<4xf32>
+ %1 = stablehlo.multiply %arg0, %arg1 : tensor<4xf32>
+ return %0, %1 : tensor<4xf32>, tensor<4xf32>
+}
