forked from iree-org/iree
-
Notifications
You must be signed in to change notification settings - Fork 11
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[LLVMCPU][ArmSME] Rework how Arm streaming mode is set on dispatches (i…
…ree-org#17646) Previously, when the `+sme` feature flag was set Armv9 streaming SVE mode would be enabled for all dispatch regions lowered with the following experts: - `CPUBufferOpsTileAndVectorize` - `CPUConvTileAndDecomposeExpert` - `CPUDoubleTilingExpert` This was not ideal as meant streaming mode could be added to dispatch regions that made no use of scalable vectors, where the (possibly) larger streaming vector length provides no benefit, and there may be a cost due to other overheads. There was also a flag `--iree-experimental-llvmcpu-arm-force-ssve` which contrary to its name _did not_ force streaming SVE mode. What this flag did do was disable tiling for 2D scalable ArmSME operations, then rely on something else later on setting the streaming mode (but it did not control it). The patch aims to add clearer and more directed ways to enable streaming mode. First, streaming mode is no longer set in any lowering experts (it's a fairly low-level concept, that does not need to be configured early in the pipeline). Second, the old `--iree-experimental-llvmcpu-arm-force-ssve` flag is removed. Now to control tiling for ArmSME and using streaming mode there are two new flags. `iree-llvmcpu-disable-arm-sme-tiling`: This disables tiling for ArmSME (i.e. using 2D scalable tile sizes), even when the `+sme` feature flag is set. This results in operations instead being tiled for SVE or Neon (depending on the configuration). `iree-llvmcpu-force-arm-streaming`: This enables Arm streaming mode for any dispatch regions that contain scalable vectors. It ignores dispatches that don't contain scalable vectors as enabling streaming mode would provide no benefit. ci-extra: build_test_all_arm64 --------- Signed-off-by: Benjamin Maxwell <[email protected]>
- Loading branch information
Showing
7 changed files
with
157 additions
and
132 deletions.
There are no files selected for viewing
This file contains 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 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 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 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
94 changes: 94 additions & 0 deletions
94
compiler/src/iree/compiler/Codegen/LLVMCPU/test/pipeline_arm_sme_streaming_mode_tests.mlir
This file contains 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,94 @@ | ||
| // RUN: iree-opt --iree-codegen-linalg-to-llvm-pipeline=enable-arm-sme --split-input-file %s | FileCheck %s | ||
| // RUN: iree-opt --iree-codegen-linalg-to-llvm-pipeline=enable-arm-sme --iree-llvmcpu-force-arm-streaming --split-input-file %s | FileCheck %s -check-prefixes=FORCE-ARM-STREAMING | ||
|
|
||
| module { | ||
| module { | ||
| func.func @fixed_size_dispatch() attributes {hal.executable.target = #hal.executable.target<"llvm-cpu", "embedded-elf-arm_64", {cpu_features = "+sve,+sme", data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", native_vector_size = 16 : index, target_triple = "aarch64-none-elf"}>, | ||
| translation_info = #iree_codegen.translation_info<CPUDoubleTilingExpert>} { | ||
| %c0 = arith.constant 0 : index | ||
| %c1 = arith.constant 1 : index | ||
| %cst = arith.constant 0.000000e+00 : f32 | ||
| %0 = hal.interface.constant.load[0] : i32 | ||
| %1 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<readwrite:tensor<1xf32>> | ||
| %2 = tensor.empty() : tensor<1xf32> | ||
| %3 = linalg.fill {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[0], [1], [0], [0]]>} | ||
| ins(%cst : f32) outs(%2 : tensor<1xf32>) -> tensor<1xf32> | ||
| flow.dispatch.tensor.store %3, %1, offsets = [0], sizes = [1], strides = [1] : tensor<1xf32> -> !flow.dispatch.tensor<readwrite:tensor<1xf32>> | ||
| return | ||
| } | ||
| } | ||
| } | ||
|
|
||
| /// A dispatch region that only uses fixed-size vectors should never use | ||
| /// streaming mode. | ||
|
|
||
| // CHECK: @fixed_size_dispatch | ||
| // CHECK-NOT: arm_locally_streaming | ||
| // CHECK-NOT: arm_new_za | ||
|
|
||
| // FORCE-ARM-STREAMING: @fixed_size_dispatch | ||
| // FORCE-ARM-STREAMING-NOT: arm_locally_streaming | ||
| // FORCE-ARM-STREAMING-NOT: arm_new_za | ||
|
|
||
| // ----- | ||
|
|
||
| module { | ||
| module { | ||
| func.func @scalable_dispatch() attributes {hal.executable.target = #hal.executable.target<"llvm-cpu", "embedded-elf-arm_64", {cpu_features = "+sve,+sme", data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", native_vector_size = 16 : index, target_triple = "aarch64-none-elf"}>, | ||
| translation_info = #iree_codegen.translation_info<CPUDoubleTilingExpert>} { | ||
| %c0 = arith.constant 0 : index | ||
| %c1 = arith.constant 1 : index | ||
| %cst = arith.constant 0.000000e+00 : f32 | ||
| %0 = hal.interface.constant.load[0] : i32 | ||
| %1 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<readwrite:tensor<1xf32>> | ||
| %2 = tensor.empty() : tensor<1xf32> | ||
| %3 = linalg.fill {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[0], [[1]], [0], [0]]>} | ||
| ins(%cst : f32) outs(%2 : tensor<1xf32>) -> tensor<1xf32> | ||
| flow.dispatch.tensor.store %3, %1, offsets = [0], sizes = [1], strides = [1] : tensor<1xf32> -> !flow.dispatch.tensor<readwrite:tensor<1xf32>> | ||
| return | ||
| } | ||
| } | ||
| } | ||
|
|
||
| /// A dispatch region that uses scalable vectors (but not ArmSME dialect | ||
| /// operations) should only use streaming if `--iree-llvmcpu-force-arm-streaming` | ||
| /// is set. | ||
|
|
||
| // CHECK: @scalable_dispatch | ||
| // CHECK-NOT: arm_locally_streaming | ||
| // CHECK-NOT: arm_new_za | ||
|
|
||
| // FORCE-ARM-STREAMING: @scalable_dispatch | ||
| // FORCE-ARM-STREAMING-NOT: arm_new_za | ||
| // FORCE-ARM-STREAMING-SAME: arm_locally_streaming | ||
|
|
||
| // ----- | ||
|
|
||
| module { | ||
| module { | ||
| func.func @scalable_dispatch_using_za() attributes {hal.executable.target = #hal.executable.target<"llvm-cpu", "embedded-elf-arm_64", {cpu_features = "+sve,+sme", data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", native_vector_size = 16 : index, target_triple = "aarch64-none-elf"}>, | ||
| translation_info = #iree_codegen.translation_info<CPUDoubleTilingExpert>} { | ||
| %c0 = arith.constant 0 : index | ||
| %c1 = arith.constant 1 : index | ||
| %cst = arith.constant 0.000000e+00 : f32 | ||
| %0 = hal.interface.constant.load[0] : i32 | ||
| %1 = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<readwrite:tensor<100x100xf32>> | ||
| %2 = tensor.empty() : tensor<100x100xf32> | ||
| %3 = linalg.fill {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[0, 0], [[4], [4]], [0, 0], [0, 0]]>} | ||
| ins(%cst : f32) outs(%2 : tensor<100x100xf32>) -> tensor<100x100xf32> | ||
| flow.dispatch.tensor.store %3, %1, offsets = [0, 0], sizes = [100, 100], strides = [100, 1] : tensor<100x100xf32> -> !flow.dispatch.tensor<readwrite:tensor<100x100xf32>> | ||
| return | ||
| } | ||
| } | ||
| } | ||
|
|
||
| /// A dispatch region that uses ArmSME operations (that require the ZA state) | ||
| /// should ways have streaming mode and ZA enabled. | ||
|
|
||
| // CHECK: @scalable_dispatch_using_za | ||
| // CHECK-SAME: arm_locally_streaming | ||
| // CHECK-SAME: arm_new_za | ||
|
|
||
| // FORCE-ARM-STREAMING: @scalable_dispatch_using_za | ||
| // FORCE-ARM-STREAMING-SAME: arm_locally_streaming | ||
| // FORCE-ARM-STREAMING-SAME: arm_new_za |
90 changes: 0 additions & 90 deletions
90
compiler/src/iree/compiler/Codegen/LLVMCPU/test/pipeline_ssve_tests.mlir
This file was deleted.
Oops, something went wrong.
Oops, something went wrong.