Skip to content

[mlir][vector] Refactor linearize.mlir #86648

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Mar 26, 2024
Merged

[mlir][vector] Refactor linearize.mlir #86648

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
bcf0eeb
[mlir][vector] Refactor linearize.mlir
banach-space Mar 26, 2024
d1dfd91
fixup! [mlir][vector] Refactor linearize.mlir
banach-space Mar 26, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
128 changes: 61 additions & 67 deletions mlir/test/Dialect/Vector/linearize.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,92 +1,87 @@
// 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
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why did you add a new file-prefix instead of keeping the default one (by default I meant just using "CHECK" not the "DEFAULT" tag)?

Copy link
Contributor Author

@banach-space banach-space Mar 26, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

DEFAULT documents the fact that the default config is being tested. And -test-vector-linearize means "default". When using CHECK, we are missing an opportunity to document what configuration is being tested (we know that it's a "check" label regardless of the actual name).

I appreciate that it's a matter of style/preference and both approaches are 100% valid. Happy to revert this change if that's your preference.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No need to revert. Just asking for my knowledge :).

// 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>
// ALL: %[[ADD:.*]] = arith.addi {{.*}} : vector<2x2xindex>
%0 = arith.addi %arg0, %arg1 : vector<2x2xindex>
return %0 : vector<2x2xindex>
}
Expand All @@ -95,10 +90,9 @@ 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>) {
// ALL: %[[MULF:.*]] = arith.mulf %arg0, %arg1 : tensor<2x2xf32>
%0 = arith.mulf %arg0, %arg1 : tensor<2x2xf32>

return %0, %arg0 : tensor<2x2xf32>, tensor<2x2xf32>
Expand Down