From 9be1956c96d9fd52af55c2658d6fca3e1b27c132 Mon Sep 17 00:00:00 2001 From: Andrzej Warzynski Date: Wed, 19 Jun 2024 07:57:59 +0100 Subject: [PATCH 1/4] [mlir][vector] Add tests xfer-permute-lowering (nfc)(2/n) Adds more tests to: * vector-transfer-permutation-lowering.mlir Specifically, adds tests for: * out-of-bounds access for the `TransferWritePermutationLowering` pattern * in-bounds access for `TransferWriteNonPermutationLowering` + `TransferWritePermutationLowering` Also renames `@permutation_with_mask_xfer_write_fixed_width` as `@xfer_write_non_transposing_permutation_map`. This is a part of a larger effort to make sure that all key cases for patterns under populateVectorTransferPermutationMapLoweringPatterns (*) are tested. I also want to make sure that tests use consistent function and variable names. (*) transform.apply_patterns.vector.transfer_permutation_patterns in TD parlance) Depends on #96031 --- .../vector-transfer-permutation-lowering.mlir | 74 +++++++++++++++---- 1 file changed, 61 insertions(+), 13 deletions(-) diff --git a/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir b/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir index b9dcb2b55e52e..01e12ead521c6 100644 --- a/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir +++ b/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir @@ -1,5 +1,7 @@ // RUN: mlir-opt %s --transform-interpreter --split-input-file | FileCheck %s +// TODO: Replace %arg0 with %vec + ///---------------------------------------------------------------------------------------- /// vector.transfer_write -> vector.transpose + vector.transfer_write /// [Pattern: TransferWritePermutationLowering] @@ -31,6 +33,27 @@ func.func @xfer_write_transposing_permutation_map( return } +// Even with out-of-bounds, it is safe to apply this pattern +// CHECK-LABEL: func.func @xfer_write_transposing_permutation_map_out_of_bounds +// CHECK-SAME: %[[ARG_0:.*]]: vector<4x8xi16>, +// CHECK-SAME: %[[MEM:.*]]: memref<2x2x?x?xi16>) { +// CHECK: %[[TR:.*]] = vector.transpose %[[ARG_0]], [1, 0] : vector<4x8xi16> to vector<8x4xi16> +// CHECK: vector.transfer_write +// CHECK-NOT: permutation_map +// CHECK-SAME: %[[TR]], %[[MEM]]{{.*}} {in_bounds = [false, false]} : vector<8x4xi16>, memref<2x2x?x?xi16> +func.func @xfer_write_transposing_permutation_map_out_of_bounds( + %arg0: vector<4x8xi16>, + %mem: memref<2x2x?x?xi16>) { + + %c0 = arith.constant 0 : index + vector.transfer_write %arg0, %mem[%c0, %c0, %c0, %c0] { + in_bounds = [false, false], + permutation_map = affine_map<(d0, d1, d2, d3) -> (d3, d2)> + } : vector<4x8xi16>, memref<2x2x?x?xi16> + + return +} + // CHECK-LABEL: func.func @xfer_write_transposing_permutation_map_with_mask_scalable // CHECK-SAME: %[[ARG_0:.*]]: vector<4x[8]xi16>, // CHECK-SAME: %[[MEM:.*]]: memref<2x2x?x4xi16>, @@ -83,19 +106,44 @@ func.func @xfer_write_transposing_permutation_map_masked( /// * vector.broadcast + vector.transpose + vector.transfer_write with a map /// which _is_ a permutation of a minor identity -// CHECK-LABEL: func @permutation_with_mask_xfer_write_fixed_width( -// CHECK: %[[vec:.*]] = arith.constant dense<-2.000000e+00> : vector<7x1xf32> -// CHECK: %[[mask:.*]] = arith.constant dense<[true, false, true, false, true, true, true]> : vector<7xi1> -// CHECK: %[[b:.*]] = vector.broadcast %[[mask]] : vector<7xi1> to vector<1x7xi1> -// CHECK: %[[tp:.*]] = vector.transpose %[[b]], [1, 0] : vector<1x7xi1> to vector<7x1xi1> -// CHECK: vector.transfer_write %[[vec]], %{{.*}}[%{{.*}}, %{{.*}}], %[[tp]] {in_bounds = [false, true]} : vector<7x1xf32>, memref -func.func @permutation_with_mask_xfer_write_fixed_width(%mem : memref, %base1 : index, - %base2 : index) { - - %fn1 = arith.constant -2.0 : f32 - %vf0 = vector.splat %fn1 : vector<7xf32> - %mask = arith.constant dense<[1, 0, 1, 0, 1, 1, 1]> : vector<7xi1> - vector.transfer_write %vf0, %mem[%base1, %base2], %mask +// CHECK-LABEL: func.func @xfer_write_non_transposing_permutation_map( +// CHECK-SAME: %[[MEM:.*]]: memref, +// CHECK-SAME: %[[VEC:.*]]: vector<7xf32>, +// CHECK-SAME: %[[BASE_1:.*]]: index, %[[BASE_2:.*]]: index) { +// CHECK: %[[BC:.*]] = vector.broadcast %[[VEC]] : vector<7xf32> to vector<1x7xf32> +// CHECK: %[[TR:.*]] = vector.transpose %[[BC]], [1, 0] : vector<1x7xf32> to vector<7x1xf32> +// CHECK: vector.transfer_write %[[TR]], %[[MEM]]{{\[}}%[[BASE_1]], %[[BASE_2]]] {in_bounds = [false, true]} : vector<7x1xf32>, memref +func.func @xfer_write_non_transposing_permutation_map( + %mem : memref, + %arg0 : vector<7xf32>, + %base1 : index, + %base2 : index) { + + vector.transfer_write %arg0, %mem[%base1, %base2] + {permutation_map = affine_map<(d0, d1) -> (d0)>} + : vector<7xf32>, memref + return +} + +// The broadcast dimension is in bounds, so the transformation is safe +// CHECK-LABEL: func.func @xfer_write_non_transposing_permutation_map_with_mask_out_of_bounds( +// CHECK-SAME: %[[MEM:.*]]: memref, +// CHECK-SAME: %[[VEC:.*]]: vector<7xf32>, +// CHECK-SAME: %[[BASE_1:.*]]: index, %[[BASE_2:.*]]: index, +// CHECK-SAME: %[[MASK:.*]]: vector<7xi1>) { +// CHECK: %[[BC_VEC:.*]] = vector.broadcast %[[VEC]] : vector<7xf32> to vector<1x7xf32> +// CHECK: %[[BC_MASK:.*]] = vector.broadcast %[[MASK]] : vector<7xi1> to vector<1x7xi1> +// CHECK: %[[TR_MASK:.*]] = vector.transpose %[[BC_MASK]], [1, 0] : vector<1x7xi1> to vector<7x1xi1> +// CHECK: %[[TR_VEC:.*]] = vector.transpose %[[BC_VEC]], [1, 0] : vector<1x7xf32> to vector<7x1xf32> +// CHECK: vector.transfer_write %[[TR_VEC]], %[[MEM]]{{\[}}%[[BASE_1]], %[[BASE_2]]], %[[TR_MASK]] {in_bounds = [false, true]} : vector<7x1xf32>, memref +func.func @xfer_write_non_transposing_permutation_map_with_mask_out_of_bounds( + %mem : memref, + %arg0 : vector<7xf32>, + %base1 : index, + %base2 : index, + %mask : vector<7xi1>) { + + vector.transfer_write %arg0, %mem[%base1, %base2], %mask {permutation_map = affine_map<(d0, d1) -> (d0)>, in_bounds = [false]} : vector<7xf32>, memref return From 96f30e90140252d977ec812ec9a4437b6d5d0639 Mon Sep 17 00:00:00 2001 From: Andrzej Warzynski Date: Wed, 31 Jul 2024 14:58:36 +0100 Subject: [PATCH 2/4] fixup! [mlir][vector] Add tests xfer-permute-lowering (nfc)(2/n) Rebase and updated based on recent changes --- .../Vector/vector-transfer-permutation-lowering.mlir | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir b/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir index 01e12ead521c6..e24f0ec836560 100644 --- a/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir +++ b/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir @@ -37,10 +37,14 @@ func.func @xfer_write_transposing_permutation_map( // CHECK-LABEL: func.func @xfer_write_transposing_permutation_map_out_of_bounds // CHECK-SAME: %[[ARG_0:.*]]: vector<4x8xi16>, // CHECK-SAME: %[[MEM:.*]]: memref<2x2x?x?xi16>) { +// CHECK: %[[C0:.*]] = arith.constant 0 : index // CHECK: %[[TR:.*]] = vector.transpose %[[ARG_0]], [1, 0] : vector<4x8xi16> to vector<8x4xi16> +// Expect the in_bounds attribute to be preserved. Since we don't print it when +// all flags are "false", it should not appear in the output. +// CHECK-NOT: in_bounds // CHECK: vector.transfer_write // CHECK-NOT: permutation_map -// CHECK-SAME: %[[TR]], %[[MEM]]{{.*}} {in_bounds = [false, false]} : vector<8x4xi16>, memref<2x2x?x?xi16> +// CHECK-SAME: %[[TR]], %[[MEM]][%[[C0]], %[[C0]], %[[C0]], %[[C0]]] : vector<8x4xi16>, memref<2x2x?x?xi16> func.func @xfer_write_transposing_permutation_map_out_of_bounds( %arg0: vector<4x8xi16>, %mem: memref<2x2x?x?xi16>) { @@ -125,7 +129,7 @@ func.func @xfer_write_non_transposing_permutation_map( return } -// The broadcast dimension is in bounds, so the transformation is safe +// Even with out-of-bounds, it is safe to apply this pattern // CHECK-LABEL: func.func @xfer_write_non_transposing_permutation_map_with_mask_out_of_bounds( // CHECK-SAME: %[[MEM:.*]]: memref, // CHECK-SAME: %[[VEC:.*]]: vector<7xf32>, From a57d9af9551343cac087877f5e59162579f40fbf Mon Sep 17 00:00:00 2001 From: Andrzej Warzynski Date: Thu, 1 Aug 2024 10:15:03 +0100 Subject: [PATCH 3/4] fixup! fixup! [mlir][vector] Add tests xfer-permute-lowering (nfc)(2/n) Apply suggestions from Hugo --- .../vector-transfer-permutation-lowering.mlir | 32 +++++++++++-------- 1 file changed, 18 insertions(+), 14 deletions(-) diff --git a/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir b/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir index e24f0ec836560..1cd455b63162b 100644 --- a/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir +++ b/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir @@ -1,6 +1,10 @@ // RUN: mlir-opt %s --transform-interpreter --split-input-file | FileCheck %s +// TODO: Replace %arg0 with %vec and +// TODO: Align naming with e.g. vector-transfer-flatten.mlir // TODO: Replace %arg0 with %vec +// TODO: Replace index args as %idx +// TODO: Align argument definition in CHECKS with function body. ///---------------------------------------------------------------------------------------- /// vector.transfer_write -> vector.transpose + vector.transfer_write @@ -14,8 +18,8 @@ /// _is_ a minor identity // CHECK-LABEL: func.func @xfer_write_transposing_permutation_map -// CHECK-SAME: %[[ARG_0:.*]]: vector<4x8xi16>, -// CHECK-SAME: %[[MEM:.*]]: memref<2x2x8x4xi16>) { +// CHECK-SAME: %[[ARG_0:.*]]: vector<4x8xi16>, +// CHECK-SAME: %[[MEM:.*]]: memref<2x2x8x4xi16>) { // CHECK: %[[TR:.*]] = vector.transpose %[[ARG_0]], [1, 0] : vector<4x8xi16> to vector<8x4xi16> // CHECK: vector.transfer_write // CHECK-NOT: permutation_map @@ -35,8 +39,8 @@ func.func @xfer_write_transposing_permutation_map( // Even with out-of-bounds, it is safe to apply this pattern // CHECK-LABEL: func.func @xfer_write_transposing_permutation_map_out_of_bounds -// CHECK-SAME: %[[ARG_0:.*]]: vector<4x8xi16>, -// CHECK-SAME: %[[MEM:.*]]: memref<2x2x?x?xi16>) { +// CHECK-SAME: %[[ARG_0:.*]]: vector<4x8xi16>, +// CHECK-SAME: %[[MEM:.*]]: memref<2x2x?x?xi16>) { // CHECK: %[[C0:.*]] = arith.constant 0 : index // CHECK: %[[TR:.*]] = vector.transpose %[[ARG_0]], [1, 0] : vector<4x8xi16> to vector<8x4xi16> // Expect the in_bounds attribute to be preserved. Since we don't print it when @@ -113,17 +117,17 @@ func.func @xfer_write_transposing_permutation_map_masked( // CHECK-LABEL: func.func @xfer_write_non_transposing_permutation_map( // CHECK-SAME: %[[MEM:.*]]: memref, // CHECK-SAME: %[[VEC:.*]]: vector<7xf32>, -// CHECK-SAME: %[[BASE_1:.*]]: index, %[[BASE_2:.*]]: index) { +// CHECK-SAME: %[[IDX_1:.*]]: index, %[[IDX_2:.*]]: index) { // CHECK: %[[BC:.*]] = vector.broadcast %[[VEC]] : vector<7xf32> to vector<1x7xf32> // CHECK: %[[TR:.*]] = vector.transpose %[[BC]], [1, 0] : vector<1x7xf32> to vector<7x1xf32> -// CHECK: vector.transfer_write %[[TR]], %[[MEM]]{{\[}}%[[BASE_1]], %[[BASE_2]]] {in_bounds = [false, true]} : vector<7x1xf32>, memref +// CHECK: vector.transfer_write %[[TR]], %[[MEM]]{{\[}}%[[IDX_1]], %[[IDX_2]]] {in_bounds = [false, true]} : vector<7x1xf32>, memref func.func @xfer_write_non_transposing_permutation_map( %mem : memref, %arg0 : vector<7xf32>, - %base1 : index, - %base2 : index) { + %idx_1 : index, + %idx_2 : index) { - vector.transfer_write %arg0, %mem[%base1, %base2] + vector.transfer_write %arg0, %mem[%idx_1, %idx_2] {permutation_map = affine_map<(d0, d1) -> (d0)>} : vector<7xf32>, memref return @@ -133,21 +137,21 @@ func.func @xfer_write_non_transposing_permutation_map( // CHECK-LABEL: func.func @xfer_write_non_transposing_permutation_map_with_mask_out_of_bounds( // CHECK-SAME: %[[MEM:.*]]: memref, // CHECK-SAME: %[[VEC:.*]]: vector<7xf32>, -// CHECK-SAME: %[[BASE_1:.*]]: index, %[[BASE_2:.*]]: index, +// CHECK-SAME: %[[IDX_1:.*]]: index, %[[IDX_2:.*]]: index, // CHECK-SAME: %[[MASK:.*]]: vector<7xi1>) { // CHECK: %[[BC_VEC:.*]] = vector.broadcast %[[VEC]] : vector<7xf32> to vector<1x7xf32> // CHECK: %[[BC_MASK:.*]] = vector.broadcast %[[MASK]] : vector<7xi1> to vector<1x7xi1> // CHECK: %[[TR_MASK:.*]] = vector.transpose %[[BC_MASK]], [1, 0] : vector<1x7xi1> to vector<7x1xi1> // CHECK: %[[TR_VEC:.*]] = vector.transpose %[[BC_VEC]], [1, 0] : vector<1x7xf32> to vector<7x1xf32> -// CHECK: vector.transfer_write %[[TR_VEC]], %[[MEM]]{{\[}}%[[BASE_1]], %[[BASE_2]]], %[[TR_MASK]] {in_bounds = [false, true]} : vector<7x1xf32>, memref +// CHECK: vector.transfer_write %[[TR_VEC]], %[[MEM]]{{\[}}%[[IDX_1]], %[[IDX_2]]], %[[TR_MASK]] {in_bounds = [false, true]} : vector<7x1xf32>, memref func.func @xfer_write_non_transposing_permutation_map_with_mask_out_of_bounds( %mem : memref, %arg0 : vector<7xf32>, - %base1 : index, - %base2 : index, + %idx_1 : index, + %idx_2 : index, %mask : vector<7xi1>) { - vector.transfer_write %arg0, %mem[%base1, %base2], %mask + vector.transfer_write %arg0, %mem[%idx_1, %idx_2], %mask {permutation_map = affine_map<(d0, d1) -> (d0)>, in_bounds = [false]} : vector<7xf32>, memref return From 242db26576088082a99b99659f392965dc7ee23e Mon Sep 17 00:00:00 2001 From: Andrzej Warzynski Date: Thu, 1 Aug 2024 10:54:00 +0100 Subject: [PATCH 4/4] fixup! fixup! fixup! [mlir][vector] Add tests xfer-permute-lowering (nfc)(2/n) Remove repeated comment --- .../Dialect/Vector/vector-transfer-permutation-lowering.mlir | 1 - 1 file changed, 1 deletion(-) diff --git a/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir b/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir index 1cd455b63162b..3ca430a92cf97 100644 --- a/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir +++ b/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir @@ -1,6 +1,5 @@ // RUN: mlir-opt %s --transform-interpreter --split-input-file | FileCheck %s -// TODO: Replace %arg0 with %vec and // TODO: Align naming with e.g. vector-transfer-flatten.mlir // TODO: Replace %arg0 with %vec // TODO: Replace index args as %idx