mlir/test/IR/visitors.mlir

// RUN: mlir-opt -test-ir-visitors -split-input-file %s | FileCheck %s

// Verify the different configurations of IR visitors.
// Constant, yield and other terminator ops are not matched for simplicity.
// Module and function op and their immediately nested blocks are not erased in
// callbacks with return so that the output includes more cases in pre-order.

func.func @structured_cfg() {
  %c0 = arith.constant 0 : index
  %c1 = arith.constant 1 : index
  %c10 = arith.constant 10 : index
  scf.for %i = %c1 to %c10 step %c1 {
    %cond = "test.zero"(%i) : (index) -> (i1)
    scf.if %cond {
      "test.one"(%i) : (index) -> ()
    } else {
      "test.two"(%i) : (index) -> ()
    }
    "test.three"(%i) : (index) -> ()
  } {walk_blocks, walk_regions}
  return
}

// CHECK-LABEL: Op pre-order visit
// CHECK:       Visiting op 'builtin.module'
// CHECK:       Visiting op 'func.func'
// CHECK:       Visiting op 'scf.for'
// CHECK:       Visiting op 'test.zero'
// CHECK:       Visiting op 'scf.if'
// CHECK:       Visiting op 'test.one'
// CHECK:       Visiting op 'test.two'
// CHECK:       Visiting op 'test.three'
// CHECK:       Visiting op 'func.return'

// CHECK-LABEL: Block pre-order visits
// CHECK:       Visiting block ^bb0 from region 0 from operation 'builtin.module'
// CHECK:       Visiting block ^bb0 from region 0 from operation 'func.func'
// CHECK:       Visiting block ^bb0 from region 0 from operation 'scf.for'
// CHECK:       Visiting block ^bb0 from region 0 from operation 'scf.if'
// CHECK:       Visiting block ^bb0 from region 1 from operation 'scf.if'

// CHECK-LABEL: Region pre-order visits
// CHECK:       Visiting region 0 from operation 'builtin.module'
// CHECK:       Visiting region 0 from operation 'func.func'
// CHECK:       Visiting region 0 from operation 'scf.for'
// CHECK:       Visiting region 0 from operation 'scf.if'
// CHECK:       Visiting region 1 from operation 'scf.if'

// CHECK-LABEL: Op post-order visits
// CHECK:       Visiting op 'test.zero'
// CHECK:       Visiting op 'test.one'
// CHECK:       Visiting op 'test.two'
// CHECK:       Visiting op 'scf.if'
// CHECK:       Visiting op 'test.three'
// CHECK:       Visiting op 'scf.for'
// CHECK:       Visiting op 'func.return'
// CHECK:       Visiting op 'func.func'
// CHECK:       Visiting op 'builtin.module'

// CHECK-LABEL: Block post-order visits
// CHECK:       Visiting block ^bb0 from region 0 from operation 'scf.if'
// CHECK:       Visiting block ^bb0 from region 1 from operation 'scf.if'
// CHECK:       Visiting block ^bb0 from region 0 from operation 'scf.for'
// CHECK:       Visiting block ^bb0 from region 0 from operation 'func.func'
// CHECK:       Visiting block ^bb0 from region 0 from operation 'builtin.module'

// CHECK-LABEL: Region post-order visits
// CHECK:       Visiting region 0 from operation 'scf.if'
// CHECK:       Visiting region 1 from operation 'scf.if'
// CHECK:       Visiting region 0 from operation 'scf.for'
// CHECK:       Visiting region 0 from operation 'func.func'
// CHECK:       Visiting region 0 from operation 'builtin.module'

// CHECK-LABEL: Op reverse post-order visits
// CHECK:       Visiting op 'func.return'
// CHECK:       Visiting op 'scf.yield'
// CHECK:       Visiting op 'test.three'
// CHECK:       Visiting op 'scf.yield'
// CHECK:       Visiting op 'test.two'
// CHECK:       Visiting op 'scf.yield'
// CHECK:       Visiting op 'test.one'
// CHECK:       Visiting op 'scf.if'
// CHECK:       Visiting op 'test.zero'
// CHECK:       Visiting op 'scf.for'
// CHECK:       Visiting op 'arith.constant'
// CHECK:       Visiting op 'arith.constant'
// CHECK:       Visiting op 'arith.constant'
// CHECK:       Visiting op 'func.func'
// CHECK:       Visiting op 'builtin.module'

// CHECK-LABEL: Invoke block pre-order visits on blocks
// CHECK:       Visiting block ^bb0 from region 0 from operation 'scf.for'
// CHECK:       Visiting block ^bb0 from region 0 from operation 'scf.if'
// CHECK:       Visiting block ^bb0 from region 1 from operation 'scf.if'

// CHECK-LABEL: Invoke block post-order visits on blocks
// CHECK:       Visiting block ^bb0 from region 0 from operation 'scf.if'
// CHECK:       Visiting block ^bb0 from region 1 from operation 'scf.if'
// CHECK:       Visiting block ^bb0 from region 0 from operation 'scf.for'

// CHECK-LABEL: Invoke region pre-order visits on region
// CHECK:       Visiting region 0 from operation 'scf.for'
// CHECK:       Visiting region 0 from operation 'scf.if'
// CHECK:       Visiting region 1 from operation 'scf.if'

// CHECK-LABEL: Invoke region post-order visits on region
// CHECK:       Visiting region 0 from operation 'scf.if'
// CHECK:       Visiting region 1 from operation 'scf.if'
// CHECK:       Visiting region 0 from operation 'scf.for'

// CHECK-LABEL: Op pre-order erasures
// CHECK:       Erasing op 'scf.for'
// CHECK:       Erasing op 'func.return'

// CHECK-LABEL: Block pre-order erasures
// CHECK:       Erasing block ^bb0 from region 0 from operation 'scf.for'

// CHECK-LABEL: Op post-order erasures (skip)
// CHECK:       Erasing op 'test.zero'
// CHECK:       Erasing op 'test.one'
// CHECK:       Erasing op 'test.two'
// CHECK:       Erasing op 'scf.if'
// CHECK:       Erasing op 'test.three'
// CHECK:       Erasing op 'scf.for'
// CHECK:       Erasing op 'func.return'

// CHECK-LABEL: Block post-order erasures (skip)
// CHECK:       Erasing block ^bb0 from region 0 from operation 'scf.if'
// CHECK:       Erasing block ^bb0 from region 1 from operation 'scf.if'
// CHECK:       Erasing block ^bb0 from region 0 from operation 'scf.for'

// CHECK-LABEL: Op post-order erasures (no skip)
// CHECK:       Erasing op 'test.zero'
// CHECK:       Erasing op 'test.one'
// CHECK:       Erasing op 'test.two'
// CHECK:       Erasing op 'scf.if'
// CHECK:       Erasing op 'test.three'
// CHECK:       Erasing op 'scf.for'
// CHECK:       Erasing op 'func.return'
// CHECK:       Erasing op 'func.func'
// CHECK:       Erasing op 'builtin.module'

// CHECK-LABEL: Block post-order erasures (no skip)
// CHECK:       Erasing block ^bb0 from region 0 from operation 'scf.if'
// CHECK:       Erasing block ^bb0 from region 1 from operation 'scf.if'
// CHECK:       Erasing block ^bb0 from region 0 from operation 'scf.for'
// CHECK:       Erasing block ^bb0 from region 0 from operation 'func.func'
// CHECK:       Erasing block ^bb0 from region 0 from operation 'builtin.module'

// -----

func.func @unstructured_cfg() {
  "test.region.op0"() ({
    ^bb0:
      "test.op0"() : () -> ()
      cf.br ^bb2
    ^bb1:
      "test.op1"() : () -> ()
      cf.br ^bb2
    ^bb2:
      "test.op2"() : () -> ()
  }) : () -> ()
  return
}

// CHECK-LABEL: Op pre-order visits
// CHECK:       Visiting op 'builtin.module'
// CHECK:       Visiting op 'func.func'
// CHECK:       Visiting op 'test.region.op0'
// CHECK:       Visiting op 'test.op0'
// CHECK:       Visiting op 'cf.br'
// CHECK:       Visiting op 'test.op1'
// CHECK:       Visiting op 'cf.br'
// CHECK:       Visiting op 'test.op2'
// CHECK:       Visiting op 'func.return'

// CHECK-LABEL: Block pre-order visits
// CHECK:       Visiting block ^bb0 from region 0 from operation 'builtin.module'
// CHECK:       Visiting block ^bb0 from region 0 from operation 'func.func'
// CHECK:       Visiting block ^bb0 from region 0 from operation 'test.region.op0'
// CHECK:       Visiting block ^bb1 from region 0 from operation 'test.region.op0'
// CHECK:       Visiting block ^bb2 from region 0 from operation 'test.region.op0'

// CHECK-LABEL: Region pre-order visits
// CHECK:       Visiting region 0 from operation 'builtin.module'
// CHECK:       Visiting region 0 from operation 'func.func'
// CHECK:       Visiting region 0 from operation 'test.region.op0'

// CHECK-LABEL: Op post-order visits
// CHECK:       Visiting op 'test.op0'
// CHECK:       Visiting op 'cf.br'
// CHECK:       Visiting op 'test.op1'
// CHECK:       Visiting op 'cf.br'
// CHECK:       Visiting op 'test.op2'
// CHECK:       Visiting op 'test.region.op0'
// CHECK:       Visiting op 'func.return'
// CHECK:       Visiting op 'func.func'
// CHECK:       Visiting op 'builtin.module'

// CHECK-LABEL: Block post-order visits
// CHECK:       Visiting block ^bb0 from region 0 from operation 'test.region.op0'
// CHECK:       Visiting block ^bb1 from region 0 from operation 'test.region.op0'
// CHECK:       Visiting block ^bb2 from region 0 from operation 'test.region.op0'
// CHECK:       Visiting block ^bb0 from region 0 from operation 'func.func'
// CHECK:       Visiting block ^bb0 from region 0 from operation 'builtin.module'

// CHECK-LABEL: Region post-order visits
// CHECK:       Visiting region 0 from operation 'test.region.op0'
// CHECK:       Visiting region 0 from operation 'func.func'
// CHECK:       Visiting region 0 from operation 'builtin.module'

// CHECK-LABEL: Op reverse post-order visits
// CHECK:       Visiting op 'func.return'
// CHECK:       Visiting op 'test.op2'
// CHECK:       Visiting op 'cf.br'
// CHECK:       Visiting op 'test.op1'
// CHECK:       Visiting op 'cf.br'
// CHECK:       Visiting op 'test.op0'
// CHECK:       Visiting op 'test.region.op0'
// CHECK:       Visiting op 'func.func'
// CHECK:       Visiting op 'builtin.module'

// CHECK-LABEL: Block reverse post-order visits
// CHECK:       Visiting block ^bb2 from region 0 from operation 'test.region.op0'
// CHECK:       Visiting block ^bb1 from region 0 from operation 'test.region.op0'
// CHECK:       Visiting block ^bb0 from region 0 from operation 'test.region.op0'
// CHECK:       Visiting block ^bb0 from region 0 from operation 'func.func'
// CHECK:       Visiting block ^bb0 from region 0 from operation 'builtin.module'

// CHECK-LABEL: Region reverse post-order visits
// CHECK:       Visiting region 0 from operation 'test.region.op0'
// CHECK:       Visiting region 0 from operation 'func.func'
// CHECK:       Visiting region 0 from operation 'builtin.module'

// CHECK-LABEL: Op pre-order erasures (skip)
// CHECK:       Erasing op 'test.region.op0'
// CHECK:       Erasing op 'func.return'

// CHECK-LABEL: Block pre-order erasures (skip)
// CHECK:       Erasing block ^bb0 from region 0 from operation 'test.region.op0'
// CHECK:       Erasing block ^bb0 from region 0 from operation 'test.region.op0'
// CHECK:       Erasing block ^bb0 from region 0 from operation 'test.region.op0'

// CHECK-LABEL: Op post-order erasures (skip)
// CHECK:       Erasing op 'test.op0'
// CHECK:       Erasing op 'cf.br'
// CHECK:       Erasing op 'test.op1'
// CHECK:       Erasing op 'cf.br'
// CHECK:       Erasing op 'test.op2'
// CHECK:       Erasing op 'test.region.op0'
// CHECK:       Erasing op 'func.return'

// CHECK-LABEL: Block post-order erasures (skip)
// CHECK:       Erasing block ^bb0 from region 0 from operation 'test.region.op0'
// CHECK:       Erasing block ^bb0 from region 0 from operation 'test.region.op0'
// CHECK:       Erasing block ^bb0 from region 0 from operation 'test.region.op0'

// CHECK-LABEL: Op post-order erasures (no skip)
// CHECK:       Erasing op 'test.op0'
// CHECK:       Erasing op 'cf.br'
// CHECK:       Erasing op 'test.op1'
// CHECK:       Erasing op 'cf.br'
// CHECK:       Erasing op 'test.op2'
// CHECK:       Erasing op 'test.region.op0'
// CHECK:       Erasing op 'func.return'

// CHECK-LABEL: Block post-order erasures (no skip)
// CHECK:       Erasing block ^bb0 from region 0 from operation 'test.region.op0'
// CHECK:       Erasing block ^bb0 from region 0 from operation 'test.region.op0'
// CHECK:       Erasing block ^bb0 from region 0 from operation 'test.region.op0'
// CHECK:       Erasing block ^bb0 from region 0 from operation 'func.func'
// CHECK:       Erasing block ^bb0 from region 0 from operation 'builtin.module'

// -----

func.func @unordered_cfg_with_loop() {
  "test.region.op0"() ({
    ^bb0:
      %c = "test.op0"() : () -> (i1)
      cf.cond_br %c, ^bb2, ^bb3
    ^bb1:
      "test.op1"(%val) : (i32) -> ()
      cf.br ^bb5
    ^bb2:
      %val = "test.op2"() : () -> (i32)
      cf.br ^bb1
    ^bb3:
      "test.op3"() : () -> ()
      cf.br ^bb2
    ^bb4:
      "test.op4"() : () -> ()
      cf.br ^bb2
    ^bb5:
      "test.op5"() : () -> ()
      cf.br ^bb7
    ^bb6:
      "test.op6"() : () -> ()
      cf.br ^bb6
    ^bb7:
      "test.op7"() : () -> ()
  }) : () -> ()
  return
}

//      4
//      |
//      v
// 0 -> 2 --> 1 --> 5 --> 7
// |    ^
// |    |      6 --
// |   /       ^   \
// |  /         \  /
// v /           --
//  3

// CHECK-LABEL: Op forward dominance post-order visits
// CHECK:       Visiting op 'test.op0'
// CHECK:       Visiting op 'cf.cond_br'
// CHECK:       Visiting op 'test.op2'
// CHECK:       Visiting op 'cf.br'
// CHECK:       Visiting op 'test.op1'
// CHECK:       Visiting op 'cf.br'
// CHECK:       Visiting op 'test.op5'
// CHECK:       Visiting op 'cf.br'
// CHECK:       Visiting op 'test.op7'
// CHECK:       Visiting op 'test.op3'
// CHECK:       Visiting op 'cf.br'
// CHECK-NOT:   Visiting op 'test.op6'
// CHECK:       Visiting op 'test.region.op0'
// CHECK:       Visiting op 'func.return'
// CHECK:       Visiting op 'func.func'

// CHECK-LABEL: Block forward dominance post-order visits
// CHECK:       Visiting block ^bb0 from region 0 from operation 'test.region.op0'
// CHECK:       Visiting block ^bb2 from region 0 from operation 'test.region.op0'
// CHECK:       Visiting block ^bb1 from region 0 from operation 'test.region.op0'
// CHECK:       Visiting block ^bb5 from region 0 from operation 'test.region.op0'
// CHECK:       Visiting block ^bb7 from region 0 from operation 'test.region.op0'
// CHECK:       Visiting block ^bb3 from region 0 from operation 'test.region.op0'
// CHECK:       Visiting block ^bb0 from region 0 from operation 'func.func'

// CHECK-LABEL: Region forward dominance post-order visits
// CHECK:       Visiting region 0 from operation 'test.region.op0'
// CHECK:       Visiting region 0 from operation 'func.func'

// CHECK-LABEL: Op reverse dominance post-order visits
// CHECK:       Visiting op 'func.return'
// CHECK-NOT:   Visiting op 'test.op6'
// CHECK:       Visiting op 'test.op7'
// CHECK:       Visiting op 'cf.br'
// CHECK:       Visiting op 'test.op5'
// CHECK:       Visiting op 'cf.br'
// CHECK:       Visiting op 'test.op1'
// CHECK:       Visiting op 'cf.br'
// CHECK:       Visiting op 'test.op2'
// CHECK:       Visiting op 'cf.br'
// CHECK:       Visiting op 'test.op3'
// CHECK:       Visiting op 'cf.cond_br'
// CHECK:       Visiting op 'test.op0'
// CHECK:       Visiting op 'test.region.op0'
// CHECK:       Visiting op 'func.func'

// CHECK-LABEL: Block reverse dominance post-order visits
// CHECK:       Visiting block ^bb7 from region 0 from operation 'test.region.op0'
// CHECK:       Visiting block ^bb5 from region 0 from operation 'test.region.op0'
// CHECK:       Visiting block ^bb1 from region 0 from operation 'test.region.op0'
// CHECK:       Visiting block ^bb2 from region 0 from operation 'test.region.op0'
// CHECK:       Visiting block ^bb3 from region 0 from operation 'test.region.op0'
// CHECK:       Visiting block ^bb0 from region 0 from operation 'test.region.op0'
// CHECK:       Visiting block ^bb0 from region 0 from operation 'func.func'

// CHECK-LABEL: Region reverse dominance post-order visits
// CHECK:       Visiting region 0 from operation 'test.region.op0'
// CHECK:       Visiting region 0 from operation 'func.func'

// CHECK-LABEL: Block pre-order erasures (skip)
// CHECK:       Erasing block ^bb0 from region 0 from operation 'test.region.op0'
// CHECK:       Cannot erase block ^bb0 from region 0 from operation 'test.region.op0', still has uses
// CHECK:       Cannot erase block ^bb1 from region 0 from operation 'test.region.op0', still has uses
// CHECK:       Erasing block ^bb2 from region 0 from operation 'test.region.op0'
// CHECK:       Erasing block ^bb2 from region 0 from operation 'test.region.op0'
// CHECK:       Cannot erase block ^bb2 from region 0 from operation 'test.region.op0', still has uses
// CHECK:       Cannot erase block ^bb3 from region 0 from operation 'test.region.op0', still has uses
// CHECK:       Cannot erase block ^bb4 from region 0 from operation 'test.region.op0', still has uses

// -----

// The following test should not crash while visiting the intra-op blocks (inside the top level
// function in this case). We are testing that the intra-block ops are erased after dropping their
// uses from ops with same parent region.
// CHECK-LABEL: func.func @test_no_skip_block_erasure
func.func @test_no_skip_block_erasure() {
  %c0 = arith.constant 0 : index
  %c3 = arith.constant 3 : index
  cf.br ^bb1
^bb1:
  %cond = arith.cmpi eq, %c0, %c3 : index
  cf.br ^bb4
^bb4:
  return
}

// -----

// Regression test for https://github.com/llvm/llvm-project/issues/182996:
// Block erasure must also drop uses of block arguments (e.g. function args)
// from sibling blocks in the same region before destroying the block.
// CHECK-LABEL: func.func @test_no_skip_block_erasure_block_args
func.func @test_no_skip_block_erasure_block_args(%arg0: i32, %arg1: i32) -> i32 {
  cf.br ^bb1
^bb1:
  %0 = arith.addi %arg0, %arg1 : i32
  return %0 : i32
}

// -----

// Regression test for https://github.com/llvm/llvm-project/issues/183511:
// testNoSkipErasureCallbacks should not crash when visiting an empty block.
// The module body block has no ops, so block->front() would previously dereference
// the ilist sentinel, causing an assertion failure.
module {}
// CHECK-LABEL: Block post-order erasures (no skip)
// CHECK-NEXT:  Erasing block ^bb0 from region 0 from operation 'builtin.module'

// -----

// Regression test for https://github.com/llvm/llvm-project/issues/116370:
// ForwardDominanceIterator<SkipGraphRegion=true> should skip graph regions
// (such as scf.forall.in_parallel's body) instead of asserting.
// CHECK-LABEL: Op forward dominance post-order visits
// CHECK: Visiting op 'scf.forall'
// CHECK-NOT: Visiting op 'tensor.parallel_insert_slice'
// CHECK: Op reverse dominance post-order visits
func.func @graph_region_skip(%fill: tensor<2xf32>, %output: tensor<2xf32>) {
  %c0 = arith.constant 0.0 : f32
  %0 = linalg.fill ins(%c0 : f32) outs(%fill : tensor<2xf32>) -> tensor<2xf32>
  %1 = scf.forall (%i) in (2) shared_outs(%arg1 = %output) -> (tensor<2xf32>) {
    %2 = tensor.extract_slice %0[%i][1][1] : tensor<2xf32> to tensor<1xf32>
    %3 = tensor.extract_slice %arg1[%i][1][1] : tensor<2xf32> to tensor<1xf32>
    scf.forall.in_parallel {
      tensor.parallel_insert_slice %3 into %arg1[%i][1][1] : tensor<1xf32> into tensor<2xf32>
    }
  }
  return
}

// -----

// Regression test for https://github.com/llvm/llvm-project/issues/205717:
// Block erasure in skip callbacks must drop uses of values defined in the block
// before erasing it.
// CHECK-LABEL: gpu.module @issue205717
module attributes {gpu.container_module, llvm.data_layout = ""} {
  gpu.module @issue205717 {
    func.func @test_skip_block_erasure_nested_uses(%arg0: i32) -> i32 {
      %c0 = arith.constant 0 : i32
      %c1 = arith.constant 1 : i32
      %cond = arith.cmpi eq, %arg0, %c0 : i32
      cf.cond_br %cond, ^bb1, ^bb2
    ^bb1:
      gpu.barrier
      cf.br ^bb3(%c1 : i32)
    ^bb2:
      gpu.barrier
      %plus = arith.addi %arg0, %c1 : i32
      cf.br ^bb3(%plus : i32)
    ^bb3(%result: i32):
      return %result : i32
    }
  }
}