[mlir][sparse][ArmSVE] Add sparse integration tests for ArmSVE

[awarzynski]

Summary

LLVM backend for AArch64 does not currently support product reductions
so it requires some test code duplication to have a version for the
addition reductions. For all the other tests, we can run the vanilla
version with VLA compilation options.

[jsetoain]

Include scalable vector tests with the others

[jsetoain]

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions.mlir

11–12 ↗

(On Diff #425205)

I need to remove this

[jsetoain]

Add partial reduction tests for VLA testing

[jsetoain]

Fix broken test. All VLA sparse tests have been verified.

[aartbik]

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions_vla.mlir

2

For the fully duplicated test (new files), would it make sense to put them in

mlir/test/Integration/Dialect/SparseTensor/CPU/ArmSVE

Similar to what we did for

mlir/test/Integration/Dialect/Vector/CPU/ArmSVE

[jsetoain]

Move test to ArmSVE-specific directory

[awarzynski]

Hi @jsetoain , thanks very much for this patch! I know that you have no access to SVE hardware ATM, so it might be tricky for you to finish this. And we (Arm) would love to have this merged :) Shall I take over? I do have access to the required testing infra.

Btw, these tests run mostly fine for me, but lli behaves slightly differently to mlir-cpu-runner - I had to add return 0 for the tests to pass. So, why not just use mlir-cpu-runner instead? I suspect that that's something to do with SVE support?

[tschuett]

In D121304#3860608, @awarzynski wrote:

Hi @jsetoain , thanks very much for this patch! I know that you have no access to SVE hardware ATM, so it might be tricky for you to finish this. And we (Arm) would love to have this merged :) Shall I take over? I do have access to the required testing infra.

Btw, these tests run mostly fine for me, but lli behaves slightly differently to mlir-cpu-runner - I had to add return 0 for the tests to pass. So, why not just use mlir-cpu-runner instead? I suspect that that's something to do with SVE support?

You can rent SVE hardware at Amazon:
https://aws.amazon.com/ec2/instance-types/c7g/

[jsetoain]

In D121304#3860608, @awarzynski wrote:

Hi @jsetoain , thanks very much for this patch! I know that you have no access to SVE hardware ATM, so it might be tricky for you to finish this. And we (Arm) would love to have this merged :) Shall I take over? I do have access to the required testing infra.

Btw, these tests run mostly fine for me, but lli behaves slightly differently to mlir-cpu-runner - I had to add return 0 for the tests to pass. So, why not just use mlir-cpu-runner instead? I suspect that that's something to do with SVE support?

Hi Andrzej! By all means, please, feel free to take over. I'm trying to figure out a way to transfer the patch to you. The reason why I'm using lli here is because mlir-cpu-runner was not working. I don't remember the issue, exactly, and it might have been fixed by now. If you manage to get the test running with mlir-cpu-runner, that's definitely preferable to lli.

Thanks! 😊

[awarzynski]

I'm trying to figure out a way to transfer the patch to you.

Done :)

[awarzynski]

@aartbik Given https://reviews.llvm.org/D136183, we should probably park this for now?

CC @jsetoain

[jsetoain]

Agreed!

[awarzynski]

Abandoning this for now. Let's wait for updates on the sparse compiler before we decide what to do next: https://discourse.llvm.org/t/mlir-sparse-compiler-progress/60479/8.

[aartbik]

In D121304#3883847, @awarzynski wrote:

Abandoning this for now. Let's wait for updates on the sparse compiler before we decide what to do next: https://discourse.llvm.org/t/mlir-sparse-compiler-progress/60479/8.

We will soon be back online with a much better SIMD approach!
Apologies for the slight detour in the meantime.

[awarzynski]

With https://reviews.llvm.org/D138236 in-tree we can re-open this.

[awarzynski]

Rebased on top of main.

I switched from lli to mlir-cpu-runner as that was the path of least resitance for now. However, there's no -march and -mattr in mlir-cpu-runner. I need to double check that indeed SVE code is generated and run. If not, we either need to update mlir-cpu-runner or I need to use lli instead.

[aartbik]

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_flatten.mlir

28

Just a very quick comment, this will not be the right way to run the sparse vectorizer.

We will have to pass in the flags to the pipeline command sparse-compiler (see original vector example), but adding this support is still TBD. I was waiting for the vectorizer to be "production" ready before adding those ;-)

[awarzynski]

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_flatten.mlir

28

Thanks @aartbik ! Yeah, I jumped the gun here a bit :)

[aartbik]

Revision https://reviews.llvm.org/D139581 should enable you to restart the efforts again!

[awarzynski]

Rebase on top of main

@aartbik Could you take a quick and let me know whether this makes sense? 3
tests are currently failing for me:

Failed Tests (3):
  MLIR :: Integration/Dialect/SparseTensor/CPU/sparse_quantized_matmul.mlir
  MLIR :: Integration/Dialect/SparseTensor/CPU/sparse_sampled_matmul.mlir
  MLIR :: Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir

I didn't get a chance to triage just yet.

[awarzynski]

Mostly just rebase on top of main

[aartbik]

Can you please rebase until at least https://reviews.llvm.org/D139983 and try again.

Also, please note you may need enable-buffer-initialization=true if you use vector_transfer/print to CHECK the output if you use codegen instead of lib.

[awarzynski]

Rebase on top of main

I've also simplified my changes a bit to make them less intrusive.

[aartbik]

mlir/test/Integration/Dialect/SparseTensor/CPU/ArmSVE/lit.local.cfg

20 ↗

(On Diff #483214)

period at end

mlir/test/Integration/Dialect/SparseTensor/CPU/ArmSVE/sparse_reductions_vla.mlir

1 ↗

(On Diff #483214)

refresh my memory, why do we need the ArmSVE specific directory?
It seems that most test fit well at the same level, so we do no need this test and the lit local file?

mlir/test/Integration/Dialect/SparseTensor/CPU/lit.local.cfg

21

period at end

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir

13

Phrase this a bit in the same style as the others,

If SVE is available, do the same run, but now with direct IR generation and VLA vectorization.

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir

85

Do you see a way to avoid adding this to *all* tests? Can we put this in its own file and let the command line compile take care of it?

[aartbik]

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir

85

Here is an idea, add a new file entry.mlir with just this entry method, and then do

RUN: cat $s entry.mlir | rest of pipeline

so that the two files are concatenated into a single stream

[awarzynski]

Address PR comments

Thanks for the comments @aartbik :)

• Moved entry_lli to a dedicated file (avoids code duplication)
• Unified tests to use direct IR generation
• Updated comments

refresh my memory, why do we need the ArmSVE specific directory?

Let me get back to you tomorrow - first I need to go over the discussion and refresh my own memory.

[awarzynski]

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir

85

Thanks! I've derived a slightly different solution - see the latest update. it turns out lli will happily consume multiple files. Let me know whether you have any preference - I'm happy to try your suggestion instead.

[aartbik]

mlir/test/Integration/Dialect/SparseTensor/CPU/Inputs/main_for_lli.ll

1 ↗

(On Diff #485408)

I am okay keeping this at the same level. No need to introduce a whole new dir for just this

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir

80

please make sure to avoid unrelated diffs

85

even better

[awarzynski]

mlir/test/Integration/Dialect/SparseTensor/CPU/Inputs/main_for_lli.ll

1 ↗

(On Diff #485408)

We need to make sure that this file is not considered as a yet another test file. By putting it in "Inputs", we don't require any additional logic for this - see the definition of config.excludes.

I've not been able to find any documentation for this, but I've seen it used in various other sub-projects.

[awarzynski]

mlir/test/Integration/Dialect/SparseTensor/CPU/ArmSVE/sparse_reductions_vla.mlir

1 ↗

(On Diff #483214)

The answer is in the original commit summary:

LLVM backend for AArch64 does not currently support product reductions so it requires some test code duplication to have a version for the addition reductions.

I did check the backend for a clear indication of what is currently supported and what is not, but it's a bit convoluted:

• relevant snippet in LegalizeVectorOps.cpp, which leads to ...
• ... an error in expandVecReduce.

I will reach out internally to learn a bit more about the status of SVE in the context of reductions, but need to wait for folks to return from their Xmas breaks first.

I do think that we should avoid code duplication though and creating a new sub-directory here. Instead, I suggest splitting "sparse_reductions.mlir" into 2 files:

• one that tests reductions supported by SVE
• another one that contains the remaining reductions.

Let me upload this so that you get a better idea.

[awarzynski]

Split "sparse_reductions.mlir" into 2 test files

I split "sparse_reductions.mlir" into:

• sparse_reductions_1.mlir (reductions supported by SVE)
• sparse_reductions_2.mlir (all other reductions)

This way we can avoid adding "ArmSVE/sparse_reductions_vla.mlir", which
duplicated SVE-supported reductions from "sparse_reductions.mlir".

[awarzynski]

Remove the remaining unrelated changes

Also made sure that all SVE tests use direct IR generation

[aartbik]

Yeah, this is LGTM after making sure the emulator works and addressing my last nit request.

mlir/test/Integration/Dialect/SparseTensor/CPU/Inputs/main_for_lli.ll

1 ↗

(On Diff #485408)

Ah, okay, makes sense.

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions_2.mlir

1 ↗

(On Diff #485704)

I like this solution of splitting up the files much better than having yet another directory. Thanks!

But one last nit, I don't really like the _1 _2 naming, as it does not add much. Perhaps we can some up with a better name, like just sparse_reductions.mlir for the original and sparse_reductions_prod.mlir for this file, so it becomes more descriptive?

[awarzynski]

In D121304#4024181, @aartbik wrote:

Yeah, this is LGTM after making sure the emulator works and addressing my last nit request.

I've experimented with both QEMU and ArmIE. While QEMU works, it crashes at the end of a simulation when running lli. I've tried to reduce my issue, but I can only reproduce it with lli (which is a relatively huge binary). As I haven't been able to find any evidence that QEMU actually ever worked here (end-to-end), I suggest leaving it for now. lli is rather too large as a reproducer. Any suggestions how to reduce this or where to report it?

Instead, I used ArmIE. It worked without any issues (I needed to replace lli with %lli in tests). It's a free emulator and the official instructions worked for me just fine. So I am hoping that this is sufficient here. I used the following CMake set-up:

cmake -DMLIR_INCLUDE_INTEGRATION_TESTS=True -DMLIR_RUN_ARM_SVE_TESTS=True -DARM_EMULATOR_OPTIONS="-msve-vector-bits=128" -DARM_EMULATOR_EXECUTABLE="armie" <other cmake options>

I'm just about to send a small update that will address the outstanding comment - rename "sparse_reductions_{1|2}.mlir". So, is this ready to be merged? :)

[awarzynski]

Replace lli with %lli to fix substition, rename test files

[Peiming]

I'm just about to send a small update that will address the outstanding comment - rename "sparse_reductions_{1|2}.mlir". So, is this ready to be merged? :)

Aart is currently on vacation and will be back after Jan 20th. I am in his team and I would suggest that we wait until he is back because I am afraid that I do not have enough context/knowledge to make the decision by myself for this patch ;-)

mlir/test/Integration/Dialect/SparseTensor/CPU/lit.local.cfg

 import sys
 
 # No JIT on win32.
 if sys.platform == 'win32':
     config.unsupported = True
+
+# ArmSVE tests must be enabled via build flag.
+if config.mlir_run_arm_sve_tests == 'ON':
+    config.substitutions.append(('%ENABLE_VLA', 'true'))
+    config.substitutions.append(('%VLA_ARCH_ATTR_OPTIONS', '--march=aarch64 --mattr="+sve"'))
+    lli_cmd = 'lli'
+    if config.arm_emulator_lli_executable:
+       lli_cmd = config.arm_emulator_lli_executable
+
+    if config.arm_emulator_utils_lib_dir:
+        config.substitutions.append(('%mlir_native_utils_lib_dir', config.arm_emulator_utils_lib_dir))
+    else:
+        config.substitutions.append(('%mlir_native_utils_lib_dir', config.mlir_integration_test_dir))
+
+    if config.arm_emulator_executable:
+        # Run test in emulator (qemu or armie)
+        emulation_cmd = config.arm_emulator_executable
+        if config.arm_emulator_options:
+            emulation_cmd = emulation_cmd + ' ' + config.arm_emulator_options
+        emulation_cmd = emulation_cmd + ' ' + lli_cmd
+        config.substitutions.append(('%lli', emulation_cmd))
+    else:
+        config.substitutions.append(('%lli', lli_cmd))
+else:
+    config.substitutions.append(('%lli', 'lli'))
+    config.substitutions.append(('%mlir_native_utils_lib_dir', config.mlir_integration_test_dir))
+    config.substitutions.append(('%ENABLE_VLA', 'false'))
+    config.substitutions.append(('%VLA_ARCH_ATTR_OPTIONS', ''))

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
 // RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=2" | \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
+//
+// If SVE is available, test VLA vectorization.
+//
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=2 enable-vla-vectorization=%ENABLE_VLA" | \
+// RUN: mlir-translate -mlir-to-llvmir | \
+// RUN: %lli --entry-function=entry %VLA_ARCH_ATTR_OPTIONS --dlopen=%mlir_native_utils_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
 
 #SV = #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>
 
 #trait_cast = {
   indexing_maps = [
     affine_map<(i) -> (i)>,  // A (in)
     affine_map<(i) -> (i)>   // X (out)
   ],

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
 // RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=2" | \
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
+//
+// If SVE is available, test VLA vectorization.
+//
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=2 enable-vla-vectorization=%ENABLE_VLA" | \
+// RUN: mlir-translate -mlir-to-llvmir | \
+// RUN: %lli --entry-function=entry %VLA_ARCH_ATTR_OPTIONS --dlopen=%mlir_native_utils_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
 
 #DCSR = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>
 
 // An example of a 2D convolution with a sparse filter.
 module {
 
   func.func @conv2d(%input:  tensor<8x8xi32>,
                %filter: tensor<3x3xi32, #DCSR>,
func.func @entry() {
     //
     %m = bufferization.to_memref %0 : memref<6x6xi32>
     %v = vector.transfer_read %m[%c0, %c0], %i0
       : memref<6x6xi32>, vector<6x6xi32>
     vector.print %v : vector<6x6xi32>
 
     // Release the resources.
     sparse_tensor.release %sparse_filter : tensor<3x3xi32, #DCSR>
     memref.dealloc %m : memref<6x6xi32>
 
     return
   }
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_flatten.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/test.tns" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
 // RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=4" | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/test.tns" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
+//
+// If SVE is available, test VLA vectorization.
+//
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=4 enable-vla-vectorization=%ENABLE_VLA" | \
+// RUN: mlir-translate -mlir-to-llvmir | \
+// RUN: TENSOR0="%mlir_integration_test_dir/data/test.tns" \
+// RUN: %lli --entry-function=entry %VLA_ARCH_ATTR_OPTIONS --dlopen=%mlir_native_utils_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
 
 !Filename = !llvm.ptr<i8>
 
 #SparseTensor = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed", "compressed", "compressed",
                    "compressed", "compressed", "compressed", "compressed" ],
   // Note that any dimOrdering permutation should give the same results
   // since, even though it impacts the sparse storage scheme layout,
   // it should not change the semantics.
   dimOrdering = affine_map<(i,j,k,l,m,n,o,p) -> (p,o,j,k,i,l,m,n)>
 }>
 
 #trait_flatten = {

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_index_dense.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
 // RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=4" | \
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
+//
+// If SVE is available, test VLA vectorization.
+//
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=4 enable-vla-vectorization=%ENABLE_VLA" | \
+// RUN: mlir-translate -mlir-to-llvmir | \
+// RUN: %lli --entry-function=entry %VLA_ARCH_ATTR_OPTIONS --dlopen=%mlir_native_utils_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
 
 #SparseVector = #sparse_tensor.encoding<{
   dimLevelType = ["compressed"]
 }>
 
 #SparseMatrix = #sparse_tensor.encoding<{
   dimLevelType = ["compressed", "compressed"]
 }>

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/wide.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
 // RUN: mlir-opt %s \
 // RUN:   --sparse-compiler="vectorization-strategy=2 vl=16 enable-simd-index32" | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/wide.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
+//
+// If SVE is available, test VLA vectorization.
+//
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=16 enable-vla-vectorization=%ENABLE_VLA enable-simd-index32" | \
+// RUN: mlir-translate -mlir-to-llvmir | \
+// RUN: TENSOR0="%mlir_integration_test_dir/data/wide.mtx" \
+// RUN: %lli --entry-function=entry %VLA_ARCH_ATTR_OPTIONS --dlopen=%mlir_native_utils_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
 
 !Filename = !llvm.ptr<i8>
 
 #SparseMatrix = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "compressed" ],
   pointerBitWidth = 8,
   indexBitWidth = 8
 }>

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_mttkrp.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/mttkrp_b.tns" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
 // RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=4" | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/mttkrp_b.tns" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
+//
+// If SVE is available, test VLA vectorization.
+//
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=4 enable-vla-vectorization=%ENABLE_VLA" | \
+// RUN: mlir-translate -mlir-to-llvmir | \
+// RUN: TENSOR0="%mlir_integration_test_dir/data/mttkrp_b.tns" \
+// RUN: %lli --entry-function=entry %VLA_ARCH_ATTR_OPTIONS --dlopen=%mlir_native_utils_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
 
 !Filename = !llvm.ptr<i8>
 
 #SparseTensor = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed", "compressed" ]
 }>
 
 #mttkrp = {

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_simple.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/test.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
 // RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=4" | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/test.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
+//
+// If SVE is available, test VLA vectorization.
+//
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=4 enable-vla-vectorization=%ENABLE_VLA" | \
+// RUN: mlir-translate -mlir-to-llvmir | \
+// RUN: TENSOR0="%mlir_integration_test_dir/data/test.mtx" \
+// RUN: %lli --entry-function=entry %VLA_ARCH_ATTR_OPTIONS --dlopen=%mlir_native_utils_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
 
 !Filename = !llvm.ptr<i8>
 
 #DCSR = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed" ],
   dimOrdering = affine_map<(i,j) -> (i,j)>
 }>
 

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_quantized_matmul.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
 // RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=2" | \
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
+//
+// If SVE is available, test VLA vectorization.
+//
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=2 enable-vla-vectorization=%ENABLE_VLA" | \
+// RUN: mlir-translate -mlir-to-llvmir | \
+// RUN: %lli --entry-function=entry %VLA_ARCH_ATTR_OPTIONS --dlopen=%mlir_native_utils_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
 
 #DCSR = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>
 
 // An example of a quantized sparse matmul. With the zero offset for the
 // sparse input, the sparse compiler generates very efficient code for the
 //      x(i,j) += (ext(a(i,k)) - 2) * ext(b(k,j))
 // operation.
 module {

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions_vla.mlir

+//
+// If SVE is available, test VLA vectorization.
+//
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=8 enable-vla-vectorization=%ENABLE_VLA" | \
+// RUN: mlir-translate -mlir-to-llvmir | \
+// RUN: %lli --entry-function=entry %VLA_ARCH_ATTR_OPTIONS --dlopen=%mlir_native_utils_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
+#SV = #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>
+#DV = #sparse_tensor.encoding<{ dimLevelType = [ "dense"      ] }>
+
+#trait_reduction = {
+  indexing_maps = [
+    affine_map<(i) -> (i)>,  // a
+    affine_map<(i) -> ()>    // x (scalar out)
+  ],
+  iterator_types = ["reduction"],
+  doc = "x += OPER_i a(i)"
+}
+
+// An example of vector reductions.
+module {
+
+  func.func @sum_reduction_i32(%arga: tensor<32xi32, #SV>,
+                          %argx: tensor<i32>) -> tensor<i32> {
+    %0 = linalg.generic #trait_reduction
+      ins(%arga: tensor<32xi32, #SV>)
+      outs(%argx: tensor<i32>) {
+        ^bb(%a: i32, %x: i32):
+          %0 = arith.addi %x, %a : i32
+          linalg.yield %0 : i32
+    } -> tensor<i32>
+    return %0 : tensor<i32>
+  }
+
+  func.func @sum_reduction_f32(%arga: tensor<32xf32, #SV>,
+                          %argx: tensor<f32>) -> tensor<f32> {
+    %0 = linalg.generic #trait_reduction
+      ins(%arga: tensor<32xf32, #SV>)
+      outs(%argx: tensor<f32>) {
+        ^bb(%a: f32, %x: f32):
+          %0 = arith.addf %x, %a : f32
+          linalg.yield %0 : f32
+    } -> tensor<f32>
+    return %0 : tensor<f32>
+  }
+
+  func.func @and_reduction_i32(%arga: tensor<32xi32, #DV>,
+                          %argx: tensor<i32>) -> tensor<i32> {
+    %0 = linalg.generic #trait_reduction
+      ins(%arga: tensor<32xi32, #DV>)
+      outs(%argx: tensor<i32>) {
+        ^bb(%a: i32, %x: i32):
+          %0 = arith.andi %x, %a : i32
+          linalg.yield %0 : i32
+    } -> tensor<i32>
+    return %0 : tensor<i32>
+  }
+
+  func.func @or_reduction_i32(%arga: tensor<32xi32, #SV>,
+                         %argx: tensor<i32>) -> tensor<i32> {
+    %0 = linalg.generic #trait_reduction
+      ins(%arga: tensor<32xi32, #SV>)
+      outs(%argx: tensor<i32>) {
+        ^bb(%a: i32, %x: i32):
+          %0 = arith.ori %x, %a : i32
+          linalg.yield %0 : i32
+    } -> tensor<i32>
+    return %0 : tensor<i32>
+  }
+
+  func.func @xor_reduction_i32(%arga: tensor<32xi32, #SV>,
+                          %argx: tensor<i32>) -> tensor<i32> {
+    %0 = linalg.generic #trait_reduction
+      ins(%arga: tensor<32xi32, #SV>)
+      outs(%argx: tensor<i32>) {
+        ^bb(%a: i32, %x: i32):
+          %0 = arith.xori %x, %a : i32
+          linalg.yield %0 : i32
+    } -> tensor<i32>
+    return %0 : tensor<i32>
+  }
+
+  func.func @dump_i32(%arg0 : memref<i32>) {
+    %v = memref.load %arg0[] : memref<i32>
+    vector.print %v : i32
+    return
+  }
+
+  func.func @dump_f32(%arg0 : memref<f32>) {
+    %v = memref.load %arg0[] : memref<f32>
+    vector.print %v : f32
+    return
+  }
+
+  func.func @entry() {
+    %ri = arith.constant dense< 7   > : tensor<i32>
+    %rf = arith.constant dense< 2.0 > : tensor<f32>
+
+    %c_0_i32 = arith.constant dense<[
+      0, 2, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 4, 0, 0, 0,
+      0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 9, 0
+    ]> : tensor<32xi32>
+
+    %c_0_f32 = arith.constant dense<[
+      0.0, 1.0, 0.0, 0.0, 4.0, 0.0, 0.0, 0.0,
+      0.0, 0.0, 3.0, 0.0, 0.0, 0.0, 0.0, 0.0,
+      0.0, 0.0, 0.0, 0.0, 2.5, 0.0, 0.0, 0.0,
+      2.0, 0.0, 0.0, 0.0, 0.0, 4.0, 0.0, 9.0
+    ]> : tensor<32xf32>
+
+    %c_1_i32 = arith.constant dense<[
+      1, 1, 7, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
+      1, 1, 1, 1, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1, 7, 3
+    ]> : tensor<32xi32>
+
+    %c_1_f32 = arith.constant dense<[
+      1.0, 1.0, 1.0, 3.5, 1.0, 1.0, 1.0, 1.0,
+      1.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0,
+      1.0, 1.0, 1.0, 1.0, 3.0, 1.0, 1.0, 1.0,
+      1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 4.0
+    ]> : tensor<32xf32>
+
+    // Convert constants to annotated tensors.
+    %sparse_input_i32 = sparse_tensor.convert %c_0_i32
+      : tensor<32xi32> to tensor<32xi32, #SV>
+    %sparse_input_f32 = sparse_tensor.convert %c_0_f32
+      : tensor<32xf32> to tensor<32xf32, #SV>
+    %dense_input_i32 = sparse_tensor.convert %c_1_i32
+      : tensor<32xi32> to tensor<32xi32, #DV>
+    %dense_input_f32 = sparse_tensor.convert %c_1_f32
+      : tensor<32xf32> to tensor<32xf32, #DV>
+
+    // Call the kernels.
+    %0 = call @sum_reduction_i32(%sparse_input_i32, %ri)
+       : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>
+    %1 = call @sum_reduction_f32(%sparse_input_f32, %rf)
+       : (tensor<32xf32, #SV>, tensor<f32>) -> tensor<f32>
+    %4 = call @and_reduction_i32(%dense_input_i32, %ri)
+       : (tensor<32xi32, #DV>, tensor<i32>) -> tensor<i32>
+    %5 = call @or_reduction_i32(%sparse_input_i32, %ri)
+       : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>
+    %6 = call @xor_reduction_i32(%sparse_input_i32, %ri)
+       : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>
+
+    // Verify results.
+    //
+    // CHECK: 26
+    // CHECK: 27.5
+    // CHECK: 1
+    // CHECK: 15
+    // CHECK: 10
+    //
+    %m0 = bufferization.to_memref %0 : memref<i32>
+    call @dump_i32(%m0) : (memref<i32>) -> ()
+    %m1 = bufferization.to_memref %1 : memref<f32>
+    call @dump_f32(%m1) : (memref<f32>) -> ()
+    %m4 = bufferization.to_memref %4 : memref<i32>
+    call @dump_i32(%m4) : (memref<i32>) -> ()
+    %m5 = bufferization.to_memref %5 : memref<i32>
+    call @dump_i32(%m5) : (memref<i32>) -> ()
+    %m6 = bufferization.to_memref %6 : memref<i32>
+    call @dump_i32(%m6) : (memref<i32>) -> ()
+
+    // Release the resources.
+    sparse_tensor.release %sparse_input_i32 : tensor<32xi32, #SV>
+    sparse_tensor.release %sparse_input_f32 : tensor<32xf32, #SV>
+    sparse_tensor.release %dense_input_i32  : tensor<32xi32, #DV>
+    sparse_tensor.release %dense_input_f32  : tensor<32xf32, #DV>
+    memref.dealloc %m0 : memref<i32>
+    memref.dealloc %m1 : memref<f32>
+    memref.dealloc %m4 : memref<i32>
+    memref.dealloc %m5 : memref<i32>
+    memref.dealloc %m6 : memref<i32>
+
+    return
+  }
+}

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_matmul.mlir

 // RUN: mlir-opt %s \
 // RUN:   --sparse-compiler="vectorization-strategy=2 vl=4 enable-simd-index32" | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/test.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
+// If SVE is available, test VLA vectorization.
+//
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=4 enable-vla-vectorization=%ENABLE_VLA enable-simd-index32" | \
+// RUN: mlir-translate -mlir-to-llvmir | \
+// RUN: TENSOR0="%mlir_integration_test_dir/data/test.mtx" \
+// RUN: %lli --entry-function=entry %VLA_ARCH_ATTR_OPTIONS --dlopen=%mlir_native_utils_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
 
 !Filename = !llvm.ptr<i8>
 
 #SparseMatrix = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed" ],
   pointerBitWidth = 32,
   indexBitWidth = 32
 }>

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
 // RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=8" | \
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
+//
+// If SVE is available, test VLA vectorization.
+//
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=8 enable-vla-vectorization=%ENABLE_VLA" | \
+// RUN: mlir-translate -mlir-to-llvmir | \
+// RUN: %lli --entry-function=entry %VLA_ARCH_ATTR_OPTIONS --dlopen=%mlir_native_utils_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
 
 #SM = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>
 
 #trait_sampled_dense_dense = {
   indexing_maps = [
     affine_map<(i,j,k) -> (i,j)>,  // S
     affine_map<(i,j,k) -> (i,k)>,  // A
     affine_map<(i,j,k) -> (k,j)>,  // B

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_scale.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
 // RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=4" | \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
+//
+// If SVE is available, test VLA vectorization.
+//
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=4 enable-vla-vectorization=%ENABLE_VLA" | \
+// RUN: mlir-translate -mlir-to-llvmir | \
+// RUN: %lli --entry-function=entry %VLA_ARCH_ATTR_OPTIONS --dlopen=%mlir_native_utils_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
 
 #CSR = #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>
 
 #trait_scale = {
   indexing_maps = [
     affine_map<(i,j) -> (i,j)>   // X (out)
   ],
   iterator_types = ["parallel", "parallel"],

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_spmm.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/wide.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
 // RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=2" | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/wide.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
+//
+// If SVE is available, test VLA vectorization.
+//
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=2 enable-vla-vectorization=%ENABLE_VLA" | \
+// RUN: mlir-translate -mlir-to-llvmir | \
+// RUN: TENSOR0="%mlir_integration_test_dir/data/wide.mtx" \
+// RUN: %lli --entry-function=entry %VLA_ARCH_ATTR_OPTIONS --dlopen=%mlir_native_utils_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
 
 !Filename = !llvm.ptr<i8>
 
 #SparseMatrix = #sparse_tensor.encoding<{
   dimLevelType = [ "dense", "compressed" ]
 }>
 
 #spmm = {

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/test_symmetric.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
 // RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=2" | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/test_symmetric.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
+//
+// If SVE is available, test VLA vectorization.
+//
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=2 enable-vla-vectorization=%ENABLE_VLA" | \
+// RUN: mlir-translate -mlir-to-llvmir | \
+// RUN: TENSOR0="%mlir_integration_test_dir/data/test_symmetric.mtx" \
+// RUN: %lli --entry-function=entry %VLA_ARCH_ATTR_OPTIONS --dlopen=%mlir_native_utils_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
 
 !Filename = !llvm.ptr<i8>
 
 #SparseMatrix = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed" ]
 }>
 
 #trait_sum_reduce = {