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

[Han-Chung Wang]

hanhanW wrote:


> Even with the original code, I'm puzzled: reading/writing `%mem` is hoisted, but reading/writing `%sv` is not - I would have expected neither to be hoisted:
> 
> ```mlir
> func.func @example(
>     %mem: memref<?x?xf32>,
>     %lb : index, %ub : index, %step: index, %in: vector<1xf32>) {
>   %c0 = arith.constant 0 : index
>   %cst = arith.constant 0.0 : f32
>   %sv = memref.subview %mem[0, 0][1, 1][1, 1] : memref<?x?xf32> to memref<1x1xf32, strided<[?, 1]>>
> 
> 
>   scf.for %i = %lb to %ub step %step {
>       %r0 = vector.transfer_read %mem[%c0, %c0], %cst: memref<?x?xf32>, vector<1xf32>
>       %r1 = vector.transfer_read %sv[%c0, %c0], %cst: memref<1x1xf32, strided<[?, 1]>>, vector<1xf32>
> 
>       %u0 = "some_use"(%r0) : (vector<1xf32>) -> vector<1xf32>
>       %u1 = "some_use"(%r1) : (vector<1xf32>) -> vector<1xf32>
> 
>       vector.transfer_write %u0, %mem[%c0, %c0] : vector<1xf32>, memref<?x?xf32>
>       vector.transfer_write %u1, %sv[%c0, %c0] : vector<1xf32>, memref<1x1xf32, strided<[?, 1]>>
> 
>       // NOTE: Using %in - some input argument. 
>       // 1. Hoisting still happens with this uncommented
>       // vector.transfer_write %in, %sv[%c0, %c0] : vector<1xf32>, memref<1x1xf32, strided<[?, 1]>>
>       // 2. Hoisting does not happen with this uncommented
>       // vector.transfer_write %in, %mem[%c0, %c0] : vector<1xf32>, memref<?x?xf32>
>   }
>   "unrelated_use"(%mem) : (memref<?x?xf32>) -> ()
>   return
> }
> 
> module attributes {transform.with_named_sequence} {
>   transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
>     %0 = transform.structured.match ops{["func.func"]} in %arg1
>       : (!transform.any_op) -> !transform.any_op
>     transform.structured.hoist_redundant_vector_transfers %0
>       : (!transform.any_op) -> !transform.any_op
>     transform.yield
>   }
> }
> ```
> 
> And this is without even introducing `memref.assume_alignment`. So yes - maybe it's just me being slow on a Friday, but I think something isn't adding up here.

It looks like a bug. It's been a long time since the last time I looked at the code, so I'm quite outdated in this area.

> > The intention of memref.assume_alignment is allowing the compiler to generate more efficient code based on the alignment assumption; i
> 
> Right —-but aren't we missing some kind of explicit “assume” mechanism here?
> 
> Diego raised something similar (though in the context of vector masking):
> 
> * [[mlir] How to best avoid masking in this case? #143920 (comment)](https://github.com/llvm/llvm-project/issues/143920#issuecomment-2968376118)
> 
> That’s a topic I’m actively exploring, though it'll take time - and it may or may not translate cleanly to the `memref` case.

I can't translate the idea to the `memref` case easily. I think they are slightly different. My first thought about your topic is more about [integer range analysis](https://github.com/llvm/llvm-project/blob/main/mlir/lib/Analysis/DataFlow/IntegerRangeAnalysis.cpp); if you go with `vector.assume`, you may hit issues like this when people make it have return values. 

> My overall feeling is this:
> 
> * The change in [PR #139521](https://github.com/llvm/llvm-project/pull/139521) seems to have introduced a real perf regression. We should consider reverting and re-evaluating the design.

Reverting the change could be a pain, since I already landed some improvements in upstream: https://github.com/llvm/llvm-project/pull/142425 https://github.com/llvm/llvm-project/pull/142358

and downstream projects: https://github.com/iree-org/iree/pull/20913 https://github.com/iree-org/iree/pull/20984 https://github.com/iree-org/iree/pull/20973 https://github.com/iree-org/iree/pull/20925

> * Before reverting, though, it would be good to have minimal repros. I’ve tried creating some (see above), but I’m still a bit stuck.

My repro is here, just in case if it is helpful: https://github.com/iree-org/iree/issues/20912#issuecomment-2931544107 It is not IREE specific, you can run `mlir-opt -transform-interpreter -canonicalize --split-input-file --allow-unregistered-dialect repro.mlir`.
 
> In general, I don't see an obvious fix - and it feels like `memref.assume_alignment` shouldn’t cause this kind of trouble.
> 
> For context, this isn’t really my area of deep expertise, and I’ve only skimmed through the original PR. That said, happy to keep digging and discussing!

So many thanks for spending your time here! I added you as a reviewer because the transformation is not updated for a while, and you're one of the most recent contributors in terms of this file. I don't have an obvious fix either. Can we add an option for now? E.g., we pass `bool allowAssumeAlignmentAlias` and we only perform the checks when it is true. The default value is false, so it'd work as expected.

> But it seems the line 340 and 347 are tring to slove this alias problem:
[Line340](https://github.com/llvm/llvm-project/blob/d8db47ac5a620bd2d755cc432c84b3a50a104e24/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp#L340)
          if (!vector::isDisjointTransferSet( cast<VectorTransferOpInterface>(*transferWrite), cast<VectorTransferOpInterface>(*transferWriteUse), /*testDynamicValueUsingBounds=*/true))
maybe we should enhance the isDisjointTransferSet for which have viewlikeOp source.

I think this is a little different. IIUC, it is mainly for checking whether there are overlapping between read and writes? It may be the place where we inject alias analysis checks though. We may use simple analysis that only allows assume_alignment as a start, and improve it when there are other needs.

-----

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.

I'd like to go with (1) and can help with (2), like providing guidance. Note that the issue has been there for ~1 month, and I only work on these fixes in my 20% time. It'd take more time if we want to go with (2) directly, while I'd like to re-enable tests in our downstream projects. Our cooperative matrix support on SPIR-V side is disabled for a while if assume_alignment feature is requested. It is only enabled without assume_alignment feature.

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