[MLIR][XeVM] Update XeVM prefetch ops.

[silee2]

xevm.blockprefetch2d op has pointer operand marked as MemRead.
And that causes the op got get folded away be canonicalize pass.
Remove the side effect mark and update XeGPU to XeVM prefetch op conversion test cases to use canonicalize pass.

[llvmbot]

@llvm/pr-subscribers-mlir-gpu
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-llvm

Author: Sang Ik Lee (silee2)

Changes

Prefetch ops need pointer operand marked as MemWrite to avoid getting dead code eliminated. As a reference point, memref.prefetch is handled in a similar way.

---

Full diff: https://github.com/llvm/llvm-project/pull/166445.diff

1 Files Affected:

• (modified) mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td (+3-2)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td
index 2dd612139fa2d..91e46d68673ed 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/XeVMOps.td
@@ -444,7 +444,8 @@ def XeVM_MemfenceOp
 def XeVM_PrefetchOp
     : XeVM_Op<"prefetch">,
       Arguments<(ins Arg<AnyTypeOf<[LLVM_PointerInAddressSpace<1>,
-                                    LLVM_PointerInAddressSpace<4>]>>:$ptr,
+                                    LLVM_PointerInAddressSpace<4>]>,
+                         "", [MemWrite]>:$ptr,
           OptionalAttr<XeVM_LoadCacheControlAttr>:$cache_control)> {
   let summary = "Prefetch data into a cache subsystem.";
   let description = [{
@@ -463,7 +464,7 @@ def XeVM_PrefetchOp
 
 def XeVM_BlockPrefetch2dOp
     : XeVM_Op<"blockprefetch2d">,
-      Arguments<(ins Arg<LLVM_AnyPointer, "", [MemRead]>:$ptr, I32:$base_width,
+      Arguments<(ins Arg<LLVM_AnyPointer, "", [MemWrite]>:$ptr, I32:$base_width,
           I32:$base_height, I32:$base_pitch, I32:$x, I32:$y,
           I32Attr:$elem_size_in_bits, I32Attr:$tile_width, I32Attr:$tile_height,
           I32Attr:$v_blocks,

[adam-smnk]

Not wrong but doesn't seem strictly necessary.
Do you have any test or reproducer where the op gets folded away?

[silee2]

Not wrong but doesn't seem strictly necessary. Do you have any test or reproducer where the op gets folded away?

For example,
https://github.com/llvm/llvm-project/blob/main/mlir/test/Conversion/XeGPUToXeVM/prefetch_nd.mlir
Adding -canonicalize removes op.

Can you elaborate why it may not be strictly necessary?

[adam-smnk]

Thanks for the example, now I see the issue.

I quickly checked on an example using only xevm.prefetch which originally doesn't define any side effects.
Ops without known side effects are generally not eliminated due to conservative evaluation. Same applies to the XeGPU prefetch ops that also don't define side effects.
As such, I think assigning write side effect might be stronger "guard" against unwanted modifications as it also forces passes to be more pessimistic.

Could you add a small test case just to avoid regression in the future?

[silee2]

Thanks for the example, now I see the issue.

I quickly checked on an example using only xevm.prefetch which originally doesn't define any side effects. Ops without known side effects are generally not eliminated due to conservative evaluation. Same applies to the XeGPU prefetch ops that also don't define side effects. As such, I think assigning write side effect might be stronger "guard" against unwanted modifications as it also forces passes to be more pessimistic.

Could you add a small test case just to avoid regression in the future?

Instead of adding new tests, updated XeGPU to XeVM conversion tests to call canonicalize pass.
This change will make sure XeVM prefetch ops do not get folded away.

[silee2]

@adam-smnk Turns out you were right about conservative evaluation.
I did more experiment and only blockprefetch2d was getting folded away before this PR.
blockprefetch2d has arg ptr marked with MemRead side effect.
The op does not get folded away after removing that side effect.

[silee2]

@adam-smnk Updated PR to just remove previous MemRead side effect.
I feel the stronger assumption, MemWrite, is not needed until a usage case that shows a difference pops up.

[adam-smnk]

LGTM
Either approach sounds fine 👍

[charithaintc]

xevm.blockprefetch2d op has pointer operand marked as MemRead.
And that causes the op got get folded away be canonicalize pass.
Remove the side effect mark and update XeGPU to XeVM prefetch op conversion test cases to use canonicalize pass.

does this mean prefetch was not active at all? :-D

[charithaintc]

LGTM

[charithaintc]

so canonicalize will fold it (before)?

[silee2]

canonicalize folds constant values and simplifies index compute, hoist constant value definitions, and removes dead code in this case.

[silee2]

xevm.blockprefetch2d op has pointer operand marked as MemRead.
And that causes the op got get folded away be canonicalize pass.
Remove the side effect mark and update XeGPU to XeVM prefetch op conversion test cases to use canonicalize pass.

does this mean prefetch was not active at all? :-D

Yes, if canonicalize is called in later pass pipeline.