diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp index e9acda657b3c2..215e1b1b87452 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>(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 361919c47b097..076760fe21dc8 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 // -----