[Mlir-commits] [mlir] [mlir] targeted verification for transform "inlining" (PR #192956)

[William Moses]

wsmoses wrote:

I think this PR broke the Reactant raising compilation pipeline:

```
// RUN: enzymexlamlir-opt %s --transform-interpreter

module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg0: !transform.any_op) {
    %0 = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
    transform.apply_patterns to %0 {
      transform.apply_patterns.enzyme_hlo.compare_op_canon {benefit = 16 : i64}
    } : !transform.any_op
    transform.yield 
  }
  func.func @main(%arg0: tensor<64xbf16>) -> tensor<64xbf16> {
    %c = stablehlo.constant dense<0> : tensor<i64>
    %c_0 = stablehlo.constant dense<64> : tensor<i64>
    %c_1 = stablehlo.constant dense<1> : tensor<i64>
    %0 = enzymexla.kernel_call @f blocks in(%c_1, %c_1, %c_1) threads in(%c_0, %c_1, %c_1) shmem = %c (%arg0) {output_operand_aliases = [#stablehlo.output_operand_alias<output_tuple_indices = [], operand_index = 0, operand_tuple_indices = []>], xla_side_effect_free} : (tensor<64xbf16>) -> tensor<64xbf16>
    return %0 : tensor<64xbf16>
  }
  llvm.func ptx_kernelcc @f(%arg0: !llvm.ptr<1> {llvm.align = 128 : i64, llvm.dereferenceable = 128 : i64, llvm.dereferenceable_or_null = 128 : i64, llvm.noalias, llvm.nocapture, llvm.nofree, llvm.noundef, llvm.writeonly}) {
    llvm.return
  }
}
```

specifically if there was an op that normally wasn't inlinable, it now throws an error
```
wmoses-mac:Enzyme-JaX wmoses$ ./bazel-bin/enzymexlamlir-opt test/lit_tests/raising/mwe.mlir --transform-interpreter
test/lit_tests/raising/mwe.mlir:15:10: error: merged call is not legal to inline into its caller
    %0 = enzymexla.kernel_call @f blocks in(%c_1, %c_1, %c_1) threads in(%c_0, %c_1, %c_1) shmem = %c (%arg0) {output_operand_aliases = [#stablehlo.output_operand_alias<output_tuple_indices = [], operand_index = 0, operand_tuple_indices = []>], xla_side_effect_free} : (tensor<64xbf16>) -> tensor<64xbf16>
         ^
test/lit_tests/raising/mwe.mlir:15:10: note: see current operation: %3 = "enzymexla.kernel_call"(%2, %2, %2, %1, %2, %2, %0, %arg0) <{backend_config = "", fn = @f, operandSegmentSizes = array<i32: 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 1>, output_operand_aliases = [#stablehlo.output_operand_alias<output_tuple_indices = [], operand_index = 0, operand_tuple_indices = []>], xla_side_effect_free}> : (tensor<i64>, tensor<i64>, tensor<i64>, tensor<i64>, tensor<i64>, tensor<i64>, tensor<i64>, tensor<64xbf16>) -> tensor<64xbf16>
test/lit_tests/raising/mwe.mlir:18:3: note: callee defined here
  llvm.func ptx_kernelcc @f(%arg0: !llvm.ptr<1> {llvm.align = 128 : i64, llvm.dereferenceable = 128 : i64, llvm.dereferenceable_or_null = 128 : i64, llvm.noalias, llvm.nocapture, llvm.nofree, llvm.noundef, llvm.writeonly}) {
  ^
test/lit_tests/raising/mwe.mlir:3:1: error: failed to merge symbols
module attributes {transform.with_named_sequence} {
^
test/lit_tests/raising/mwe.mlir:3:1: note: see current operation: 
"builtin.module"() ({
  "transform.named_sequence"() <{function_type = (!transform.any_op) -> (), sym_name = "__transform_main"}> ({
  ^bb0(%arg2: !transform.any_op):
    %4 = "transform.structured.match"(%arg2) <{ops = ["func.func"]}> : (!transform.any_op) -> !transform.any_op
    "transform.apply_patterns"(%4) <{max_iterations = -1 : i64, max_num_rewrites = -1 : i64}> ({
      "transform.apply_patterns.enzyme_hlo.compare_op_canon"() <{benefit = 16 : i64}> : () -> ()
    }) : (!transform.any_op) -> ()
    "transform.yield"() : () -> ()
  }) : () -> ()
  "func.func"() <{function_type = (tensor<64xbf16>) -> tensor<64xbf16>, sym_name = "main"}> ({
  ^bb0(%arg1: tensor<64xbf16>):
    %0 = "stablehlo.constant"() <{value = dense<0> : tensor<i64>}> : () -> tensor<i64>
    %1 = "stablehlo.constant"() <{value = dense<64> : tensor<i64>}> : () -> tensor<i64>
    %2 = "stablehlo.constant"() <{value = dense<1> : tensor<i64>}> : () -> tensor<i64>
    %3 = "enzymexla.kernel_call"(%2, %2, %2, %1, %2, %2, %0, %arg1) <{backend_config = "", fn = @f, operandSegmentSizes = array<i32: 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 1>, output_operand_aliases = [#stablehlo.output_operand_alias<output_tuple_indices = [], operand_index = 0, operand_tuple_indices = []>], xla_side_effect_free}> : (tensor<i64>, tensor<i64>, tensor<i64>, tensor<i64>, tensor<i64>, tensor<i64>, tensor<i64>, tensor<64xbf16>) -> tensor<64xbf16>
    "func.return"(%3) : (tensor<64xbf16>) -> ()
  }) : () -> ()
  "llvm.func"() <{CConv = #llvm.cconv<ptx_kernelcc>, arg_attrs = [{llvm.align = 128 : i64, llvm.dereferenceable = 128 : i64, llvm.dereferenceable_or_null = 128 : i64, llvm.noalias, llvm.nocapture, llvm.nofree, llvm.noundef, llvm.writeonly}], function_type = !llvm.func<void (ptr<1>)>, linkage = #llvm.linkage<external>, sym_name = "f", unnamed_addr = 0 : i64, visibility_ = 0 : i64}> ({
  ^bb0(%arg0: !llvm.ptr<1>):
    "llvm.return"() : () -> ()
  }) : () -> ()
}) {transform.with_named_sequence} : () -> ()
```


https://github.com/llvm/llvm-project/pull/192956