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

[llvmlistbot at llvm.org]

Author: Adam Siemieniuk
Date: 2024-10-09T18:59:14+02:00
New Revision: ec450b19004a653f3db3ad50e88fbf6529a9d841

URL: https://github.com/llvm/llvm-project/commit/ec450b19004a653f3db3ad50e88fbf6529a9d841
DIFF: https://github.com/llvm/llvm-project/commit/ec450b19004a653f3db3ad50e88fbf6529a9d841.diff

LOG: [mlir][xegpu] Allow out-of-bounds writes (#110811)

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 during block stores.

Added: 
    

Modified: 
    mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
    mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index e9acda657b3c27..215e1b1b874520 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -218,18 +218,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
 
 // -----