[mlir][ArmSME] Lower vector.outerproduct to FMOPA/BFMOPA #65621
+430
−17
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,116 @@ | ||
| // DEFINE: %{entry_point} = test_outerproduct_no_accumulator_4x4xf32 | ||
| // DEFINE: %{compile} = mlir-opt %s \ | ||
| // DEFINE: -enable-arm-streaming="mode=locally enable-za" \ | ||
| // DEFINE: -convert-vector-to-arm-sme -convert-arm-sme-to-scf \ | ||
| // DEFINE: -convert-vector-to-llvm="enable-arm-sme" -cse -canonicalize \ | ||
| // DEFINE: -allocate-arm-sme-tiles -test-lower-to-llvm | ||
| // DEFINE: %{run} = %mcr_aarch64_cmd \ | ||
| // DEFINE: -march=aarch64 -mattr=+sve,+sme \ | ||
| // DEFINE: -e %{entry_point} -entry-point-result=void \ | ||
| // DEFINE: -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils | ||
|
|
||
| // RUN: %{compile} | %{run} | FileCheck %s --check-prefix=WITHOUT-ACC | ||
|
|
||
| // REDEFINE: %{entry_point} = test_outerproduct_with_accumulator_4x4xf32 | ||
| // RUN: %{compile} | %{run} | FileCheck %s --check-prefix=WITH-ACC | ||
|
|
||
| llvm.func @printCString(!llvm.ptr<i8>) | ||
|
|
||
| func.func @printTileBegin() { | ||
| %0 = llvm.mlir.addressof @str_tile_begin : !llvm.ptr<array<11 x i8>> | ||
| %1 = llvm.mlir.constant(0 : index) : i64 | ||
| %2 = llvm.getelementptr %0[%1, %1] | ||
| : (!llvm.ptr<array<11 x i8>>, i64, i64) -> !llvm.ptr<i8> | ||
| llvm.call @printCString(%2) : (!llvm.ptr<i8>) -> () | ||
| return | ||
| } | ||
|
|
||
| func.func @printTileEnd() { | ||
| %0 = llvm.mlir.addressof @str_tile_end : !llvm.ptr<array<9 x i8>> | ||
| %1 = llvm.mlir.constant(0 : index) : i64 | ||
| %2 = llvm.getelementptr %0[%1, %1] | ||
| : (!llvm.ptr<array<9 x i8>>, i64, i64) -> !llvm.ptr<i8> | ||
| llvm.call @printCString(%2) : (!llvm.ptr<i8>) -> () | ||
| return | ||
| } | ||
|
|
||
| func.func @test_outerproduct_no_accumulator_4x4xf32() { | ||
| %c0 = arith.constant 0 : index | ||
|
|
||
| %vector_i32 = llvm.intr.experimental.stepvector : vector<[4]xi32> | ||
| %vector = arith.sitofp %vector_i32 : vector<[4]xi32> to vector<[4]xf32> | ||
| %tile = vector.outerproduct %vector, %vector : vector<[4]xf32>, vector<[4]xf32> | ||
|
|
||
| // Calculate the size of a 32-bit tile, e.g. ZA{n}.s. | ||
| %vscale = vector.vscale | ||
| %min_elts_s = arith.constant 4 : index | ||
| %svl_s = arith.muli %min_elts_s, %vscale : index | ||
| %za_s_size = arith.muli %svl_s, %svl_s : index | ||
|
|
||
| // Allocate memory. | ||
| %mem = memref.alloca(%za_s_size) : memref<?xf32> | ||
|
|
||
| // Store the tile to memory. | ||
| vector.store %tile, %mem[%c0] : memref<?xf32>, vector<[4]x[4]xf32> | ||
|
|
||
| // Reload and print. The smallest SVL is 128-bits so the tile will be at | ||
| // least 4x4xf32. | ||
| // | ||
| // WITHOUT-ACC: TILE BEGIN | ||
| // WITHOUT-ACC-NEXT: ( 0, 0, 0, 0 | ||
| // WITHOUT-ACC-NEXT: ( 0, 1, 2, 3 | ||
| // WITHOUT-ACC-NEXT: ( 0, 2, 4, 6 | ||
| // WITHOUT-ACC-NEXT: ( 0, 3, 6, 9 | ||
| // WITHOUT-ACC: TILE END | ||
| func.call @printTileBegin() : () -> () | ||
| scf.for %i = %c0 to %za_s_size step %svl_s { | ||
| %tileslice = vector.load %mem[%i] : memref<?xf32>, vector<[4]xf32> | ||
| vector.print %tileslice : vector<[4]xf32> | ||
| } | ||
| func.call @printTileEnd() : () -> () | ||
|
|
||
| return | ||
| } | ||
|
|
||
| func.func @test_outerproduct_with_accumulator_4x4xf32() { | ||
| %c0 = arith.constant 0 : index | ||
| %f10 = arith.constant 10.0 : f32 | ||
|
|
||
| %acc = vector.broadcast %f10 : f32 to vector<[4]x[4]xf32> | ||
| %vector_i32 = llvm.intr.experimental.stepvector : vector<[4]xi32> | ||
| %vector = arith.sitofp %vector_i32 : vector<[4]xi32> to vector<[4]xf32> | ||
| %tile = vector.outerproduct %vector, %vector, %acc : vector<[4]xf32>, vector<[4]xf32> | ||
|
|
||
| // Calculate the size of a 32-bit tile, e.g. ZA{n}.s. | ||
| %vscale = vector.vscale | ||
| %min_elts_s = arith.constant 4 : index | ||
| %svl_s = arith.muli %min_elts_s, %vscale : index | ||
| %za_s_size = arith.muli %svl_s, %svl_s : index | ||
|
|
||
| // Allocate memory. | ||
| %mem = memref.alloca(%za_s_size) : memref<?xf32> | ||
|
|
||
| // Store the tile to memory. | ||
| vector.store %tile, %mem[%c0] : memref<?xf32>, vector<[4]x[4]xf32> | ||
|
|
||
| // Reload and print. The smallest SVL is 128-bits so the tile will be at | ||
| // least 4x4xf32. | ||
| // | ||
| // WITH-ACC: TILE BEGIN | ||
| // WITH-ACC-NEXT: ( 10, 10, 10, 10 | ||
| // WITH-ACC-NEXT: ( 10, 11, 12, 13 | ||
| // WITH-ACC-NEXT: ( 10, 12, 14, 16 | ||
| // WITH-ACC-NEXT: ( 10, 13, 16, 19 | ||
| // WITH-ACC: TILE END | ||
| func.call @printTileBegin() : () -> () | ||
| scf.for %i = %c0 to %za_s_size step %svl_s { | ||
| %tileslice = vector.load %mem[%i] : memref<?xf32>, vector<[4]xf32> | ||
| vector.print %tileslice : vector<[4]xf32> | ||
| } | ||
| func.call @printTileEnd() : () -> () | ||
|
|
||
| return | ||
| } | ||
|
|
||
| llvm.mlir.global internal constant @str_tile_begin("TILE BEGIN\0A") | ||
| llvm.mlir.global internal constant @str_tile_end("TILE END\0A") |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,77 @@ | ||
| // DEFINE: %{entry_point} = test_outerproduct_with_accumulator_2x2xf64 | ||
| // DEFINE: %{compile} = mlir-opt %s \ | ||
| // DEFINE: -enable-arm-streaming="mode=locally enable-za" \ | ||
| // DEFINE: -convert-vector-to-arm-sme -convert-arm-sme-to-scf \ | ||
| // DEFINE: -convert-vector-to-llvm="enable-arm-sme" -cse -canonicalize \ | ||
| // DEFINE: -allocate-arm-sme-tiles -test-lower-to-llvm | ||
| // DEFINE: %{run} = %mcr_aarch64_cmd \ | ||
| // DEFINE: -march=aarch64 -mattr=+sve,+sme-f64f64 \ | ||
| // DEFINE: -e %{entry_point} -entry-point-result=void \ | ||
| // DEFINE: -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils | ||
|
|
||
| // RUN: %{compile} | %{run} | FileCheck %s | ||
|
|
||
| llvm.func @printCString(!llvm.ptr<i8>) | ||
|
|
||
| func.func @printTileBegin() { | ||
| %0 = llvm.mlir.addressof @str_tile_begin : !llvm.ptr<array<11 x i8>> | ||
| %1 = llvm.mlir.constant(0 : index) : i64 | ||
| %2 = llvm.getelementptr %0[%1, %1] | ||
| : (!llvm.ptr<array<11 x i8>>, i64, i64) -> !llvm.ptr<i8> | ||
| llvm.call @printCString(%2) : (!llvm.ptr<i8>) -> () | ||
| return | ||
| } | ||
|
|
||
| func.func @printTileEnd() { | ||
| %0 = llvm.mlir.addressof @str_tile_end : !llvm.ptr<array<9 x i8>> | ||
| %1 = llvm.mlir.constant(0 : index) : i64 | ||
| %2 = llvm.getelementptr %0[%1, %1] | ||
| : (!llvm.ptr<array<9 x i8>>, i64, i64) -> !llvm.ptr<i8> | ||
| llvm.call @printCString(%2) : (!llvm.ptr<i8>) -> () | ||
| return | ||
| } | ||
|
|
||
| func.func @test_outerproduct_with_accumulator_2x2xf64() { | ||
| %c0 = arith.constant 0 : index | ||
| %f1 = arith.constant 1.0 : f64 | ||
| %f2 = arith.constant 2.0 : f64 | ||
| %f10 = arith.constant 10.0 : f64 | ||
|
|
||
| %a = vector.splat %f1 : vector<[2]xf64> | ||
| %b = vector.splat %f2 : vector<[2]xf64> | ||
| // TODO: vector.splat doesn't support ArmSME. | ||
| %c = vector.broadcast %f10 : f64 to vector<[2]x[2]xf64> | ||
|
|
||
| %tile = vector.outerproduct %a, %b, %c : vector<[2]xf64>, vector<[2]xf64> | ||
|
|
||
| // Calculate the size of a 64-bit tile, e.g. ZA{n}.d. | ||
| %vscale = vector.vscale | ||
| %min_elts_d = arith.constant 2 : index | ||
| %svl_d = arith.muli %min_elts_d, %vscale : index | ||
| %za_d_size = arith.muli %svl_d, %svl_d : index | ||
|
|
||
| // Allocate memory. | ||
| %mem = memref.alloca(%za_d_size) : memref<?xf64> | ||
|
|
||
| // Store the tile to memory. | ||
| vector.store %tile, %mem[%c0] : memref<?xf64>, vector<[2]x[2]xf64> | ||
|
|
||
| // Reload and print. The smallest SVL is 128-bits so the tile will be at | ||
| // least 2x2xf64. | ||
| // | ||
| // CHECK: TILE BEGIN | ||
| // CHECK-NEXT: ( 12, 12 | ||
| // CHECK-NEXT: ( 12, 12 | ||
| // CHECK: TILE END | ||
| func.call @printTileBegin() : () -> () | ||
| scf.for %i = %c0 to %za_d_size step %svl_d { | ||
| %tileslice = vector.load %mem[%i] : memref<?xf64>, vector<[2]xf64> | ||
| vector.print %tileslice : vector<[2]xf64> | ||
| } | ||
| func.call @printTileEnd() : () -> () | ||
|
|
||
| return | ||
| } | ||
|
|
||
| llvm.mlir.global internal constant @str_tile_begin("TILE BEGIN\0A") | ||
| llvm.mlir.global internal constant @str_tile_end("TILE END\0A") |
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
these two have the same name?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This intrinsic takes two operands to mask the inputs in the op itself so to support masking you would have to propagate the mask from the producer ops... That's interesting because it looks like the op knows how to merge both masks without requiring independent mask manipulation operations.
How do we plan to implement proper support for this? I see two options:
I think doing all of that as part of the lowering (1) might be too much for a lowering, esp. if finding the masks through the use-def chain is not trivial. (2) seems simpler to me but I wouldn't implement that on top of an llvm intrinsic. I think for that we should have a proper sme op, which should be fine now that we have the sme dialect.
Thoughts?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If you're referring to the predicate (?) that's because they're the same, both all active
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The only examples I've seen of masking (from grepping around the codebase) are where the mask is applied to the result of the outerproduct e.g.
vector.mask { vector.outerproduct ... }, I just figured we'd need some way to correlate this to the inputs, but hadn't given it much thought.Appreciate your input, I'll add a custom op that way there's more flexibility when it comes to masking, and will also look into how it would be used.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Given the expected complexity with generating correct masks, I am also leaning towards a custom op. Having said that, IMHO this PR is fine as is and we could iterate in the follow-up patches.
I guess that for this to work, we'd need something like `
So, we'd allow 2 optional masks, both of which would be optional:
ptrue(all lanes are active).WDYT?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We perhaps could use
vector.maskfor masking the result so that we don't have to disambiguate the semantics based on the number of masks...