[Mlir-commits] [mlir] [mlir][vector] Refactor linearize.mlir (PR #86648)

Andrzej WarzyƄski llvmlistbot at llvm.org
Tue Mar 26 08:55:41 PDT 2024


https://github.com/banach-space updated https://github.com/llvm/llvm-project/pull/86648

>From bcf0eeb553e4bc03de888e4186f52d65a021a4ac Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Tue, 26 Mar 2024 09:50:40 +0000
Subject: [PATCH 1/2] [mlir][vector] Refactor linearize.mlir

This patch refactors the `linearize.mlir` test - currently it contains
some duplication and can be tricky to follow.

Summary of changes:
  * reduce duplication by introducing a shared check prefix (`ALL`) and
    by introducing `-check-prefixes`,
  * make sure that every "check" line is directly above the
    corresponding line of input MLIR,
  * group check lines corresponding to a particular prefix together (so
    that it's easier to see the expected output for a particular
    prefix),
  * remove `CHECK` from prefix names (with multiple prefixes that's just
    noise that can be avoided) and use a bit more descriptive prefixes
    instead (`CHECK0` -> `BW-0`, where `BW` stands for bitwidth),
  * unify indentation,
  * `nonvec_result` -> `test_tensor_no_linearize` (for consistency with
    `test_index_no_linearize`).

NOTE: This change only updates the format of the "CHECK" lines and
doesn't affect what's being tested.

This change is intended as preparation for adding support for scalable
vectors to `LinearizeConstant` and `LinearizeVectorizable` - i.e.
patterns that linearlize.mlir is meant to test.
---
 mlir/test/Dialect/Vector/linearize.mlir | 132 ++++++++++++------------
 1 file changed, 65 insertions(+), 67 deletions(-)

diff --git a/mlir/test/Dialect/Vector/linearize.mlir b/mlir/test/Dialect/Vector/linearize.mlir
index e865fcb7a419c5..84325a3ccd319b 100644
--- a/mlir/test/Dialect/Vector/linearize.mlir
+++ b/mlir/test/Dialect/Vector/linearize.mlir
@@ -1,92 +1,89 @@
-// RUN: mlir-opt %s -split-input-file -test-vector-linearize | FileCheck %s
-// RUN: mlir-opt %s -split-input-file -test-vector-linearize=target-vector-bitwidth=128 | FileCheck %s --check-prefix=CHECK128
-// RUN: mlir-opt %s -split-input-file -test-vector-linearize=target-vector-bitwidth=0 | FileCheck %s --check-prefix=CHECK0
-
-// CHECK-LABEL: test_linearize
-// CHECK128-LABEL: test_linearize
-// CHECK0-LABEL: test_linearize
-//  CHECK-SAME: (%[[ORIG_ARG:.*]]: vector<2x2xf32>)
-//  CHECK128-SAME: (%[[ORIG_ARG:.*]]: vector<2x2xf32>)
-//       CHECK: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x2xf32> to vector<4xf32>
-//       CHECK128: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x2xf32> to vector<4xf32>
+// RUN: mlir-opt %s -split-input-file -test-vector-linearize | FileCheck %s --check-prefixes=ALL,DEFAULT
+// RUN: mlir-opt %s -split-input-file -test-vector-linearize=target-vector-bitwidth=128 | FileCheck %s --check-prefixes=ALL,BW-128
+// RUN: mlir-opt %s -split-input-file -test-vector-linearize=target-vector-bitwidth=0 | FileCheck %s --check-prefixes=ALL,BW-0
+
+// ALL-LABEL: test_linearize
+// ALL-SAME: (%[[ORIG_ARG:.*]]: vector<2x2xf32>)
 func.func @test_linearize(%arg0: vector<2x2xf32>) -> vector<2x2xf32> {
-//       CHECK: %[[C1:.*]] = arith.constant dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00]> : vector<4xf32>
-//       CHECK128: %[[C1:.*]] = arith.constant dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00]> : vector<4xf32>
-//       CHECK0: %[[C1:.*]] = arith.constant dense<{{.*}}> : vector<2x2xf32>
+  // DEFAULT: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x2xf32> to vector<4xf32>
+  // DEFAULT: %[[CST:.*]] = arith.constant dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00]> : vector<4xf32>
+  // DEFAULT: %[[RES:.*]] = vector.shape_cast %[[CST]] : vector<4xf32> to vector<2x2xf32>
+
+  // BW-128: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x2xf32> to vector<4xf32>
+  // BW-128: %[[CST:.*]] = arith.constant dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00]> : vector<4xf32>
+  // BW-128: %[[RES:.*]] = vector.shape_cast %[[CST]] : vector<4xf32> to vector<2x2xf32>
 
+  // BW-0: %[[RES:.*]] = arith.constant dense<{{.*}}> : vector<2x2xf32>
   %0 = arith.constant dense<[[1.0, 2.0], [3.0, 4.0]]> : vector<2x2xf32>
-//       CHECK: %[[RES:.*]] = vector.shape_cast %[[C1]] : vector<4xf32> to vector<2x2xf32>
-//       CHECK128: %[[RES:.*]] = vector.shape_cast %[[C1]] : vector<4xf32> to vector<2x2xf32>
-// Arith and math ops are handled in generic way, check some of them
-//       CHECK: %{{.*}} =  math.sin %[[ARG]] : vector<4xf32>
-//       CHECK128: %{{.*}} =  math.sin %[[ARG]] : vector<4xf32>
-//       CHECK0: %{{.*}} =  math.sin %{{.*}} : vector<2x2xf32>
+
+  // DEFAULT: %{{.*}} =  math.sin %[[ARG]] : vector<4xf32>
+  // BW-128: %{{.*}} =  math.sin %[[ARG]] : vector<4xf32>
+  // BW-0: %{{.*}} =  math.sin %{{.*}} : vector<2x2xf32>
   %1 = math.sin %arg0 : vector<2x2xf32>
-//       CHECK: %{{.*}} = arith.addf %[[ARG]], %[[C1]] : vector<4xf32>
-//       CHECK128: %{{.*}} = arith.addf %[[ARG]], %[[C1]] : vector<4xf32>
-//       CHECK0: %{{.*}} = arith.addf %{{.*}} : vector<2x2xf32>
 
+  // DEFAULT: %{{.*}} = arith.addf %[[ARG]], %[[CST]] : vector<4xf32>
+  // BW-128: %{{.*}} = arith.addf %[[ARG]], %[[CST]] : vector<4xf32>
+  // BW-0: %{{.*}} = arith.addf %{{.*}} : vector<2x2xf32>
   %2 = arith.addf %arg0, %0 :  vector<2x2xf32>
 
-//       CHECK: return %[[RES]] : vector<2x2xf32>
-//       CHECK128: return %[[RES]] : vector<2x2xf32>
+  // ALL: return %[[RES]] : vector<2x2xf32>
   return %0 : vector<2x2xf32>
 }
 
-// CHECK-LABEL: test_partial_linearize
-// CHECK128-LABEL: test_partial_linearize
-// CHECK0-LABEL: test_partial_linearize
-//  CHECK-SAME: (%[[ORIG_ARG:.*]]: vector<2x2xf32>, %[[ORIG_ARG2:.*]]: vector<4x4xf32>)
-//  CHECK128-SAME: (%[[ORIG_ARG:.*]]: vector<2x2xf32>, %[[ORIG_ARG2:.*]]: vector<4x4xf32>)
-//  CHECK0-SAME: (%[[ORIG_ARG:.*]]: vector<2x2xf32>, %[[ORIG_ARG2:.*]]: vector<4x4xf32>)
-//       CHECK: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x2xf32> to vector<4xf32>
-//       CHECK128: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x2xf32> to vector<4xf32>
-//       CHECK: %[[ARG2:.*]] = vector.shape_cast %[[ORIG_ARG2]] : vector<4x4xf32> to vector<16xf32>
+// -----
+
+// ALL-LABEL: test_partial_linearize
+// ALL-SAME: (%[[ORIG_ARG:.*]]: vector<2x2xf32>, %[[ORIG_ARG2:.*]]: vector<4x4xf32>)
 func.func @test_partial_linearize(%arg0: vector<2x2xf32>, %arg1: vector<4x4xf32>) -> vector<2x2xf32> {
-//       CHECK: %[[C1:.*]] = arith.constant dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00]> : vector<4xf32>
-//       CHECK128: %[[C1:.*]] = arith.constant dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00]> : vector<4xf32>
-//       CHECK0: %[[C1:.*]] = arith.constant dense<{{.*}}> : vector<2x2xf32>
+  // DEFAULT: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x2xf32> to vector<4xf32>
+  // DEFAULT: %[[ARG2:.*]] = vector.shape_cast %[[ORIG_ARG2]] : vector<4x4xf32> to vector<16xf32>
+  // DEFAULT: %[[CST:.*]] = arith.constant dense<{{.*}}> : vector<4xf32>
+  // DEFAULT: %[[RES:.*]] = vector.shape_cast %[[CST]] : vector<4xf32> to vector<2x2xf32>
 
+  // BW-128: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x2xf32> to vector<4xf32>
+  // BW-128: %[[CST:.*]] = arith.constant dense<{{.*}}> : vector<4xf32>
+  // BW-128: %[[RES:.*]] = vector.shape_cast %[[CST]] : vector<4xf32> to vector<2x2xf32>
+
+  // BW-0: %[[RES:.*]] = arith.constant dense<{{.*}}> : vector<2x2xf32>
   %0 = arith.constant dense<[[1.0, 2.0], [3.0, 4.0]]> : vector<2x2xf32>
-//       CHECK: %[[RES:.*]] = vector.shape_cast %[[C1]] : vector<4xf32> to vector<2x2xf32>
-//       CHECK128: %[[RES:.*]] = vector.shape_cast %[[C1]] : vector<4xf32> to vector<2x2xf32>
 
-  // CHECK: %[[C2:.*]] = arith.constant dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 1.000000e+00, 2.000000e+00, 5.000000e+00, 6.000000e+00]> : vector<16xf32>
-  // CHECK128: %[[C2:.*]] = arith.constant dense<{{.*}}> : vector<4x4xf32>
-  // CHECK0: %[[C2:.*]] = arith.constant dense<{{.*}}> : vector<4x4xf32>
+  // DEFAULT: %[[C2:.*]] = arith.constant dense<{{.*}}> : vector<16xf32>
+  // BW-128: %[[C2:.*]] = arith.constant dense<{{.*}}> : vector<4x4xf32>
+  // BW-0: %[[C2:.*]] = arith.constant dense<{{.*}}> : vector<4x4xf32>
   %5 = arith.constant dense<[[1.0, 2.0, 3.0, 4.0], [1.0, 2.0,3.0, 4.0], [1.0, 2.0, 3.0, 4.0], [1.0, 2.0, 5.0, 6.0]]> : vector<4x4xf32>
-// Arith and math ops are handled in generic way, check some of them
-//       CHECK: %[[SIN:.*]] =  math.sin %[[ARG]] : vector<4xf32>
-//       CHECK128: %[[SIN:.*]] =  math.sin %[[ARG]] : vector<4xf32>
-//       CHECK0: %[[SIN:.*]] =  math.sin %[[ORIG_ARG]] : vector<2x2xf32>
+
+  // Arith and math ops are handled in generic way, check some of them
+  // DEFAULT: %[[SIN:.*]] =  math.sin %[[ARG]] : vector<4xf32>
+  // BW-128: %[[SIN:.*]] =  math.sin %[[ARG]] : vector<4xf32>
+  // BW-0: %[[SIN:.*]] =  math.sin %[[ORIG_ARG]] : vector<2x2xf32>
   %1 = math.sin %arg0 : vector<2x2xf32>
 
-  //     CHECK: %[[SIN1:.*]] =  math.sin %[[ARG2]] : vector<16xf32>
-//       CHECK128: %[[SIN1:.*]] =  math.sin %[[ORIG_ARG2]] : vector<4x4xf32>
-//       CHECK0: %[[SIN1:.*]] =  math.sin %[[ORIG_ARG2]] : vector<4x4xf32>
+  // DEFAULT: %[[SIN1:.*]] =  math.sin %[[ARG2]] : vector<16xf32>
+  // BW-128: %[[SIN1:.*]] =  math.sin %[[ORIG_ARG2]] : vector<4x4xf32>
+  // BW-0: %[[SIN1:.*]] =  math.sin %[[ORIG_ARG2]] : vector<4x4xf32>
   %6 = math.sin %arg1 : vector<4x4xf32>
-//       CHECK: %{{.*}} = arith.addf %[[ARG]], %[[C1]] : vector<4xf32>
-//       CHECK128: %{{.*}} = arith.addf %[[ARG]], %[[C1]] : vector<4xf32>
-//       CHECK0: %{{.*}} = arith.addf %{{.*}} : vector<2x2xf32>
 
+  // DEFAULT: %{{.*}} = arith.addf %[[ARG]], %[[CST]] : vector<4xf32>
+  // BW-128: %{{.*}} = arith.addf %[[ARG]], %[[CST]] : vector<4xf32>
+  // BW-0: %{{.*}} = arith.addf %{{.*}} : vector<2x2xf32>
   %2 = arith.addf %arg0, %0 :  vector<2x2xf32>
 
-  // CHECK: %[[ADD2:.*]] = arith.addf %[[ARG2]], %[[C2]] : vector<16xf32>
-  // CHECK128: %[[ADD2:.*]] = arith.addf %[[ORIG_ARG2]], %[[C2]] : vector<4x4xf32>
-  // CHECK0: %[[ADD2:.*]] = arith.addf %[[ORIG_ARG2]], %[[C2]] : vector<4x4xf32>
+  // DEFAULT: %[[ADD2:.*]] = arith.addf %[[ARG2]], %[[C2]] : vector<16xf32>
+  // BW-128: %[[ADD2:.*]] = arith.addf %[[ORIG_ARG2]], %[[C2]] : vector<4x4xf32>
+  // BW-0: %[[ADD2:.*]] = arith.addf %[[ORIG_ARG2]], %[[C2]] : vector<4x4xf32>
   %7 = arith.addf %arg1, %5 : vector<4x4xf32>
-//       CHECK: return %[[RES]] : vector<2x2xf32>
-//       CHECK128: return %[[RES]] : vector<2x2xf32>
+
+  // ALL: return %[[RES]] : vector<2x2xf32>
   return %0 : vector<2x2xf32>
 }
 
-// CHECK-LABEL: test_index_no_linearize
-// CHECK128-LABEL: test_index_no_linearize
-// CHECK0-LABEL: test_index_no_linearize
+// -----
+
+// ALL-LABEL: test_index_no_linearize
 func.func @test_index_no_linearize(%arg0: vector<2x2xindex>, %arg1: vector<2x2xindex>) -> vector<2x2xindex> {
-    // CHECK: %[[ADD:.*]] = arith.addi {{.*}} : vector<2x2xindex>
-    // CHECK128: %[[ADD:.*]] = arith.addi {{.*}} : vector<2x2xindex>
-    // CHECK0: %[[ADD:.*]] = arith.addi {{.*}} : vector<2x2xindex>
+    // DEFAULT: %[[ADD:.*]] = arith.addi {{.*}} : vector<2x2xindex>
+    // BW-128: %[[ADD:.*]] = arith.addi {{.*}} : vector<2x2xindex>
+    // BW-0: %[[ADD:.*]] = arith.addi {{.*}} : vector<2x2xindex>
     %0 = arith.addi %arg0, %arg1 : vector<2x2xindex>
     return %0 : vector<2x2xindex>
 }
@@ -95,10 +92,11 @@ func.func @test_index_no_linearize(%arg0: vector<2x2xindex>, %arg1: vector<2x2xi
 
 // vectorizable operation (arith.mulf) with tensor result types.
 
-func.func @nonvec_result(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> (tensor<2x2xf32>, tensor<2x2xf32>) {
-    // CHECK: %[[MULF:.*]] = arith.mulf %arg0, %arg1 : tensor<2x2xf32>
-    // CHECK128: %[[MULF:.*]] = arith.mulf %arg0, %arg1 : tensor<2x2xf32>
-    // CHECK0: %[[MULF:.*]] = arith.mulf %arg0, %arg1 : tensor<2x2xf32>
+// ALL-LABEL: test_tensor_no_linearize
+func.func @test_tensor_no_linearize(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> (tensor<2x2xf32>, tensor<2x2xf32>) {
+    // DEFAULT: %[[MULF:.*]] = arith.mulf %arg0, %arg1 : tensor<2x2xf32>
+    // BW-128: %[[MULF:.*]] = arith.mulf %arg0, %arg1 : tensor<2x2xf32>
+    // BW-0: %[[MULF:.*]] = arith.mulf %arg0, %arg1 : tensor<2x2xf32>
     %0 = arith.mulf %arg0, %arg1 : tensor<2x2xf32>
 
     return %0, %arg0 : tensor<2x2xf32>, tensor<2x2xf32>

>From d1dfd91497b870b9dd7c04279be0e7736640d04f Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Tue, 26 Mar 2024 15:55:16 +0000
Subject: [PATCH 2/2] fixup! [mlir][vector] Refactor linearize.mlir

---
 mlir/test/Dialect/Vector/linearize.mlir | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/mlir/test/Dialect/Vector/linearize.mlir b/mlir/test/Dialect/Vector/linearize.mlir
index 84325a3ccd319b..1b225c7a97d233 100644
--- a/mlir/test/Dialect/Vector/linearize.mlir
+++ b/mlir/test/Dialect/Vector/linearize.mlir
@@ -81,9 +81,7 @@ func.func @test_partial_linearize(%arg0: vector<2x2xf32>, %arg1: vector<4x4xf32>
 
 // ALL-LABEL: test_index_no_linearize
 func.func @test_index_no_linearize(%arg0: vector<2x2xindex>, %arg1: vector<2x2xindex>) -> vector<2x2xindex> {
-    // DEFAULT: %[[ADD:.*]] = arith.addi {{.*}} : vector<2x2xindex>
-    // BW-128: %[[ADD:.*]] = arith.addi {{.*}} : vector<2x2xindex>
-    // BW-0: %[[ADD:.*]] = arith.addi {{.*}} : vector<2x2xindex>
+    // ALL: %[[ADD:.*]] = arith.addi {{.*}} : vector<2x2xindex>
     %0 = arith.addi %arg0, %arg1 : vector<2x2xindex>
     return %0 : vector<2x2xindex>
 }
@@ -94,9 +92,7 @@ func.func @test_index_no_linearize(%arg0: vector<2x2xindex>, %arg1: vector<2x2xi
 
 // ALL-LABEL: test_tensor_no_linearize
 func.func @test_tensor_no_linearize(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> (tensor<2x2xf32>, tensor<2x2xf32>) {
-    // DEFAULT: %[[MULF:.*]] = arith.mulf %arg0, %arg1 : tensor<2x2xf32>
-    // BW-128: %[[MULF:.*]] = arith.mulf %arg0, %arg1 : tensor<2x2xf32>
-    // BW-0: %[[MULF:.*]] = arith.mulf %arg0, %arg1 : tensor<2x2xf32>
+    // ALL: %[[MULF:.*]] = arith.mulf %arg0, %arg1 : tensor<2x2xf32>
     %0 = arith.mulf %arg0, %arg1 : tensor<2x2xf32>
 
     return %0, %arg0 : tensor<2x2xf32>, tensor<2x2xf32>



More information about the Mlir-commits mailing list