[Mlir-commits] [mlir] [mlir][vector] Hoist transfer pairs when the source is AssumeAlignmentOp (PR #144809)

[Andrzej Warzyński]

banach-space wrote:

> So I think there are two approaches:
> 
> 1. Introduce an option and only apply the check when it is on.
> 2. Inject basic alias and use it in isDisjointTransferSet checks -- I don't have direct prototype, so I'm not sure if it solves all > the problem or not.

Re (1), ideally, I would prefer to avoid that (if possible). Lets briefly chat about alternatives.

Re (2), as discussed in this [comment](https://github.com/llvm/llvm-project/pull/145235#discussion_r2166274880), that hook is to look at two Xfer ops operating on the same `base`.  Adding support for different `base`s would require growing that in a non-trivial way. IMO, we'd need some dedicated logic elsewhere (e.g. in the `MemRef` dialect). But we don't necessarily need it _yet_.

In the meantime, there's another option:
3.  Look at how `memref.assume_alignment` is used and add basic data-flow logic around that.

Looking at a specific example (taken from https://godbolt.org/z/Tj19noof1, which was shared in https://github.com/llvm/llvm-project/issues/144825)
```mlir
 func.func @hoist_vector_transfer_read_write() {
    %c0 = arith.constant 0 : index
    %c64 = arith.constant 64 : index
    %c256 = arith.constant 256 : index
    %c4096 = arith.constant 4096 : index
    %cst = arith.constant 0.000000e+00 : f16
    %alloc = memref.alloc() : memref<4096x4096xf16>
    %alloc_0 = memref.alloc() : memref<4096x4096xf16>
    %assume_align = memref.assume_alignment %alloc, 64 : memref<4096x4096xf16>
    %assume_align_1 = memref.assume_alignment %alloc_0, 64 : memref<4096x4096xf16>
    scf.for %arg0 = %c256 to %c4096 step %c256 {
      %0 = vector.transfer_read %assume_align[%c0, %c0], %cst {in_bounds = [true, true]} : memref<4096x4096xf16>, vector<16x16xf16>
      %1 = vector.transfer_read %alloc_0[%arg0, %arg0], %cst {in_bounds = [true, true]} : memref<4096x4096xf16>, vector<16x16xf16>
      %2 = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %1, %1, %0 : vector<16x16xf16>, vector<16x16xf16> into vector<16x16xf16>
      vector.transfer_write %2, %assume_align[%c0, %c0] {in_bounds = [true, true]} : vector<16x16xf16>, memref<4096x4096xf16>
    }
    return
  }
```

Note that:
* `%assume_align = memref.assume_alignment %alloc` has two users (the Xfer pair that we want to hoist)
* The defining op for `%alloc` (the argument to `memref.assume_alignment`) has only one use (`memref.assume_alignment`).

In this case, it is safe to hoist. Would this be sufficient for you? I implemented this logic [here](https://github.com/llvm/llvm-project/pull/145235/commits/36bbbdc464df57e43d15f5046c35a799e705eb0f). WDYT?

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