mlir/lib/ExecutionEngine/ArmRunnerUtils.cpp

   1 //===- ArmRunnerUtils.cpp - Utilities for configuring architecture properties //
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8
   9 #include "llvm/Support/MathExtras.h"
  10 #include <iostream>
  11 #include <stdint.h>
  12 #include <string_view>
  13
  14 #if (defined(_WIN32) || defined(__CYGWIN__))
  15 #define MLIR_ARMRUNNERUTILS_EXPORTED __declspec(dllexport)
  16 #else
  17 #define MLIR_ARMRUNNERUTILS_EXPORTED __attribute__((visibility("default")))
  18 #endif
  19
  20 #ifdef __linux__
  21 #include <sys/prctl.h>
  22 #endif
  23
  24 extern "C" {
  25
  26 // Defines for prctl() calls. These may not necessarily exist in the host
  27 // <sys/prctl.h>, but will still be useable under emulation.
  28 //
  29 // https://www.kernel.org/doc/html/v5.3/arm64/sve.html#prctl-extensions
  30 #ifndef PR_SVE_SET_VL
  31 #define PR_SVE_SET_VL 50
  32 #endif
  33 // https://docs.kernel.org/arch/arm64/sme.html#prctl-extensions
  34 #ifndef PR_SME_SET_VL
  35 #define PR_SME_SET_VL 63
  36 #endif
  37 // Note: This mask is the same as both PR_SME_VL_LEN_MASK and
  38 // PR_SVE_VL_LEN_MASK.
  39 #define PR_VL_LEN_MASK 0xffff
  40
  41 static void setArmVectorLength(std::string_view helper_name, int option,
  42                                uint32_t bits) {
  43 #if defined(__linux__) && defined(__aarch64__)
  44   if (bits < 128 || bits > 2048 || !llvm::isPowerOf2_32(bits)) {
  45     std::cerr << "[error] Attempted to set an invalid vector length (" << bits
  46               << "-bit)" << std::endl;
  47     abort();
  48   }
  49   uint32_t vl = bits / 8;
  50   if (auto ret = prctl(option, vl & PR_VL_LEN_MASK); ret < 0) {
  51     std::cerr << "[error] prctl failed (" << ret << ")" << std::endl;
  52     abort();
  53   }
  54 #else
  55   std::cerr << "[error] " << helper_name << " is unsupported" << std::endl;
  56   abort();
  57 #endif
  58 }
  59
  60 /// Sets the SVE vector length (in bits) to `bits`.
  61 void MLIR_ARMRUNNERUTILS_EXPORTED setArmVLBits(uint32_t bits) {
  62   setArmVectorLength(__func__, PR_SVE_SET_VL, bits);
  63 }
  64
  65 /// Sets the SME streaming vector length (in bits) to `bits`.
  66 void MLIR_ARMRUNNERUTILS_EXPORTED setArmSVLBits(uint32_t bits) {
  67   setArmVectorLength(__func__, PR_SME_SET_VL, bits);
  68 }
  69 }