Pass to wrap StableHLO ops in composite (#2722)

Wraps StableHLO operations in `stablehlo.composite` operations.

    For instance, consider a simple StableHLO program:

    ```mlir
func.func @main(%arg0 : tensor<2xf32>, %arg1 : tensor<2xf32>) ->
tensor<2xf32> {
        %0 = stablehlo.add %arg0, %arg1 : tensor<2xf32>
        return %0 : tensor<2xf32>
      }
    ```

Applying this pass to wrap `stablehlo.add` operations will result in the
    following program:

    ```mlir
func.func @main(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) ->
tensor<2xf32> {
%0 = stablehlo.composite "stablehlo.add" %arg0, %arg1 {decomposition =
@stablehlo.add.impl} : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
        return %0 : tensor<2xf32>
      }
func.func private @stablehlo.add.impl(%arg0: tensor<2xf32>, %arg1:
tensor<2xf32>) -> tensor<2xf32> {
        %0 = stablehlo.add %arg0, %arg1 : tensor<2xf32>
        return %0 : tensor<2xf32>
      }
    ```

    Notes:

- The `name` attribute of the generated `stablehlo.composite` operation
will always be the same as the name of the original operation that was
wrapped (e.g., if you wrap a `stablehlo.add` operation, the composite
        will also be named `"stablehlo.add"`).
      - The private function that encapsulates the original operation
        (referenced by the `decomposition` attribute of the
        `stablehlo.composite` operation) will be named using the pattern
`<op_name>.impl[.N]`, where `<op_name>` is the name of the original
operation, and `N` is a unique integer identifier generated to prevent
        naming conflicts within the module.

    This pass can be used in three distinct ways:

    **Mode 1: Command-line Usage**

This mode is the simplest, using the `stablehlo-opt` utility with the
`op-names` (a comma-separated list of operation names) and `version` (an
integer version number) options. It wraps **all instances** of specified
operations. The attributes of the newly created `stablehlo.composite`
operation will be the same as the attributes of the original operation.

    **Usage Example:**

    ```bash
stablehlo-opt input.mlir
--stablehlo-wrap-in-composite=op-names='stablehlo.add,stablehlo.mul' -o
output.mlir
    ```

    **Mode 2: Programmatic Single-Op Wrapping**

    This mode provides programmatic control to wrap
    **a specific operation instance** and returns a pointer to the newly
    created `stablehlo.composite` operation.

    **Example (C++):**

    ```cpp
      // To wrap a specific stablehlo.add instance

mlir::stablehlo::AddOp addOp = ...; // The op instanced to be wrapped.
      mlir::ModuleOp module = addOp->getParentOfType<mlir::ModuleOp>();
      mlir::OpBuilder builder(addOp);
mlir::NamedAttrList attrs = ...; // Attributes to be set on the
composite op.
      int32_t version = 0; // Composite version.

mlir::stablehlo::CompositeOp compositeOp =
mlir::stablehlo::wrapOperationInComposite(builder, addOp, attrs,
version, module);
      addOp.replaceAllUsesWith(compositeOp);
    ```

**Mode 3: Programmatic Module-Wide Wrapping with Attribute Predicates**

This mode extends programmatic wrapping to the entire module, offering
fine-grained control over which operations are wrapped and their
attributes.
This is achieved by using the `createStablehloWrapInCompositePass` API,
    which takes an `AttributePredicateMap` as an argument.

The `AttributePredicateMap` is a map that dictates which operations
should
be considered for wrapping and how their attributes should be handled.
Its
    semantics are as follows:

- **Keys (mlir::TypeID):** `TypeID` of an MLIR operation. If an
operation's
`TypeID` matches a key in the map, it becomes a candidate for wrapping.
    - **Values (Lambda Functions):** Lambda function of type
`std::function<std::optional<NamedAttrList>(Operation*)>`. This function
        is applied to each candidate operation.
        - **Input:** An `mlir::Operation*`, which is an instance of the
          operation type corresponding to the `TypeID` key.
        - **Return Value:** An `std::optional<NamedAttrList>`.
          - If the lambda returns a `NamedAttrList` (wrapped in
            `std::optional`), the operation is wrapped in a
`stablehlo::composite` operation, and the returned attributes are
            used to set the composite's attributes.
- If the lambda returns `std::nullopt`, the operation is **not**
            wrapped. This allows for selective wrapping based on custom
            criteria.

    **Example (C++):**

    ```cpp

    // ... inside a pass or function ...

    stablehlo::AttributePredicateMap attributePredicateMap;

    attributePredicateMap[mlir::TypeID::get<mlir::stablehlo::AddOp>()] =
      [](mlir::Operation* op) -> std::optional<mlir::NamedAttrList> {
      // Custom logic to determine if and how to wrap the operation.
      // Example: Only wrap if it's on a specific type.
      if (op->getOperand(0).getType().isa<mlir::Float32Type>()) {
        return mlir::NamedAttrList(op->getAttrs());
      }
      return std::nullopt; // Do not wrap.
    };

pm.addPass(createStablehloWrapInCompositePass(attributePredicateMap,
compositeVersion));
    if (mlir::failed(pm.run(module))) {
      return;
    }
    ```

Author: sdasgup3

BUILD.bazel

@@ -1246,6 +1246,7 @@ cc_library(
         "stablehlo/transforms/StablehloLegalizeToVhlo.cpp",
         "stablehlo/transforms/StablehloRefineArguments.cpp",
         "stablehlo/transforms/StablehloRefineShapes.cpp",
+        "stablehlo/transforms/StablehloWrapInComposite.cpp",
         "stablehlo/transforms/VhloLegalizeToStablehlo.cpp",
         "stablehlo/transforms/VhloToVersion.cpp",
     ],

docs/generated/stablehlo_passes.md

@@ -301,8 +301,8 @@ _Refines shapes across a StableHLO program._
   The flagship use case for this pass is specializing dynamically-shaped
   programs to static shapes. If a dynamically-shaped StableHLO program has the
   right structure, then updating its argument types from dynamic shapes to
-  static shapes and running this pass will propagate static shapes across
-  the program.
+  static shapes and running this pass will propagate static shapes across the
+  program.
 
   This pass removes `custom_call @shape_refinement_operand_wrapper` by
   replacing uses of the result with the operand directly, and propagates
@@ -338,6 +338,121 @@ Modules valid for shape refinement must have the following properties:
   * All calls to a single function resolve to the same argument shapes, and no
     recursive / co-recursive function calls are made.
 
+### `-stablehlo-wrap-in-composite`
+
+_Wraps a non-composite  StableHLO op in a composite op._
+
+Wraps StableHLO operations in `stablehlo.composite` operations.
+
+For instance, consider a simple StableHLO program:
+
+```mlir
+func.func @main(%arg0 : tensor<2xf32>, %arg1 : tensor<2xf32>) -> tensor<2xf32> {
+  %0 = stablehlo.add %arg0, %arg1 : tensor<2xf32>
+  return %0 : tensor<2xf32>
+}
+```
+
+Applying this pass to wrap `stablehlo.add` operations will result in the
+following program:
+
+```mlir
+func.func @main(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {
+  %0 = stablehlo.composite "stablehlo.add" %arg0, %arg1 {decomposition = @stablehlo.add.impl} : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
+  return %0 : tensor<2xf32>
+}
+func.func private @stablehlo.add.impl(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {
+  %0 = stablehlo.add %arg0, %arg1 : tensor<2xf32>
+  return %0 : tensor<2xf32>
+}
+```
+
+Notes:
+
+  - The `name` attribute of the generated `stablehlo.composite` operation
+    will always be the same as the name of the original operation that was
+    wrapped (e.g., if you wrap a `stablehlo.add` operation, the composite
+    will also be named `"stablehlo.add"`).
+  - The private function that encapsulates the original operation
+    (referenced by the `decomposition` attribute of the
+    `stablehlo.composite` operation) will be named using the pattern
+    `<op_name>.impl[.N]`, where `<op_name>` is the name of the original
+    operation, and `N` is a unique integer identifier generated to prevent
+    naming conflicts within the module.
+
+This pass can be used in two distinct ways:
+
+**Mode 1: Command-line Usage**
+
+This mode is intended for debugging or testing, as it offers minimal control
+over the attributes of the generated `stablehlo.composite` operations.
+It wraps **all instances** of operations specified using the `op-names`
+(a comma-separated list of operation names) options. The attributes of the
+newly created `stablehlo.composite` operation will be the same as the
+attributes of the original operation.
+
+**Usage Example:**
+
+```bash
+stablehlo-opt input.mlir --stablehlo-wrap-in-composite=op-names='stablehlo.add,stablehlo.mul' -o output.mlir
+```
+
+**Mode 2: Programmatic Module-Wide Wrapping with customized Attribute Handling**
+
+This mode extends programmatic wrapping to the entire module, offering
+fine-grained control over which operations are wrapped and their attributes.
+This is achieved by using the `createStablehloWrapInCompositePass` API,
+which takes an `CompositeAttributeProviderMap` as an argument.
+
+The `CompositeAttributeProviderMap` is a map that dictates which operations
+should be considered for wrapping and how their attributes should be
+handled. Its semantics are as follows:
+
+- **Keys (mlir::TypeID):** `TypeID` of an MLIR operation. If an operation's
+    `TypeID` matches a key in the map, it becomes a candidate for wrapping.
+- **Values (Lambda Functions):** Lambda function of type
+    `std::function<std::optional<NamedAttrList>(Operation*)>`. This function
+    is applied to each candidate operation.
+    - **Input:** An `mlir::Operation*`, which is an instance of the
+      operation type corresponding to the `TypeID` key.
+    - **Return Value:** An `std::optional<NamedAttrList>`.
+      - If the lambda returns a `NamedAttrList` (wrapped in
+        `std::optional`), the operation is wrapped in a
+        `stablehlo::composite` operation, and the returned attributes are
+        used to set the composite's attributes.
+      - If the lambda returns `std::nullopt`, the operation is **not**
+        wrapped. This allows for selective wrapping based on custom
+        criteria.
+
+**Example (C++):**
+
+```cpp
+
+stablehlo::CompositeAttributeProviderMap compositeAttributeProviderMap;
+
+compositeAttributeProviderMap[mlir::TypeID::get<mlir::stablehlo::AddOp>()] =
+  [](mlir::Operation* op) -> std::optional<mlir::NamedAttrList> {
+  // Custom logic to determine if and how to wrap the operation.
+  // Example: Only wrap if it's on a specific type.
+  if (op->getOperand(0).getType().isa<mlir::Float32Type>()) {
+    return mlir::NamedAttrList(op->getAttrs());
+  }
+  return std::nullopt; // Do not wrap.
+};
+
+pm.addPass(createStablehloWrapInCompositePass(compositeAttributeProviderMap, compositeVersion));
+if (mlir::failed(pm.run(module))) {
+  return;
+}
+```
+
+#### Options
+
+```
+-op-names : The names of the ops to wrap.
+-version  : The version number of the composite op.
+```
+
 ### `-vhlo-legalize-to-stablehlo`
 
 _Legalize VHLO to StableHLO._

stablehlo/tests/transforms/stablehlo_wrap_in_composite.mlir

@@ -0,0 +1,88 @@
+// RUN: stablehlo-opt --stablehlo-wrap-in-composite='op-names=stablehlo.add,stablehlo.convolution,stablehlo.reduce version=1' --split-input-file --verify-diagnostics %s | FileCheck %s
+
+// CHECK-LABEL: func.func @wrap_in_composite
+// CHECK-SAME: %[[ARG_0:.*]]: tensor<64x8x8x8xi8>,
+// CHECK-SAME: %[[ARG_1:.*]]: tensor<4x4x8x32xi8>,
+// CHECK-SAME: %[[ARG_2:.*]]: tensor<64x3x3x32xi32>) -> tensor<64x3x3x32xi32> {
+// CHECK: %[[CONV:.*]] = stablehlo.composite "stablehlo.convolution" %[[ARG_0]], %[[ARG_1]] {
+// CHECK-SAME: composite_attributes = {batch_group_count = 1 : i64,
+// CHECK-SAME: dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
+// CHECK-SAME: feature_group_count = 1 : i64,
+// CHECK-SAME{LITERAL}: padding = dense<[[0, 1], [0, 1]]> : tensor<2x2xi64>,
+// CHECK-SAME{LITERAL}: rhs_dilation = array<i64: 2, 2>,
+// CHECK-SAME{LITERAL}: window_strides = array<i64: 1, 1>},
+// CHECK-SAME: decomposition = @stablehlo.convolution.impl,
+// CHECK-SAME: version = 1 : i32} : (tensor<64x8x8x8xi8>, tensor<4x4x8x32xi8>) -> tensor<64x3x3x32xi32>
+// CHECK: %[[ADD:.*]] = stablehlo.composite "stablehlo.add" %[[CONV]], %[[ARG_2]] {
+// CHECK-SAME: decomposition = @stablehlo.add.impl,
+// CHECK-SAME: version = 1 : i32} : (tensor<64x3x3x32xi32>, tensor<64x3x3x32xi32>) -> tensor<64x3x3x32xi32>
+// CHECK-NEXT: return %[[ADD]]
+
+// CHECK-LABEL: func.func private @stablehlo.add.impl
+// CHECK-SAME: %[[ARG_0:.*]]: tensor<64x3x3x32xi32>,
+// CHECK-SAME: %[[ARG_1:.*]]: tensor<64x3x3x32xi32>) -> tensor<64x3x3x32xi32> {
+// CHECK: %[[VAL:.*]] = stablehlo.add %[[ARG_0]], %[[ARG_1]] : tensor<64x3x3x32xi32>
+// CHECK-NEXT: return %[[VAL]]
+
+// CHECK-LABEL: func.func private @stablehlo.convolution.impl
+// CHECK-SAME: %[[ARG_0:.*]]: tensor<64x8x8x8xi8>,
+// CHECK-SAME: %[[ARG_1:.*]]: tensor<4x4x8x32xi8>) -> tensor<64x3x3x32xi32> {
+// CHECK: %[[VAL:.*]] = stablehlo.convolution(%[[ARG_0]], %[[ARG_1]])
+// CHECK-SAME{LITERAL}: dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f],
+// CHECK-SAME{LITERAL}: stride = [1, 1],
+// CHECK-SAME{LITERAL}: pad = [[0, 1], [0, 1]],
+// CHECK-SAME{LITERAL}: rhs_dilate = [2, 2]}
+// CHECK-SAME: batch_group_count = 1 : i64
+// CHECK-SAME: feature_group_count = 1 : i64
+// CHECK-SAME: : (tensor<64x8x8x8xi8>, tensor<4x4x8x32xi8>) -> tensor<64x3x3x32xi32>
+// CHECK-NEXT: return %[[VAL]]
+
+func.func @wrap_in_composite(
+    %arg0: tensor<64x8x8x8xi8>,
+    %arg1: tensor<4x4x8x32xi8>,
+    %arg2: tensor<64x3x3x32xi32>) -> tensor<64x3x3x32xi32> {
+  %0 = stablehlo.convolution(%arg0, %arg1)
+         dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f],
+         window = {stride = [1, 1], pad = [[0, 1], [0, 1]], rhs_dilate = [2, 2]}
+         {batch_group_count = 1 : i64, feature_group_count = 1 : i64} :
+       (tensor<64x8x8x8xi8>, tensor<4x4x8x32xi8>) -> tensor<64x3x3x32xi32>
+  %1 = stablehlo.add %0, %arg2 : tensor<64x3x3x32xi32>
+  func.return %1 : tensor<64x3x3x32xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @wrap_in_composite_op_with_region
+// CHECK-SAME: %[[ARG_0:.*]]: tensor<4x3xf32>) -> tensor<4xf32>
+// CHECK: %[[CONST:.*]] = stablehlo.constant
+// CHECK-NEXT: %[[COMPOSITE_REDUCE:.*]] = stablehlo.composite "stablehlo.reduce" %[[ARG_0]], %[[CONST]] {
+// CHECK-SAME: composite_attributes = {
+// CHECK-SAME: dimensions = array<i64: 1>},
+// CHECK-SAME: decomposition = @stablehlo.reduce.impl,
+// CHECK-SAME: version = 1 : i32} : (tensor<4x3xf32>, tensor<f32>) -> tensor<4xf32>
+// CHECK-NEXT: return %[[COMPOSITE_REDUCE]]
+
+// CHECK-LABEL: func.func private @stablehlo.reduce.impl
+// CHECK-SAME: %[[ARG_0:.*]]: tensor<4x3xf32>,
+// CHECK-SAME: %[[ARG_1:.*]]: tensor<f32>) -> tensor<4xf32> {
+// CHECK: %[[REDUCE:.*]] = stablehlo.reduce(%[[ARG_0]] init: %[[ARG_1]])
+// CHECK-SAME{LITERAL}: applies stablehlo.add across dimensions = [1]
+// CHECK-SAME: (tensor<4x3xf32>, tensor<f32>) -> tensor<4xf32>
+// CHECK-NEXT: return %[[REDUCE]]
+func.func @wrap_in_composite_op_with_region(%x : tensor<4x3xf32>) -> tensor<4xf32> {
+  %cst = stablehlo.constant dense<2.7> : tensor<f32>
+  %res = stablehlo.reduce(%x init: %cst) applies stablehlo.add across dimensions = [1] : (tensor<4x3xf32>, tensor<f32>) -> tensor<4xf32>
+  func.return %res : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @cannot_be_wrapped_ops_does_not_match
+// CHECK-SAME: %[[ARG_0:.*]]: tensor<2xf32>,
+// CHECK-SAME: %[[ARG_1:.*]]: tensor<2xf32>) -> tensor<2xf32> {
+// CHECK: %[[VAL:.*]] = stablehlo.multiply %[[ARG_0]], %[[ARG_1]] : tensor<2xf32>
+// CHECK-NEXT: return %[[VAL]] : tensor<2xf32>
+func.func @cannot_be_wrapped_ops_does_not_match(%arg0: tensor<2xf32>, %arg1: tensor<2xf32>) -> tensor<2xf32> {
+  %0 = stablehlo.multiply %arg0, %arg1 : tensor<2xf32>
+  func.return %0 : tensor<2xf32>
+}

stablehlo/transforms/CMakeLists.txt

@@ -57,6 +57,7 @@ add_mlir_dialect_library(StablehloPasses
   StablehloLegalizeToVhlo.cpp
   StablehloRefineArguments.cpp
   StablehloRefineShapes.cpp
+  StablehloWrapInComposite.cpp
   VhloLegalizeToStablehlo.cpp
   VhloToVersion.cpp
   PassUtils.cpp

stablehlo/transforms/Passes.h

@@ -16,16 +16,25 @@ limitations under the License.
 #ifndef STABLEHLO_TRANSFORMS_PASSES_H
 #define STABLEHLO_TRANSFORMS_PASSES_H
 
+#include <cstdint>
+#include <functional>
 #include <memory>
+#include <optional>
 
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/Quant/IR/Quant.h"
-#include "mlir/Dialect/Shape/IR/Shape.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // IWYU pragma: keep
+#include "mlir/Dialect/Quant/IR/Quant.h"   // IWYU pragma: keep
+#include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/TypeRange.h"
 #include "mlir/Pass/Pass.h"
-#include "mlir/Support/LogicalResult.h"
+#include "mlir/Pass/PassOptions.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Support/TypeID.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "stablehlo/dialect/StablehloOps.h"
 #include "stablehlo/dialect/Version.h"
 
 namespace mlir {
@@ -102,6 +111,43 @@ void populateStablehloCompatibilityExpanderPatterns(
 std::unique_ptr<OperationPass<ModuleOp>> createStablehloRefineArgumentsPass(
     TypeRange refinedTypes);
 
+/// Creates a pass that wraps StableHLO ops in CompositeOp.
+/// The pass takes in a map from op's type id to a function that returns the
+/// attributes to be added to the CompositeOp. The pass also takes in a
+/// version number for the CompositeOp.
+using CompositeAttributeProvider =
+    std::function<std::optional<NamedAttrList>(Operation *)>;
+using CompositeAttributeProviderMap =
+    llvm::DenseMap<mlir::TypeID, CompositeAttributeProvider>;
+std::unique_ptr<OperationPass<ModuleOp>> createStablehloWrapInCompositePass(
+    const CompositeAttributeProviderMap &compositeAttributeProviderMap,
+    int32_t compositeVersion);
+
+/// Wraps the given operation in a CompositeOp with the specified NamedAttrs and
+/// version and returns the CompositeOp.
+///
+/// **A typical usage **
+///
+/// ```cpp
+/// // To wrap a specific stablehlo.add instance
+///
+/// mlir::stablehlo::AddOp addOp = ...; // The op instanced to be wrapped.
+/// mlir::ModuleOp module = addOp->getParentOfType<mlir::ModuleOp>();
+/// mlir::OpBuilder builder(addOp);
+/// mlir::NamedAttrList attrs = ...; // Attributes to be set on the
+///                                  // composite op.
+/// int32_t version = 0; // Composite version.
+///
+/// mlir::stablehlo::CompositeOp compositeOp =
+///   mlir::stablehlo::wrapOperationInComposite(builder, addOp, attrs,
+///                                             version, module);
+/// addOp.replaceAllUsesWith(compositeOp);
+/// ```
+stablehlo::CompositeOp wrapOperationInComposite(OpBuilder &builder,
+                                                Operation *op,
+                                                const NamedAttrList &attrs,
+                                                int32_t compositeVersion,
+                                                ModuleOp module);
 //// Pass pipelines ////
 
 // StableHLO consumers can add this pipeline to convert portable artifacts to