[Mlir-commits] [mlir] [mlir][vector] Keep gather alignment when unrolling and lowering vector.mask (PR #221382)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Hung-Kuan Tseng (Tim096)

<details>
<summary>Changes</summary>

`UnrollGatherPattern` rebuilds a `vector.gather` per tile from the same
base and offsets, slicing only the index, mask and pass-through vectors,
and `MaskedGatherOpPattern` rebuilds the gather inside a `vector.mask`
with the mask of the region. Neither forwards the `alignment` attribute:

```mlir
%0 = vector.mask %m { vector.gather %base[%c0] [%idx], %all_true, %pt alignment = 8 : memref<64xf32>, vector<4xindex>, vector<4xi1>, vector<4xf32> into vector<4xf32> } : vector<4xi1> -> vector<4xf32>
// -lower-vector-mask today
%0 = vector.gather %base[%c0] [%idx], %m, %cst : memref<64xf32>, vector<4xindex>, vector<4xi1>, vector<4xf32> into vector<4xf32>
```

The rebuilt gathers access the same addresses, so the alignment still
holds. This passes it through the builder, the same way the gather
unrolling in `LowerVectorGather` already does.

The bufferization of `vector.gather`/`vector.scatter` rebuilds them too,
but the verifier rejects `alignment` on a tensor base, so there is
nothing to forward there.


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


4 Files Affected:

- (modified) mlir/lib/Dialect/Vector/Transforms/LowerVectorMask.cpp (+1-1) 
- (modified) mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp (+1-1) 
- (modified) mlir/test/Dialect/Vector/lower-vector-mask.mlir (+12) 
- (modified) mlir/test/Dialect/Vector/vector-transfer-unroll.mlir (+15) 


``````````diff
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorMask.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorMask.cpp
index 7730c4e7c950a..3d631efc12cfa 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorMask.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorMask.cpp
@@ -270,7 +270,7 @@ struct MaskedGatherOpPattern : public MaskOpRewritePattern<GatherOp> {
     rewriter.replaceOpWithNewOp<GatherOp>(
         maskingOp.getOperation(), gatherOp.getVectorType(), gatherOp.getBase(),
         gatherOp.getOffsets(), gatherOp.getIndices(), maskingOp.getMask(),
-        passthru);
+        passthru, gatherOp.getMaybeAlign());
     return success();
   }
 };
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index a95bfba0814bd..7b908fefa21de 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -693,7 +693,7 @@ struct UnrollGatherPattern : public OpRewritePattern<vector::GatherOp> {
               strides);
       auto slicedGather = vector::GatherOp::create(
           rewriter, loc, targetType, gatherOp.getBase(), gatherOp.getOffsets(),
-          indexSubVec, maskSubVec, passThruSubVec);
+          indexSubVec, maskSubVec, passThruSubVec, gatherOp.getMaybeAlign());
 
       result = rewriter.createOrFold<vector::InsertStridedSliceOp>(
           loc, slicedGather, result, elementOffsets, strides);
diff --git a/mlir/test/Dialect/Vector/lower-vector-mask.mlir b/mlir/test/Dialect/Vector/lower-vector-mask.mlir
index a8a1164e2f762..fd6f3eea49f04 100644
--- a/mlir/test/Dialect/Vector/lower-vector-mask.mlir
+++ b/mlir/test/Dialect/Vector/lower-vector-mask.mlir
@@ -88,3 +88,15 @@ func.func @empty_vector_mask_with_return(%a : vector<8xf32>, %mask : vector<8xi1
   return %0 : vector<8xf32>
 }
 
+
+// -----
+
+// CHECK-LABEL:   func.func @vector_gather_alignment(
+//  CHECK-SAME:     %[[BASE:.*]]: memref<64xf32>, %[[IDX:.*]]: vector<4xindex>, %[[MASK:.*]]: vector<4xi1>,
+//       CHECK:     vector.gather %[[BASE]][%{{.*}}] [%[[IDX]]], %[[MASK]], %{{.*}} alignment = 8 : memref<64xf32>, vector<4xindex>, vector<4xi1>, vector<4xf32> into vector<4xf32>
+func.func @vector_gather_alignment(%base: memref<64xf32>, %idx: vector<4xindex>, %mask: vector<4xi1>, %pass_thru: vector<4xf32>) -> vector<4xf32> {
+  %c0 = arith.constant 0 : index
+  %all_true = arith.constant dense<true> : vector<4xi1>
+  %0 = vector.mask %mask { vector.gather %base[%c0] [%idx], %all_true, %pass_thru alignment = 8 : memref<64xf32>, vector<4xindex>, vector<4xi1>, vector<4xf32> into vector<4xf32> } : vector<4xi1> -> vector<4xf32>
+  return %0 : vector<4xf32>
+}
diff --git a/mlir/test/Dialect/Vector/vector-transfer-unroll.mlir b/mlir/test/Dialect/Vector/vector-transfer-unroll.mlir
index bfa76ecaf77f7..53f30ef6e7f35 100644
--- a/mlir/test/Dialect/Vector/vector-transfer-unroll.mlir
+++ b/mlir/test/Dialect/Vector/vector-transfer-unroll.mlir
@@ -416,3 +416,18 @@ func.func @transfer_write_unroll_dynamic_index(%mem : memref<4x4xf32>, %vec : ve
   vector.transfer_write %vec, %mem[%idx, %idx] : vector<4x2xf32>, memref<4x4xf32>
   return
 }
+
+// -----
+
+// ALL-LABEL: func @vector_gather_unroll_alignment(
+//       ALL:   vector.gather {{.*}} alignment = 8 : memref<?x?x?xf32>, vector<2x2xindex>, vector<2x2xi1>, vector<2x2xf32> into vector<2x2xf32>
+//       ALL:   vector.gather {{.*}} alignment = 8 : memref<?x?x?xf32>, vector<2x2xindex>, vector<2x2xi1>, vector<2x2xf32> into vector<2x2xf32>
+//   ALL-NOT:   vector.gather
+func.func @vector_gather_unroll_alignment(%mem : memref<?x?x?xf32>,
+                                          %indices : vector<2x4xindex>,
+                                          %mask : vector<2x4xi1>,
+                                          %pass_thru : vector<2x4xf32>) -> vector<2x4xf32> {
+  %c0 = arith.constant 0 : index
+  %res = vector.gather %mem[%c0, %c0, %c0] [%indices], %mask, %pass_thru alignment = 8 : memref<?x?x?xf32>, vector<2x4xindex>, vector<2x4xi1>, vector<2x4xf32> into vector<2x4xf32>
+  return %res : vector<2x4xf32>
+}

``````````

</details>


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