[mlir][SparseTensor][ArmSVE] Conditionally disable SVE RUN line

[awarzynski]

Summary

This patch updates one SparseTensor integration test so that the VLA
vectorisation is run conditionally based on the value of the
MLIR_RUN_ARM_SME_TESTS CMake variable.

This change opens the path to reduce the duplication of RUN lines in
"mlir/test/Integration/Dialect/SparseTensor/CPU/". ATM, there are
usually 2 RUN lines to test vectorization in SparseTensor integration
tests:

• one for VLS vectorisation,
• one for VLA vectorisation whenever that's available and which reduces to VLS vectorisation when VLA is not supported.

When VLA is not available, VLS vectorisation is verified twice. This
duplication should be avoided - integration test are relatively
expansive to run.

This patch makes sure that the 2nd vectorisation RUN line becomes:

if (SVE integration tests are enabled)
  run VLA vectorisation
else
  return

This logic is implemented using LIT's (relatively new) conditional
substitution [1]. It enables us to guarantee that all RUN lines are
unique and that the VLA vectorisation is only enabled when supported.

This patch updates only 1 test to set-up and to demonstrate the logic.
Subsequent patches will update the remaining tests.

[1] https://www.llvm.org/docs/TestingGuide.html

[c-rhodes]

Thanks Andrzej, this is nice improvement

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

20

the vl has changed for the VLA test, does this matter?

mlir/test/lit.site.cfg.py.in

42

Just curious, is the idea here to support %if %ENABLE_VLA ...?

[awarzynski]

Thanks for reviewing!

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

20

Not really. The intention was always for both vectorisations to be identical (*), but things got out of sync during one of many copy & paste exercises :)

(*) The only intended difference is VLA vs VLS.

mlir/test/lit.site.cfg.py.in

42

Not quite.

The goal is to remove %ENABLE_VLA altogether (once all tests have been ported) and to use LIT's conditional substitution altogether. This brings us to:

%if <feature>

I wouldn't use ENABLE_VLA for feature, because that's not a ... feature. But there's also an option to update LIT's TestRunner to allow other things (i.e. not only "features") in <feature>. Also, perhaps quoting the docs will help:

%if feature %{<if branch>%} %else %{<else branch>%}

    Conditional substitution: if feature is available it expands to <if branch>, otherwise it expands to <else branch>. %else %{<else branch>%} is optional and treated like %else %{%} if not present.

Does this answer your question?

[c-rhodes]

LGTM cheers

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

20

Not really. The intention was always for both vectorisations to be identical (*), but things got out of sync during one of many copy & paste exercises :)

(*) The only intended difference is VLA vs VLS.

No worries, cheers!

mlir/test/lit.site.cfg.py.in

42

Not quite.

The goal is to remove %ENABLE_VLA altogether (once all tests have been ported) and to use LIT's conditional substitution altogether. This brings us to:

%if <feature>

I wouldn't use ENABLE_VLA for feature, because that's not a ... feature. But there's also an option to update LIT's TestRunner to allow other things (i.e. not only "features") in <feature>.

I see, so it could query config rather than config.available_features, where it could find mlir_run_arm_sve_tests, which could then be used as the <cond>.

Does this answer your question?

It does, thanks.

[awarzynski]

mlir/test/lit.site.cfg.py.in

42

I see, so it could query config rather than config.available_features, where it could find mlir_run_arm_sve_tests, which could then be used as the <cond>.

Indeed. Just not sure whether that would be the right design 🤔 .

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

+// DEFINE: %{option_vec} = 
 // DEFINE: %{option} = enable-runtime-library=true
+
 // DEFINE: %{compile} = mlir-opt %s --sparse-compiler=%{option}
 // DEFINE: %{run} = mlir-cpu-runner \
 // DEFINE:  -e entry -entry-point-result=void  \
 // DEFINE:  -shared-libs=%mlir_c_runner_utils | \
 // DEFINE: FileCheck %s
 //
 // RUN: %{compile} | %{run}
 //
 // Do the same run, but now with direct IR generation.
 // REDEFINE: %{option} = enable-runtime-library=false
 // RUN: %{compile} | %{run}
 //
 // Do the same run, but now with direct IR generation and vectorization.
-// REDEFINE: %{option} = "enable-runtime-library=false vl=2 reassociate-fp-reductions=true enable-index-optimizations=true"
+// REDEFINE: %{option_vec} = enable-runtime-library=false vl=2 reassociate-fp-reductions=true enable-index-optimizations=true
+// REDEFINE: %{option} = "%{option_vec}"
 // RUN: %{compile} | %{run}
 
-// Do the same run, but now with direct IR generation and, if available, VLA
-// vectorization.
-// REDEFINE: %{option} = "enable-runtime-library=false vl=4 enable-arm-sve=%ENABLE_VLA"
+// Do the same run, but with VLA vectorization.
+// REDEFINE: %{option} = "enable-arm-sve=true %{option_vec}"
 // REDEFINE: %{run} = %lli_host_or_aarch64_cmd \
 // REDEFINE:   --entry-function=entry_lli \
 // REDEFINE:   --extra-module=%S/Inputs/main_for_lli.ll \
 // REDEFINE:   %VLA_ARCH_ATTR_OPTIONS \
 // REDEFINE:   --dlopen=%mlir_native_utils_lib_dir/libmlir_c_runner_utils%shlibext | \
 // REDEFINE: FileCheck %s
-// RUN: %{compile} | mlir-translate -mlir-to-llvmir | %{run}
+// RUN: %if mlir_arm_sve_tests %{ %{compile} | mlir-translate -mlir-to-llvmir | %{run} %}
 
 #SparseVector = #sparse_tensor.encoding<{
   lvlTypes = ["compressed"]
 }>
 
 #SparseMatrix = #sparse_tensor.encoding<{
   lvlTypes = ["compressed", "compressed"]
 }>

mlir/test/lit.site.cfg.py.in

 config.enable_rocm_runner = @MLIR_ENABLE_ROCM_RUNNER@
 config.rocm_test_chipset = "@ROCM_TEST_CHIPSET@"
 config.enable_spirv_cpu_runner = @MLIR_ENABLE_SPIRV_CPU_RUNNER@
 config.enable_vulkan_runner = @MLIR_ENABLE_VULKAN_RUNNER@
 config.enable_bindings_python = @MLIR_ENABLE_BINDINGS_PYTHON@
 config.intel_sde_executable = "@INTEL_SDE_EXECUTABLE@"
 config.mlir_run_amx_tests = @MLIR_RUN_AMX_TESTS@
 config.mlir_run_arm_sve_tests = @MLIR_RUN_ARM_SVE_TESTS@
+# This is a workaround for the fact that LIT's:
+#   %if <cond>
+# requires <cond> to be in the set of available features. 
+# TODO: Update LIT's TestRunner so that this is not required.
+if config.mlir_run_arm_sve_tests:
+    config.available_features.add("mlir_arm_sve_tests")
 config.mlir_run_arm_sme_tests = @MLIR_RUN_ARM_SME_TESTS@
 config.mlir_run_x86vector_tests = @MLIR_RUN_X86VECTOR_TESTS@
 config.mlir_run_riscv_vector_tests = "@MLIR_RUN_RISCV_VECTOR_TESTS@"
 config.mlir_run_cuda_tensor_core_tests = @MLIR_RUN_CUDA_TENSOR_CORE_TESTS@
 config.mlir_run_cuda_sm80_tests = @MLIR_RUN_CUDA_SM80_TESTS@
 config.mlir_run_cuda_sm80_lt_tests = @MLIR_RUN_CUDA_SM80_LT_TESTS@
 config.mlir_include_integration_tests = @MLIR_INCLUDE_INTEGRATION_TESTS@
 config.arm_emulator_executable = "@ARM_EMULATOR_EXECUTABLE@"