[TTI][TLI][AArch64] Add isLegalAddScalableImmediate

Adds an interface to determine whether an immediate would be legal within
an add instruction, when said immediate is multiplied by vscale.

Author: huntergr-arm

llvm/include/llvm/Analysis/TargetTransformInfo.h

@@ -696,6 +696,12 @@ class TargetTransformInfo {
   /// immediate without having to materialize the immediate into a register.
   bool isLegalAddImmediate(int64_t Imm) const;
 
+  /// Return true if the specified immediate is legal add of a scalable
+  /// immediate, that is the target has add instructions which can add a
+  /// register with the immediate (multiplied by vscale) without having to
+  /// materialize the immediate into a register.
+  bool isLegalAddScalableImmediate(int64_t Imm) const;
+
   /// Return true if the specified immediate is legal icmp immediate,
   /// that is the target has icmp instructions which can compare a register
   /// against the immediate without having to materialize the immediate into a
@@ -1835,6 +1841,7 @@ class TargetTransformInfo::Concept {
       std::function<void(Instruction *, unsigned, APInt, APInt &)>
           SimplifyAndSetOp) = 0;
   virtual bool isLegalAddImmediate(int64_t Imm) = 0;
+  virtual bool isLegalAddScalableImmediate(int64_t Imm) = 0;
   virtual bool isLegalICmpImmediate(int64_t Imm) = 0;
   virtual bool isLegalAddressingMode(Type *Ty, GlobalValue *BaseGV,
                                      int64_t BaseOffset, bool HasBaseReg,
@@ -2295,6 +2302,9 @@ class TargetTransformInfo::Model final : public TargetTransformInfo::Concept {
   bool isLegalAddImmediate(int64_t Imm) override {
     return Impl.isLegalAddImmediate(Imm);
   }
+  bool isLegalAddScalableImmediate(int64_t Imm) override {
+    return Impl.isLegalAddScalableImmediate(Imm);
+  }
   bool isLegalICmpImmediate(int64_t Imm) override {
     return Impl.isLegalICmpImmediate(Imm);
   }

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

@@ -216,6 +216,8 @@ class TargetTransformInfoImplBase {
 
   bool isLegalAddImmediate(int64_t Imm) const { return false; }
 
+  bool isLegalAddScalableImmediate(int64_t Imm) const { return false; }
+
   bool isLegalICmpImmediate(int64_t Imm) const { return false; }
 
   bool isLegalAddressingMode(Type *Ty, GlobalValue *BaseGV, int64_t BaseOffset,

llvm/include/llvm/CodeGen/BasicTTIImpl.h

@@ -328,6 +328,10 @@ class BasicTTIImplBase : public TargetTransformInfoImplCRTPBase<T> {
     return getTLI()->isLegalAddImmediate(imm);
   }
 
+  bool isLegalAddScalableImmediate(int64_t Imm) {
+    return getTLI()->isLegalAddScalableImmediate(Imm);
+  }
+
   bool isLegalICmpImmediate(int64_t imm) {
     return getTLI()->isLegalICmpImmediate(imm);
   }

llvm/include/llvm/CodeGen/TargetLowering.h

@@ -2770,6 +2770,12 @@ class TargetLoweringBase {
     return true;
   }
 
+  /// Return true if the specified immediate is legal add of a scalable
+  /// immediate, that is the target has add instructions which can add a
+  /// register with the immediate (multiplied by vscale) without having to
+  /// materialize the immediate into a register.
+  virtual bool isLegalAddScalableImmediate(int64_t) const { return false; }
+
   /// Return true if the specified immediate is legal for the value input of a
   /// store instruction.
   virtual bool isLegalStoreImmediate(int64_t Value) const {

llvm/lib/Analysis/TargetTransformInfo.cpp

@@ -395,6 +395,10 @@ bool TargetTransformInfo::isLegalAddImmediate(int64_t Imm) const {
   return TTIImpl->isLegalAddImmediate(Imm);
 }
 
+bool TargetTransformInfo::isLegalAddScalableImmediate(int64_t Imm) const {
+  return TTIImpl->isLegalAddScalableImmediate(Imm);
+}
+
 bool TargetTransformInfo::isLegalICmpImmediate(int64_t Imm) const {
   return TTIImpl->isLegalICmpImmediate(Imm);
 }

llvm/unittests/Target/AArch64/CMakeLists.txt

@@ -29,6 +29,7 @@ add_llvm_target_unittest(AArch64Tests
   MatrixRegisterAliasing.cpp
   SMEAttributesTest.cpp
   AArch64SVESchedPseudoTest.cpp
+  Immediates.cpp
   )
 
 set_property(TARGET AArch64Tests PROPERTY FOLDER "Tests/UnitTests/TargetTests")

llvm/unittests/Target/AArch64/Immediates.cpp

@@ -0,0 +1,62 @@
+#include "AArch64Subtarget.h"
+#include "AArch64TargetMachine.h"
+#include "llvm/IR/DataLayout.h"
+#include "llvm/MC/TargetRegistry.h"
+#include "llvm/Support/TargetSelect.h"
+
+#include "gtest/gtest.h"
+#include <initializer_list>
+#include <memory>
+
+using namespace llvm;
+
+namespace {
+
+struct TestCase {
+  int64_t Imm;
+  bool Result;
+};
+
+const std::initializer_list<TestCase> Tests = {
+    // ScalableImm, Result
+    // No change, easily 'supported'
+    {0, false},
+
+    // addvl increments by whole registers, range [-32,31]
+    // +(16 * vscale), one register's worth
+    {16, false},
+    // +(8 * vscale), half a register's worth
+    {8, false},
+    // -(32 * 16 * vscale)
+    {-512, false},
+    // -(33 * 16 * vscale)
+    {-528, false},
+    // +(31 * 16 * vscale)
+    {496, false},
+    // +(32 * 16 * vscale)
+    {512, false},
+};
+} // namespace
+
+TEST(Immediates, Immediates) {
+  LLVMInitializeAArch64TargetInfo();
+  LLVMInitializeAArch64Target();
+  LLVMInitializeAArch64TargetMC();
+
+  std::string Error;
+  auto TT = Triple::normalize("aarch64");
+  const Target *T = TargetRegistry::lookupTarget(TT, Error);
+
+  std::unique_ptr<TargetMachine> TM(T->createTargetMachine(
+      TT, "generic", "+sve", TargetOptions(), std::nullopt, std::nullopt,
+      CodeGenOptLevel::Default));
+  AArch64Subtarget ST(TM->getTargetTriple(), TM->getTargetCPU(),
+                      TM->getTargetCPU(), TM->getTargetFeatureString(), *TM,
+                      true);
+
+  auto *TLI = ST.getTargetLowering();
+
+  for (const auto &Test : Tests) {
+    ASSERT_EQ(TLI->isLegalAddScalableImmediate(Test.Imm), Test.Result);
+  }
+}