[Mlir-commits] [mlir] [mlir][AMX] Memory effects to amx.tile_zero (PR #155403)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-amx

Author: Arun Thangamani (arun-thmn)

<details>
<summary>Changes</summary>

Updating `amx.tile_zero` with memory effect (`MemWrite`).

`CSE` optimizes multiple `amx.tile_zero` into one `amx.tile_zero`. It is okay for `reads` but for `writes` to multiple `amx.tile_zero` the `CSE` optimized code leads to a `spill`. 

Below is the example `IR`, where four  `amx.tile_zero` get `CSE` optimized into one `amx.tile_zero` but it lowers to assembly with  three `spills`. Adding, memory effect (`MemWrite`) will avoid `CSE` if there is a write.

        tilezero        %tmm0
        tilestored      %tmm0, -2032(%rbp,%rbx) # 1024-byte Folded Spill
        tileloadd       -2032(%rbp,%rbx), %tmm3 # 1024-byte Folded Reload
        tilestored      %tmm0, -3056(%rbp,%rbx) # 1024-byte Folded Spill
        tileloadd       -3056(%rbp,%rbx), %tmm2 # 1024-byte Folded Reload
        tilestored      %tmm0, -4080(%rbp,%rbx) # 1024-byte Folded Spill
        tileloadd       -4080(%rbp,%rbx), %tmm1 # 1024-byte Folded Reload

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


1 Files Affected:

- (modified) mlir/include/mlir/Dialect/AMX/AMX.td (+3-2) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/AMX/AMX.td b/mlir/include/mlir/Dialect/AMX/AMX.td
index 6bbde43e2d011..91fe1bfb5cd35 100644
--- a/mlir/include/mlir/Dialect/AMX/AMX.td
+++ b/mlir/include/mlir/Dialect/AMX/AMX.td
@@ -142,8 +142,9 @@ class AMX_Op<string mnemonic, list<Trait> traits = []> :
 // Tile reset.
 //
 
-def TileZeroOp : AMX_Op<"tile_zero", [Pure,
-    AMXIntrinsicOpInterface
+def TileZeroOp : AMX_Op<"tile_zero", [
+    AMXIntrinsicOpInterface,
+    MemoryEffects<[MemWrite]>
   ]> {
   let summary = "tile zero operation";
   let description = [{

``````````

</details>


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