[mlir] Add mlir_arm_runner_utils library for use in integration tests (#78583)

This adds a new `mlir_arm_runner_utils` library that contains utils
specific to Arm/AArch64. This is for use in MLIR integration tests.

This initial patch adds `setArmVLBits()` and `setArmSVLBits()`. This
allows changing vector length or streaming vector length at runtime (or
setting it to a known minimum, i.e. 128-bits).

Author: MacDue

mlir/lib/ExecutionEngine/ArmRunnerUtils.cpp

@@ -0,0 +1,69 @@
+//===- ArmRunnerUtils.cpp - Utilities for configuring architecture properties //
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Support/MathExtras.h"
+#include <iostream>
+#include <stdint.h>
+#include <string_view>
+
+#if (defined(_WIN32) || defined(__CYGWIN__))
+#define MLIR_ARMRUNNERUTILS_EXPORTED __declspec(dllexport)
+#else
+#define MLIR_ARMRUNNERUTILS_EXPORTED __attribute__((visibility("default")))
+#endif
+
+#ifdef __linux__
+#include <sys/prctl.h>
+#endif
+
+extern "C" {
+
+// Defines for prctl() calls. These may not necessarily exist in the host
+// <sys/prctl.h>, but will still be useable under emulation.
+//
+// https://www.kernel.org/doc/html/v5.3/arm64/sve.html#prctl-extensions
+#ifndef PR_SVE_SET_VL
+#define PR_SVE_SET_VL 50
+#endif
+// https://docs.kernel.org/arch/arm64/sme.html#prctl-extensions
+#ifndef PR_SME_SET_VL
+#define PR_SME_SET_VL 63
+#endif
+// Note: This mask is the same as both PR_SME_VL_LEN_MASK and
+// PR_SVE_VL_LEN_MASK.
+#define PR_VL_LEN_MASK 0xffff
+
+static void setArmVectorLength(std::string_view helper_name, int option,
+                               uint32_t bits) {
+#if defined(__linux__) && defined(__aarch64__)
+  if (bits < 128 || bits > 2048 || !llvm::isPowerOf2_32(bits)) {
+    std::cerr << "[error] Attempted to set an invalid vector length (" << bits
+              << "-bit)" << std::endl;
+    abort();
+  }
+  uint32_t vl = bits / 8;
+  if (auto ret = prctl(option, vl & PR_VL_LEN_MASK); ret < 0) {
+    std::cerr << "[error] prctl failed (" << ret << ")" << std::endl;
+    abort();
+  }
+#else
+  std::cerr << "[error] " << helper_name << " is unsupported" << std::endl;
+  abort();
+#endif
+}
+
+/// Sets the SVE vector length (in bits) to `bits`.
+void MLIR_ARMRUNNERUTILS_EXPORTED setArmVLBits(uint32_t bits) {
+  setArmVectorLength(__func__, PR_SVE_SET_VL, bits);
+}
+
+/// Sets the SME streaming vector length (in bits) to `bits`.
+void MLIR_ARMRUNNERUTILS_EXPORTED setArmSVLBits(uint32_t bits) {
+  setArmVectorLength(__func__, PR_SME_SET_VL, bits);
+}
+}

mlir/lib/ExecutionEngine/CMakeLists.txt

@@ -2,6 +2,7 @@
 # is a big dependency which most don't need.
 
 set(LLVM_OPTIONAL_SOURCES
+  ArmRunnerUtils.cpp
   ArmSMEStubs.cpp
   AsyncRuntime.cpp
   CRunnerUtils.cpp
@@ -186,6 +187,10 @@ if(LLVM_ENABLE_PIC)
     ArmSMEStubs.cpp)
   target_compile_definitions(mlir_arm_sme_abi_stubs PRIVATE mlir_arm_sme_abi_stubs_EXPORTS)
 
+  add_mlir_library(mlir_arm_runner_utils
+    SHARED
+    ArmRunnerUtils.cpp)
+
   if(MLIR_ENABLE_CUDA_RUNNER)
     # Configure CUDA support. Using check_language first allows us to give a
     # custom error message.

mlir/test/CMakeLists.txt

@@ -153,6 +153,10 @@ if (MLIR_RUN_ARM_SME_TESTS AND NOT ARM_SME_ABI_ROUTINES_SHLIB)
   list(APPEND MLIR_TEST_DEPENDS mlir_arm_sme_abi_stubs)
 endif()
 
+if (MLIR_RUN_ARM_SVE_TESTS OR MLIR_RUN_ARM_SME_TESTS)
+  list(APPEND MLIR_TEST_DEPENDS mlir_arm_runner_utils)
+endif()
+
 list(APPEND MLIR_TEST_DEPENDS MLIRUnitTests)
 
 if(LLVM_BUILD_EXAMPLES)

mlir/test/Integration/Dialect/Vector/CPU/ArmSME/Emulated/lit.local.cfg

@@ -0,0 +1,5 @@
+# The tests in this folder assume full control of the hardware features, such as
+# the vector length, so must be run under an emulator.
+
+if not config.arm_emulator_executable:
+    config.unsupported = True

mlir/test/Integration/Dialect/Vector/CPU/ArmSME/Emulated/test-setArmSVLBits.mlir

@@ -0,0 +1,78 @@
+// DEFINE: %{entry_point} = main
+// DEFINE: %{compile} = mlir-opt %s -convert-arm-sme-to-llvm -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,%mlir_arm_runner_utils
+
+// RUN: %{compile} | %{run} | FileCheck %s
+
+func.func @checkSVL() {
+  %svl_b = arm_sme.streaming_vl <byte>
+  %svl_h = arm_sme.streaming_vl <half>
+  %svl_w = arm_sme.streaming_vl <word>
+  %svl_d = arm_sme.streaming_vl <double>
+  vector.print str "SVL.b"
+  vector.print %svl_b : index
+  vector.print str "SVL.h"
+  vector.print %svl_h : index
+  vector.print str "SVL.w"
+  vector.print %svl_w : index
+  vector.print str "SVL.d"
+  vector.print %svl_d : index
+  return
+}
+
+func.func @setAndCheckSVL(%bits: i32) {
+  func.call @setArmSVLBits(%bits) : (i32) -> ()
+  func.call @checkSVL() : () -> ()
+  return
+}
+
+func.func @main() {
+  //      CHECK: SVL.b
+  // CHECK-NEXT: 16
+  //
+  // CHECK-NEXT: SVL.h
+  // CHECK-NEXT: 8
+  //
+  // CHECK-NEXT: SVL.w
+  // CHECK-NEXT: 4
+  //
+  // CHECK-NEXT: SVL.d
+  // CHECK-NEXT: 2
+  %c128 = arith.constant 128 : i32
+  func.call @setAndCheckSVL(%c128) : (i32) -> ()
+
+  //      CHECK: SVL.b
+  // CHECK-NEXT: 32
+  //
+  // CHECK-NEXT: SVL.h
+  // CHECK-NEXT: 16
+  //
+  // CHECK-NEXT: SVL.w
+  // CHECK-NEXT: 8
+  //
+  // CHECK-NEXT: SVL.d
+  // CHECK-NEXT: 4
+  %c256 = arith.constant 256 : i32
+  func.call @setAndCheckSVL(%c256) : (i32) -> ()
+
+  //      CHECK: SVL.b
+  // CHECK-NEXT: 64
+  //
+  // CHECK-NEXT: SVL.h
+  // CHECK-NEXT: 32
+  //
+  // CHECK-NEXT: SVL.w
+  // CHECK-NEXT: 16
+  //
+  // CHECK-NEXT: SVL.d
+  // CHECK-NEXT: 8
+  %c512 = arith.constant 512 : i32
+  func.call @setAndCheckSVL(%c512) : (i32) -> ()
+
+  return
+}
+
+func.func private @setArmSVLBits(%bits : i32)

mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/Emulated/lit.local.cfg

@@ -0,0 +1,5 @@
+# The tests in this folder assume full control of the hardware features, such as
+# the vector length, so must be run under an emulator.
+
+if not config.arm_emulator_executable:
+    config.unsupported = True

mlir/test/Integration/Dialect/Vector/CPU/ArmSVE/Emulated/test-setArmVLBits.mlir

@@ -0,0 +1,51 @@
+// DEFINE: %{entry_point} = main
+// DEFINE: %{compile} = mlir-opt %s -test-lower-to-llvm
+// DEFINE: %{run} = %mcr_aarch64_cmd -march=aarch64 -mattr=+sve \
+// DEFINE:  -e %{entry_point} -entry-point-result=void \
+// DEFINE:  -shared-libs=%mlir_runner_utils,%mlir_c_runner_utils,%mlir_arm_runner_utils
+
+// RUN: %{compile} | %{run} | FileCheck %s
+
+func.func @checkVScale() {
+  %vscale = vector.vscale
+  vector.print str "vscale"
+  vector.print %vscale : index
+  return
+}
+
+func.func @setAndCheckVL(%bits: i32) {
+  func.call @setArmVLBits(%bits) : (i32) -> ()
+  func.call @checkVScale() : () -> ()
+  return
+}
+
+func.func @main() {
+  //      CHECK: vscale
+  // CHECK-NEXT: 1
+  %c128 = arith.constant 128 : i32
+  func.call @setAndCheckVL(%c128) : (i32) -> ()
+
+  //      CHECK: vscale
+  // CHECK-NEXT: 2
+  %c256 = arith.constant 256 : i32
+  func.call @setAndCheckVL(%c256) : (i32) -> ()
+
+  //      CHECK: vscale
+  // CHECK-NEXT: 4
+  %c512 = arith.constant 512 : i32
+  func.call @setAndCheckVL(%c512) : (i32) -> ()
+
+  //      CHECK: vscale
+  // CHECK-NEXT: 8
+  %c1024 = arith.constant 1024 : i32
+  func.call @setAndCheckVL(%c1024) : (i32) -> ()
+
+  //      CHECK: vscale
+  // CHECK-NEXT: 16
+  %c2048 = arith.constant 2048 : i32
+  func.call @setAndCheckVL(%c2048) : (i32) -> ()
+
+  return
+}
+
+func.func private @setArmVLBits(%bits : i32)

mlir/test/lit.cfg.py

@@ -135,6 +135,9 @@ def add_runtime(name):
 if config.enable_sycl_runner:
     tools.extend([add_runtime("mlir_sycl_runtime")])
 
+if config.mlir_run_arm_sve_tests or config.mlir_run_arm_sme_tests:
+    tools.extend([add_runtime("mlir_arm_runner_utils")])
+
 if config.mlir_run_arm_sme_tests:
     config.substitutions.append(
         (