[Mlir-commits] [mlir] [MLIR] Fix ErasedOpsListener false positives for newly created ops/blocks (PR #192291)

[Mehdi Amini]

joker-eph wrote:

The case I'm looking at:
func.func @vector_and_reduction() {
  %v_mask = vector.constant_mask [1] : vector<2xi1>
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  %c2 = arith.constant 2 : index
  %result = scf.parallel (%i) = (%c0) to (%c2) step (%c1) init(%v_mask) -> vector<2xi1> {
    %val = vector.constant_mask [1] : vector<2xi1>
    scf.reduce (%val : vector<2xi1>) {
    ^bb0(%lhs: vector<2xi1>, %rhs: vector<2xi1>):
      %0 = arith.andi %lhs, %rhs : vector<2xi1>
      scf.reduce.return %0 : vector<2xi1>
    }
  }
  return
}

For SCFToOpenMP, to remove the scf.parallel we'll build a omp.parallel op first to move the body there:

  "omp.parallel"() <{operandSegmentSizes = array<i32: 0, 0, 0, 0, 0, 0>}> ({
  ^bb0:
  }) : () -> ()
  %6 = "scf.parallel"(%1, %3, %2, %0) <{operandSegmentSizes = array<i32: 1, 1, 1, 1>}> ({
  ^bb0(%arg0: index, %arg1: !llvm.ptr):
    %7 = "vector.constant_mask"() <{mask_dim_sizes = array<i64: 1>}> : () -> vector<2xi1>
    %8 = "llvm.load"(%arg1) <{ordering = 0 : i64}> : (!llvm.ptr) -> vector<2xi1>
    %9 = "arith.andi"(%8, %7) : (vector<2xi1>, vector<2xi1>) -> vector<2xi1>
    "llvm.store"(%9, %arg1) <{ordering = 0 : i64}> : (vector<2xi1>, !llvm.ptr) -> ()
  }) : (index, index, index, vector<2xi1>) -> vector<2xi1>then we move operations and restructure until we get to:

  "omp.parallel"() <{operandSegmentSizes = array<i32: 0, 0, 0, 0, 0, 0>}> ({
    "omp.wsloop"(%5) <{operandSegmentSizes = array<i32: 0, 0, 0, 0, 0, 1, 0>, reduction_byref = array<i1: false>, reduction_syms = [@__scf_reduction]}> ({
    ^bb0(%arg0: !llvm.ptr):
      "omp.loop_nest"(%1, %3, %2) <{collapse_num_loops = 1 : i64}> ({
      ^bb0(%arg1: index):
      ^bb1:  // no predecessors
        %7 = "vector.constant_mask"() <{mask_dim_sizes = array<i64: 1>}> : () -> vector<2xi1>
        %8 = "llvm.load"(%arg0) <{ordering = 0 : i64}> : (!llvm.ptr) -> vector<2xi1>
        %9 = "arith.andi"(%8, %7) : (vector<2xi1>, vector<2xi1>) -> vector<2xi1>
        "llvm.store"(%9, %arg0) <{ordering = 0 : i64}> : (vector<2xi1>, !llvm.ptr) -> ()
      }) : (index, index, index) -> ()
    }) : (!llvm.ptr) -> ()
    "omp.terminator"() : () -> ()
  }) : () -> ()
  %6 = "scf.parallel"(%1, %3, %2, %0) <{operandSegmentSizes = array<i32: 1, 1, 1, 1>}> ({
  }) : (index, index, index, vector<2xi1>) -> vector<2xi1>

The scf.parallel is empty, everything is under the omp.parallel. The next thing to do is to move things into a nested AllocaScopeOp inside the loop_nest op:

"omp.parallel"() <{operandSegmentSizes = array<i32: 0, 0, 0, 0, 0, 0>}> ({
    "omp.wsloop"(%5) <{operandSegmentSizes = array<i32: 0, 0, 0, 0, 0, 1, 0>, reduction_byref = array<i1: false>, reduction_syms = [@__scf_reduction]}> ({
    ^bb0(%arg0: !llvm.ptr):
      "omp.loop_nest"(%1, %3, %2) <{collapse_num_loops = 1 : i64}> ({
      ^bb0(%arg1: index):
        "memref.alloca_scope"() ({
          %7 = "vector.constant_mask"() <{mask_dim_sizes = array<i64: 1>}> : () -> vector<2xi1>
          %8 = "llvm.load"(%arg0) <{ordering = 0 : i64}> : (!llvm.ptr) -> vector<2xi1>
          %9 = "arith.andi"(%8, %7) : (vector<2xi1>, vector<2xi1>) -> vector<2xi1>
          "llvm.store"(%9, %arg0) <{ordering = 0 : i64}> : (vector<2xi1>, !llvm.ptr) -> ()
          "memref.alloca_scope.return"() : () -> ()
        }) : () -> ()
        "omp.yield"() : () -> ()
      }) : (index, index, index) -> ()
    }) : (!llvm.ptr) -> ()
    "omp.terminator"() : () -> ()
  }) : () -> ()
  %6 = "scf.parallel"(%1, %3, %2, %0) <{operandSegmentSizes = array<i32: 1, 1, 1, 1>}> ({
  }) : (index, index, index, vector<2xi1>) -> vector<2xi1>

However when we do this, we have to call rewriter.mergeBlocks to move the ops into the alloca_scope. When we do that we delete the old block, but its parent is omp.loop_nest/omp.parallel, not scf.parallel, and the check is then failing.

So the restriction on not deleting anything not nested under the visited root op seems be too strict here.

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