[Mlir-commits] [mlir] [mlir][xegpu] Allow out-of-bounds writes (PR #110811)

[Adam Siemieniuk]

https://github.com/adam-smnk updated https://github.com/llvm/llvm-project/pull/110811

>From 41bd8969e2da8493f8297e87d204d7e0c2e17d7f Mon Sep 17 00:00:00 2001
From: Adam Siemieniuk <adam.siemieniuk at intel.com>
Date: Wed, 2 Oct 2024 11:30:57 +0200
Subject: [PATCH] [mlir][xegpu] Allow out-of-bounds writes

Relaxes vector.transfer_write lowering to allow out-of-bound writes.

This aligns lowering with the current hardware specification which
does not update bytes in out-of-bound locations.
---
 .../VectorToXeGPU/VectorToXeGPU.cpp           | 11 +++----
 .../transfer-write-to-xegpu.mlir              | 29 ++++++++++++-------
 2 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 0e21e96cc3fbb9..9bfa9f5558081c 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -203,18 +203,15 @@ struct TransferWriteLowering
     if (failed(transferPreconditions(rewriter, writeOp)))
       return failure();
 
-    if (writeOp.hasOutOfBoundsDim())
-      return rewriter.notifyMatchFailure(writeOp,
-                                         "Unsupported out-of-bounds write");
     AffineMap map = writeOp.getPermutationMap();
     if (!map.isMinorIdentity())
       return rewriter.notifyMatchFailure(writeOp, "Expects identity map");
 
     VectorType vecTy = writeOp.getVectorType();
-    auto descType =
-        xegpu::TensorDescType::get(vecTy.getShape(), vecTy.getElementType(),
-                                   /*array_length=*/1, /*boundary_check=*/false,
-                                   xegpu::MemorySpace::Global);
+    auto descType = xegpu::TensorDescType::get(
+        vecTy.getShape(), vecTy.getElementType(),
+        /*array_length=*/1, /*boundary_check=*/writeOp.hasOutOfBoundsDim(),
+        xegpu::MemorySpace::Global);
     xegpu::CreateNdDescOp ndDesc = createNdDescriptor(
         rewriter, loc, descType,
         dyn_cast<TypedValue<MemRefType>>(writeOp.getSource()),
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
index 361919c47b097d..076760fe21dc86 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
@@ -66,30 +66,37 @@ func.func @store_dynamic_source(%vec: vector<8x16xf32>,
 
 // -----
 
-func.func @no_store_transposed(%vec: vector<8x16xf32>,
-    %source: memref<32x64xf32>, %offset: index) {
+func.func @store_out_of_bounds(%vec: vector<8x16xf32>,
+    %source: memref<7x64xf32>, %offset: index) {
   vector.transfer_write %vec, %source[%offset, %offset]
-    {permutation_map = affine_map<(d0, d1) -> (d1, d0)>,
-    in_bounds = [true, true]}
-    : vector<8x16xf32>, memref<32x64xf32>
+    {in_bounds = [false, true]}
+    : vector<8x16xf32>, memref<7x64xf32>
   return
 }
 
-// CHECK-LABEL: @no_store_transposed(
-// CHECK:       vector.transfer_write
+// CHECK-LABEL:   @store_out_of_bounds(
+// CHECK-SAME:  %[[VEC:.+]]: vector<8x16xf32>,
+// CHECK-SAME:  %[[SRC:.+]]: memref<7x64xf32>,
+// CHECK-SAME:  %[[OFFSET:.+]]: index
+// CHECK:       %[[DESC:.+]] = xegpu.create_nd_tdesc
+// CHECK-SAME:    %[[SRC]][%[[OFFSET]], %[[OFFSET]]]
+// CHECK-SAME:    memref<7x64xf32> -> !xegpu.tensor_desc<8x16xf32,
+// CHECK-SAME:    boundary_check = true
+// CHECK:       xegpu.store_nd %[[VEC]], %[[DESC]] : vector<8x16xf32>
 
 // -----
 
-func.func @no_store_out_of_bounds(%vec: vector<8x16xf32>,
+func.func @no_store_transposed(%vec: vector<8x16xf32>,
     %source: memref<32x64xf32>, %offset: index) {
   vector.transfer_write %vec, %source[%offset, %offset]
-    {in_bounds = [false, true]}
+    {permutation_map = affine_map<(d0, d1) -> (d1, d0)>,
+    in_bounds = [true, true]}
     : vector<8x16xf32>, memref<32x64xf32>
   return
 }
 
-// CHECK-LABEL:   @no_store_out_of_bounds(
-// CHECK:         vector.transfer_write
+// CHECK-LABEL: @no_store_transposed(
+// CHECK:       vector.transfer_write
 
 // -----