[Mlir-commits] [mlir] [mlir][vector] Update tests for xfer permutation lowering (3/N) (PR #127320)
Andrzej Warzyński
llvmlistbot at llvm.org
Sat Feb 15 05:05:12 PST 2025
https://github.com/banach-space created https://github.com/llvm/llvm-project/pull/127320
* Remove `vector.create_mask` from tests. Instead, pass masks as
arguments. This simplifies the tests without sacrificing test
coverage.
* Update `@xfer_read_minor_identity_tranposed_with_mask_scalable` to use
similar shapes as other tests and to avoid using test Ops (e.g.
`@test.some_use`). This improves consistency between tests.
* Fix some comment typos.
>From 036a005f500cf3fa0cdc1499cf9bf77b064553e4 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Sat, 15 Feb 2025 12:59:53 +0000
Subject: [PATCH] [mlir][vector] Update tests for xfer permutation lowering
(3/N)
* Remove `vector.create_mask` from tests. Instead, pass masks as
arguments. This simplifies the tests without sacrificing test
coverage.
* Update `@xfer_read_minor_identity_tranposed_with_mask_scalable` to use
similar shapes as other tests and to avoid using test Ops (e.g.
`@test.some_use`). This improves consistency between tests.
* Fix some comment typos.
---
.../vector-transfer-permutation-lowering.mlir | 75 +++++++++----------
1 file changed, 36 insertions(+), 39 deletions(-)
diff --git a/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir b/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir
index 6705905633e0f..7029c9b15de8c 100644
--- a/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir
+++ b/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir
@@ -106,8 +106,8 @@ func.func @xfer_write_minor_identity_transposed_map_masked(
/// (neither a minor identity nor transposed minor identity map)
/// OUT 1: vector.broadcast + vector.transfer_write
/// (transposed minor identity)
-/// OUT 2: vector.transfer_write -> vector.broadcast + vector.transpose + vector.transfer_write
-/// (minor identity)
+/// OUT 2: vector.transfer_write -> vector.broadcast + vector.transpose
+/// + vector.transfer_write (minor identity)
///----------------------------------------------------------------------------------------
// CHECK-LABEL: func.func @xfer_write_non_minor_identity(
@@ -233,16 +233,16 @@ func.func @xfer_write_non_minor_identity_masked_scalable(
// CHECK-LABEL: func @xfer_write_non_minor_identity_masked_2
// CHECK-SAME: %[[DEST:.*]]: tensor<?x?x?x?xf32>
// CHECK-SAME: %[[VEC:.*]]: vector<14x8x16xf32>
-// CHECK-SAME: %[[DIM:.*]]: index, %[[IDX:.*]]: index) -> tensor<?x?x?x?xf32>
+// CHECK-SAME: %[[MASK:.*]]: vector<14x8x16xi1>
+// CHECK-SAME: %[[DIM:.*]]: index
// CHECK-NOT: vector.broadcast
-// CHECK: vector.mask %0 { vector.transfer_write %[[VEC]], %[[DEST]]{{.*}} : vector<14x8x16xf32>, tensor<?x?x?x?xf32> } : vector<14x8x16xi1> -> tensor<?x?x?x?xf32>
+// CHECK: vector.mask %[[MASK]] { vector.transfer_write %[[VEC]], %[[DEST]]{{.*}} : vector<14x8x16xf32>, tensor<?x?x?x?xf32> } : vector<14x8x16xi1> -> tensor<?x?x?x?xf32>
func.func @xfer_write_non_minor_identity_masked_2(
%dest : tensor<?x?x?x?xf32>,
%vec : vector<14x8x16xf32>,
- %dim : index,
+ %mask: vector<14x8x16xi1>,
%idx: index) -> tensor<?x?x?x?xf32> {
- %mask = vector.create_mask %dim, %dim, %dim : vector<14x8x16xi1>
%res = vector.mask %mask {
vector.transfer_write %vec, %dest[%idx, %idx, %idx, %idx] {
in_bounds = [false, false, true],
@@ -259,29 +259,27 @@ func.func @xfer_write_non_minor_identity_masked_2(
///
/// IN: vector.transfer_read
/// (_transposed_ minor identity permutation map, with 0 or more broadcast dims)
-/// OUT: vector.transpose + vector.transfer_write
+/// OUT: vector.transfer_read + vector.broadcast + vector.transpose
/// (minor identity permutation map with 0 or more leading broadcast dims)
///----------------------------------------------------------------------------------------
/// TODO: Inner broadcast dim - see also the block at the bottom of this file
-// CHECK-LABEL: func.func @xfer_read_minor_identity_tranposed_with_mask
+// CHECK-LABEL: func.func @xfer_read_minor_identity_transposed_with_mask
// CHECK-SAME: %[[MEM:.*]]: memref<?x?xf32>,
-// CHECK-SAME: %[[DIM_1:.*]]: index, %[[DIM_2:.*]]: index, %[[IDX:.*]]: index) -> vector<8x4x2xf32> {
+// CHECK-SAME: %[[MASK:.*]]: vector<2x4xi1>
+// CHECK-SAME: %[[IDX:.*]]: index
// CHECK: %[[PASS_THROUGH:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK: %[[MASK:.*]] = vector.create_mask %[[DIM_2]], %[[DIM_1]] : vector<2x4xi1>
// CHECK: %[[T_READ:.*]] = vector.transfer_read %[[MEM]]{{\[}}%[[IDX]], %[[IDX]]], %[[PASS_THROUGH]], %[[MASK]] {in_bounds = [true, true]} : memref<?x?xf32>, vector<2x4xf32>
// CHECK: %[[BCAST:.*]] = vector.broadcast %[[T_READ]] : vector<2x4xf32> to vector<8x2x4xf32>
// CHECK: %[[TRANSPOSE:.*]] = vector.transpose %[[BCAST]], [0, 2, 1] : vector<8x2x4xf32> to vector<8x4x2xf32>
// CHECK: return %[[TRANSPOSE]] : vector<8x4x2xf32>
-func.func @xfer_read_minor_identity_tranposed_with_mask(
+func.func @xfer_read_minor_identity_transposed_with_mask(
%mem: memref<?x?xf32>,
- %dim_1: index,
- %dim_2: index,
+ %mask: vector<2x4xi1>,
%idx: index) -> (vector<8x4x2xf32>) {
%pad = arith.constant 0.000000e+00 : f32
- %mask = vector.create_mask %dim_2, %dim_1 : vector<2x4xi1>
%res = vector.transfer_read %mem[%idx, %idx], %pad, %mask {
in_bounds = [true, true, true],
permutation_map = affine_map<(d0, d1) -> (0, d1, d0)>
@@ -290,24 +288,22 @@ func.func @xfer_read_minor_identity_tranposed_with_mask(
return %res : vector<8x4x2xf32>
}
-// CHECK-LABEL: func.func @xfer_read_minor_identity_tranposed_with_mask_scalable(
+// CHECK-LABEL: func.func @xfer_read_minor_identity_transposed_with_mask_scalable(
// CHECK-SAME: %[[MEM:.*]]: memref<?x?xf32>,
-// CHECK-SAME: %[[DIM_1:.*]]: index, %[[DIM_2:.*]]: index, %[[IDX:.*]]: index) -> vector<8x[4]x2xf32> {
+// CHECK-SAME: %[[MASK:.*]]: vector<2x[4]xi1>
+// CHECK-SAME: %[[IDX:.*]]: index
// CHECK: %[[PAD:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK: %[[MASK:.*]] = vector.create_mask %[[DIM_2]], %[[DIM_1]] : vector<2x[4]xi1>
// CHECK: %[[T_READ:.*]] = vector.transfer_read %[[MEM]]{{\[}}%[[IDX]], %[[IDX]]], %[[PAD]], %[[MASK]] {in_bounds = [true, true]} : memref<?x?xf32>, vector<2x[4]xf32>
// CHECK: %[[BCAST:.*]] = vector.broadcast %[[T_READ]] : vector<2x[4]xf32> to vector<8x2x[4]xf32>
// CHECK: %[[TRANSPOSE:.*]] = vector.transpose %[[BCAST]], [0, 2, 1] : vector<8x2x[4]xf32> to vector<8x[4]x2xf32>
// CHECK: return %[[TRANSPOSE]] : vector<8x[4]x2xf32>
-func.func @xfer_read_minor_identity_tranposed_with_mask_scalable(
+func.func @xfer_read_minor_identity_transposed_with_mask_scalable(
%mem: memref<?x?xf32>,
- %dim_1: index,
- %dim_2: index,
+ %mask: vector<2x[4]xi1>,
%idx: index) -> (vector<8x[4]x2xf32>) {
%pad = arith.constant 0.000000e+00 : f32
- %mask = vector.create_mask %dim_2, %dim_1 : vector<2x[4]xi1>
%res = vector.transfer_read %mem[%idx, %idx], %pad, %mask {
in_bounds = [true, true, true],
permutation_map = affine_map<(d0, d1) -> (0, d1, d0)>
@@ -319,24 +315,26 @@ func.func @xfer_read_minor_identity_tranposed_with_mask_scalable(
// Masked version is not supported
// CHECK-LABEL: func @xfer_read_minor_identity_transposed_masked(
-// CHECK-SAME: %[[DEST:.*]]: tensor<?x1xf32>,
-// CHECK-SAME: %[[MASK:.*]]: vector<4x1xi1>
+// CHECK-SAME: %[[DEST:.*]]: tensor<?x?xf32>,
+// CHECK-SAME: %[[MASK:.*]]: vector<2x4xi1>
+// CHECK-SAME: %[[IDX:.*]]: index
// CHECK-NOT: vector.transpose
-// CHECK: vector.mask %[[MASK]] { vector.transfer_read %[[DEST]]{{.*}}: tensor<?x1xf32>, vector<1x4x4xf32> } : vector<4x1xi1> -> vector<1x4x4xf32>
+// CHECK: vector.mask %[[MASK]] { vector.transfer_read %[[DEST]]{{.*}}: tensor<?x?xf32>, vector<8x4x2xf32> } : vector<2x4xi1> -> vector<8x4x2xf32>
func.func @xfer_read_minor_identity_transposed_masked(
- %dest: tensor<?x1xf32>,
- %mask : vector<4x1xi1>,
- %idx: index) {
+ %dest: tensor<?x?xf32>,
+ %mask: vector<2x4xi1>,
+ %idx: index) -> (vector<8x4x2xf32>) {
%pad = arith.constant 0.000000e+00 : f32
- %3 = vector.mask %mask {
- vector.transfer_read %dest[%idx, %idx], %pad {
- permutation_map = affine_map<(d0, d1) -> (d1, 0, d0)>
- } : tensor<?x1xf32>, vector<1x4x4xf32>
- } : vector<4x1xi1> -> vector<1x4x4xf32>
- "test.some_use"(%3) : (vector<1x4x4xf32>) -> ()
- return
+ %res = vector.mask %mask {
+ vector.transfer_read %dest[%idx, %idx], %pad, %mask {
+ in_bounds = [true, true, true],
+ permutation_map = affine_map<(d0, d1) -> (0, d1, d0)>
+ } : tensor<?x?xf32>, vector<8x4x2xf32>
+ } : vector<2x4xi1> -> vector<8x4x2xf32>
+
+ return %res : vector<8x4x2xf32>
}
// CHECK-LABEL: func.func @xfer_read_minor_identity_transposed_masked_scalable(
@@ -346,7 +344,7 @@ func.func @xfer_read_minor_identity_transposed_masked(
// CHECK: %[[T_READ:.*]] = vector.mask %[[MASK]] { vector.transfer_read %[[DEST]]{{.*}} : tensor<?x?xf32>, vector<8x[4]x2xf32> } : vector<2x[4]xi1> -> vector<8x[4]x2xf32>
func.func @xfer_read_minor_identity_transposed_masked_scalable(
%dest: tensor<?x?xf32>,
- %mask : vector<2x[4]xi1>,
+ %mask: vector<2x[4]xi1>,
%idx: index) -> vector<8x[4]x2xf32> {
%pad = arith.constant 0.000000e+00 : f32
@@ -388,17 +386,16 @@ func.func @xfer_read_minor_identitiy_bcast_dims_scalable(
// CHECK-LABEL: func.func @xfer_read_minor_identitiy_bcast_dims_masked
// CHECK-SAME: %[[MEM:.*]]: memref<?x?x?x?xf32>,
-// CHECK-SAME: %[[DIM:.*]]: index,
+// CHECK-SAME: %[[MASK:.*]]: vector<[4]x3xi1>
// CHECK-SAME: %[[IDX:.*]]: index) -> vector<8x[4]x2x3xf32> {
// CHECK-NOT: vector.broadcast
-// CHECK: %[[MASK:.*]] = vector.mask %0 { vector.transfer_read %[[MEM]]{{.*}} : memref<?x?x?x?xf32>, vector<8x[4]x2x3xf32> } : vector<[4]x3xi1> -> vector<8x[4]x2x3xf32>
+// CHECK: vector.mask %[[MASK]] { vector.transfer_read %[[MEM]]{{.*}} : memref<?x?x?x?xf32>, vector<8x[4]x2x3xf32> } : vector<[4]x3xi1> -> vector<8x[4]x2x3xf32>
func.func @xfer_read_minor_identitiy_bcast_dims_masked(
%mem: memref<?x?x?x?xf32>,
- %dim: index,
+ %mask: vector<[4]x3xi1>,
%idx: index) -> vector<8x[4]x2x3xf32> {
%pad = arith.constant 0.000000e+00 : f32
- %mask = vector.create_mask %dim, %dim: vector<[4]x3xi1>
%res = vector.mask %mask {
vector.transfer_read %mem[%idx, %idx, %idx, %idx], %pad {
More information about the Mlir-commits
mailing list