[Mlir-commits] [mlir] [MLIR][LLVM] Block address support (PR #134335)

[Bruno Cardoso Lopes]

bcardosolopes wrote:

> What happens if MLIR merges two block (e.g. when running region simplify) with different tags? I guess this could still be exported but the LLVM IR before and after MLIR would be semantically different since two block addresses would be merged into one block address. Does LLVM proper also merge block addresses if two blocks can be merged?

Given the snippet:

```
llvm.func @fn(%cond : i1, %arg0 : i32, %arg1 : i32) -> i32 {
  llvm.cond_br %cond, ^bb1, ^bb2
^bb1:
  llvm.return %arg0 : i32
^bb2:
  llvm.return %arg1 : i32
}
````

Region simplify works nicely:

```
$ mlir-opt /tmp/x.mlir -pass-pipeline='builtin.module(llvm.func(canonicalize{region-simplify=aggressive}))'
module {
  llvm.func @fn(%arg0: i1, %arg1: i32, %arg2: i32) -> i32 {
    llvm.cond_br %arg0, ^bb1(%arg1 : i32), ^bb1(%arg2 : i32)
  ^bb1(%0: i32):  // 2 preds: ^bb0, ^bb0
    llvm.return %0 : i32
  }
}
```

If I introduce the blocktags, the opt is blocked:

```
$ mlir-opt /tmp/x.mlir  -pass-pipeline='builtin.module(llvm.func(canonicalize{region-simplify=aggressive}))'
module {
  llvm.mlir.global private @g() {addr_space = 0 : i32, dso_local} : !llvm.ptr {
    %0 = llvm.blockaddress <function = @fn, tag = <id = 0>> : !llvm.ptr
    llvm.return %0 : !llvm.ptr
  }
  llvm.mlir.global private @h() {addr_space = 0 : i32, dso_local} : !llvm.ptr {
    %0 = llvm.blockaddress <function = @fn, tag = <id = 1>> : !llvm.ptr
    llvm.return %0 : !llvm.ptr
  }
  llvm.func @fn(%arg0: i1, %arg1: i32, %arg2: i32) -> i32 {
    llvm.cond_br %arg0, ^bb1, ^bb2
  ^bb1:  // pred: ^bb0
    llvm.blocktag <id = 0>
    llvm.return %arg1 : i32
  ^bb2:  // pred: ^bb0
    llvm.blocktag <id = 1>
    llvm.return %arg2 : i32
  }
}
```

Seems like the optimization is already conservative enough? LLVM proper is slightly more smart but it also "blocks" the full optimization in face of blockaddresses: https://godbolt.org/z/TP4K35zo8 (comment the globals to double check)

I would like to also write a test for even more simple block merging (something like https://godbolt.org/z/c5o78o6Y9), but canonicalize can't seem to do the job, how do call `mlir-opt` to transform this:
```
llvm.func @fn() {
  llvm.br ^bb1
^bb1:
  llvm.br ^bb2
^bb2:
  llvm.return
}
```
into
```
llvm.func @fn() {
  llvm.return
}
```
?

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