Pass to wrap StableHLO ops in composite

[sdasgup3]

The pass wraps StableHLO ops, as specified by a pass option flag, in a composite op. The composite op will inherit all attributes of the original op.

For example, using the option flag --stablehlo-wrap-in-composite=op-names='stablehlo.add':

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

will become:

    func.func @add(%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>, tensor<2xf32>) -> tensor<2xf32>
      func.return %0 : tensor<2xf32>
    }

The pass is also exposed as an API createStablehloWrapInCompositePass to allow for more flexible selection of ops to wrap.
For example, the following will wrap all non-composite ops that are not stablehlo.add or stablehlo.convolution:

```c++
auto pass = createStablehloWrapInCompositePass(
   (Operation *op) {
      return (op->getName().getStringRef() == "stablehlo.add" ||
              op->getName().getStringRef() == "stablehlo.convolution") &&
             !isa<stablehlo::CompositeOp>(op);
    });
```

[gemini-code-assist]

Important

The terms of service for this installation has not been accepted. Please ask the Organization owners to visit the Gemini Code Assist Admin Console to sign it.

[GleasonK]

Let's talk about the design more. It's missing a few things to be generally useful:

• What to do with ops that have attributes?
• What should the composites be named, or should the name be user specified as a part of the callback fn, etc

Need to think more, maybe we have a little brainstorm and run through some examples

[sdasgup3]

All good points! SGTM.

What to do with ops that have attributes?

All the op attributes are preserved as composite attributes. To demonstrate that I have used convolution op as an example in stablehlo/tests/transforms/stablehlo_wrap_in_composite.mlir.

[ghpvnist]

Can we have an example with regions like reduce?
Should this pass run as part of a pass pipeline? If we run a decomposer pass that introduces stablehlo.convolution, they also need to get wrapped right?

[sdasgup3]

Added.

Should this pass run as part of a pass pipeline?

Yes, that is the more practical usage of the pass.

If we run a decomposer pass that introduces stablehlo.convolution, they also need to get wrapped right?

If we have a decomposer pass and it emit a convolution and if the user intend to wrap that conv into a composite, then yes it will be wrapped.
Let me know if that address your concern.