[InlineCost][TargetTransformInfo][AMDGPU] Consider cost of alloca instructions in the caller (1/2)

[jmmartinez]

Summary

On AMDGPU, alloca instructions have penalty that can
be avoided when SROA is applied after inlining.

This patch introduces the default implementation of
TargetTransformInfo::getCallerAllocaCost.

[jmmartinez]

• Rebase

[mtrofin]

I'm missing the AMDGPU bit of this patch, or maybe I'm not understanding something?

Is it possible to add a small unit test, maybe in unittests/Analysis/InlineCostTest.cpp?

[jmmartinez]

In D149740#4387303, @mtrofin wrote:

I'm missing the AMDGPU bit of this patch, or maybe I'm not understanding something?

Sorry, I put it on the patch stack: https://reviews.llvm.org/D149741

Is it possible to add a small unit test, maybe in unittests/Analysis/InlineCostTest.cpp?

Will do! Thanks!

[jmmartinez]

• I updated InlineCostFeatureAnalyzer::onInitializeSROAArg to match InlineCostCallAnalyzer::onInitializeSROAArg: I'm not sure this is really desired. I'd expect the ML model to figure out the cost by itself. Maybe it deserves an extra "inline-cost feature" to make the distinction (e.g. like sroa_alloca_cost or sroa_alloca_size)?

• I've added a unittest in InlineCostTest.cpp. The unittest already present in InlineCostTest.cpp only tested if getInliningCostFeatures does provide a result. I've added a test to check the cost calculated by a small example is correct. The test would be better with lit + filecheck though (I have yet to figure how that's done for InlineCostFeaturesAnalyzer).

[jmmartinez]

Ping !

[mtrofin]

llvm/unittests/Analysis/InlineCostTest.cpp

220

nit: is there something that we can test that's not 0 (to distinguish from "no effect"). Like maybe the total cost?

[jmmartinez]

llvm/unittests/Analysis/InlineCostTest.cpp

220

I think that would make the test fragile, since any change in the way the cost is calculated would break the test despite not changing how the sroa savings/losses are computed.

Would it be good if I add a test for the case where the sroa cost is lost?

[mtrofin]

llvm/unittests/Analysis/InlineCostTest.cpp

220

Not sure what the concern is - if folks change the cost function (not very frequent), a test here would highlight to them whether that side-effect is desired, which is a good thing. Assuming it is, updating should be as easy as updating the number - which is also a good thing, when tracking down regressions, it offers an additional signal ("here's one effect of this change").

Adding an additional test for where the sroa cost is lost is good, too.

[jmmartinez]

llvm/unittests/Analysis/InlineCostTest.cpp

220

Adding an additional test for where the sroa cost is lost is good, too.

Will do!

Not sure what the concern is - if folks change the cost function (not very frequent), a test here would highlight to them whether that side-effect is desired, which is a good thing. Assuming it is, updating should be as easy as updating the number - which is also a good thing, when tracking down regressions, it offers an additional signal ("here's one effect of this change").

I've just realized that there is no "total cost" equivalent in the InliningCostFeatures. The only values affected by this modification are the sroa savings and losses. These costs do not get accumulated in a single value since they are used to feed the MLInlineAdvisor. Or am I missing something?

[mtrofin]

llvm/lib/Analysis/InlineCost.cpp

720–722

nit: can you do this in 2 steps, easier to read. same in onInitializeSROAArg. thanks

llvm/unittests/Analysis/InlineCostTest.cpp

220

SROAArgCosts is used for total cost estimation - see onDisableSROA

[jmmartinez]

• Added sroa_losses unittest
• Refactored common code in unittest

[jmmartinez]

llvm/unittests/Analysis/InlineCostTest.cpp

220

There is InlineCostCallAnalyzer::onDisableSROA, where the cost is used for the total estimation,

But there is the very similar InlineCostFeatureAnalyzer::onDisableSROA where the cost is only used for the SROA savings/losses estimation.

In the unittests the second one is getting tested, where there is no total cost estimation.

[jmmartinez]

• Updated one place where I forgot to split the increment of the sroa cost in two lines.

[mtrofin]

lgtm

llvm/include/llvm/Analysis/TargetTransformInfo.h

 #include <utility>
 
 namespace llvm {
 
 namespace Intrinsic {
 typedef unsigned ID;
 }
 
+class AllocaInst;
 class AssumptionCache;
 class BlockFrequencyInfo;
 class DominatorTree;
 class BranchInst;
 class CallBase;
 class Function;
 class GlobalValue;
 class InstCombiner;
public:
   ///
   /// TODO: This is a rather blunt instrument.  Perhaps altering the costs of
   /// individual classes of instructions would be better.
   unsigned getInliningThresholdMultiplier() const;
 
   /// \returns A value to be added to the inlining threshold.
   unsigned adjustInliningThreshold(const CallBase *CB) const;
 
+  /// \returns The cost of having an Alloca in the caller if not inlined, to be
+  /// added to the threshold
+  unsigned getCallerAllocaCost(const CallBase *CB, const AllocaInst *AI) const;
+
   /// \returns Vector bonus in percent.
   ///
   /// Vector bonuses: We want to more aggressively inline vector-dense kernels
   /// and apply this bonus based on the percentage of vector instructions. A
   /// bonus is applied if the vector instructions exceed 50% and half that
   /// amount is applied if it exceeds 10%. Note that these bonuses are some what
   /// arbitrary and evolved over time by accident as much as because they are
   /// principled bonuses.
virtual InstructionCost getGEPCost(Type *PointeeType, const Value *Ptr,
                                      TTI::TargetCostKind CostKind) = 0;
   virtual InstructionCost
   getPointersChainCost(ArrayRef<const Value *> Ptrs, const Value *Base,
                        const TTI::PointersChainInfo &Info, Type *AccessTy,
                        TTI::TargetCostKind CostKind) = 0;
   virtual unsigned getInliningThresholdMultiplier() const = 0;
   virtual unsigned adjustInliningThreshold(const CallBase *CB) = 0;
   virtual int getInlinerVectorBonusPercent() const = 0;
+  virtual unsigned getCallerAllocaCost(const CallBase *CB,
+                                       const AllocaInst *AI) const = 0;
   virtual InstructionCost getMemcpyCost(const Instruction *I) = 0;
   virtual uint64_t getMaxMemIntrinsicInlineSizeThreshold() const = 0;
   virtual unsigned
   getEstimatedNumberOfCaseClusters(const SwitchInst &SI, unsigned &JTSize,
                                    ProfileSummaryInfo *PSI,
                                    BlockFrequencyInfo *BFI) = 0;
   virtual InstructionCost getInstructionCost(const User *U,
                                              ArrayRef<const Value *> Operands,
unsigned getInliningThresholdMultiplier() const override {
     return Impl.getInliningThresholdMultiplier();
   }
   unsigned adjustInliningThreshold(const CallBase *CB) override {
     return Impl.adjustInliningThreshold(CB);
   }
   int getInlinerVectorBonusPercent() const override {
     return Impl.getInlinerVectorBonusPercent();
   }
+  unsigned getCallerAllocaCost(const CallBase *CB,
+                               const AllocaInst *AI) const override {
+    return Impl.getCallerAllocaCost(CB, AI);
+  }
   InstructionCost getMemcpyCost(const Instruction *I) override {
     return Impl.getMemcpyCost(I);
   }
 
   uint64_t getMaxMemIntrinsicInlineSizeThreshold() const override {
     return Impl.getMaxMemIntrinsicInlineSizeThreshold();
   }
 

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

unsigned getEstimatedNumberOfCaseClusters(const SwitchInst &SI,
     (void)PSI;
     (void)BFI;
     JTSize = 0;
     return SI.getNumCases();
   }
 
   unsigned getInliningThresholdMultiplier() const { return 1; }
   unsigned adjustInliningThreshold(const CallBase *CB) const { return 0; }
+  unsigned getCallerAllocaCost(const CallBase *CB, const AllocaInst *AI) const {
+    return 0;
+  };
 
   int getInlinerVectorBonusPercent() const { return 150; }
 
   InstructionCost getMemcpyCost(const Instruction *I) const {
     return TTI::TCC_Expensive;
   }
 
   uint64_t getMaxMemIntrinsicInlineSizeThreshold() const {

llvm/include/llvm/CodeGen/BasicTTIImpl.h

InstructionCost getFPOpCost(Type *Ty) {
     EVT VT = TLI->getValueType(DL, Ty);
     if (TLI->isOperationLegalOrCustomOrPromote(ISD::FADD, VT))
       return TargetTransformInfo::TCC_Basic;
     return TargetTransformInfo::TCC_Expensive;
   }
 
   unsigned getInliningThresholdMultiplier() const { return 1; }
   unsigned adjustInliningThreshold(const CallBase *CB) { return 0; }
+  unsigned getCallerAllocaCost(const CallBase *CB, const AllocaInst *AI) const {
+    return 0;
+  }
 
   int getInlinerVectorBonusPercent() const { return 150; }
 
   void getUnrollingPreferences(Loop *L, ScalarEvolution &SE,
                                TTI::UnrollingPreferences &UP,
                                OptimizationRemarkEmitter *ORE) {
     // This unrolling functionality is target independent, but to provide some
     // motivation for its intended use, for x86:

llvm/lib/Analysis/InlineCost.cpp

void onFinalizeSwitch(unsigned JumpTableSize,
 
     addCost(SwitchCost);
   }
   void onMissedSimplification() override { addCost(InstrCost); }
 
   void onInitializeSROAArg(AllocaInst *Arg) override {
     assert(Arg != nullptr &&
            "Should not initialize SROA costs for null value.");
-    SROAArgCosts[Arg] = 0;
+    auto SROAArgCost = TTI.getCallerAllocaCost(&CandidateCall, Arg);
+    SROACostSavings += SROAArgCost;
+    SROAArgCosts[Arg] = SROAArgCost;
   }
 
   void onAggregateSROAUse(AllocaInst *SROAArg) override {
     auto CostIt = SROAArgCosts.find(SROAArg);
     assert(CostIt != SROAArgCosts.end() &&
            "expected this argument to have a cost");
     CostIt->second += InstrCost;
     SROACostSavings += InstrCost;
void onFinalizeSwitch(unsigned JumpTableSize,
     increment(InlineCostFeatureIndex::switch_penalty, SwitchCost);
   }
 
   void onMissedSimplification() override {
     increment(InlineCostFeatureIndex::unsimplified_common_instructions,
               InstrCost);
   }
 
-  void onInitializeSROAArg(AllocaInst *Arg) override { SROACosts[Arg] = 0; }
+  void onInitializeSROAArg(AllocaInst *Arg) override {
+    auto SROAArgCost = TTI.getCallerAllocaCost(&CandidateCall, Arg);
+    SROACosts[Arg] = SROAArgCost;
+    SROACostSavingOpportunities += SROAArgCost;
+  }
+
   void onAggregateSROAUse(AllocaInst *Arg) override {
     SROACosts.find(Arg)->second += InstrCost;
     SROACostSavingOpportunities += InstrCost;
   }
 
   void onBlockAnalyzed(const BasicBlock *BB) override {
     if (BB->getTerminator()->getNumSuccessors() > 1)
       set(InlineCostFeatureIndex::is_multiple_blocks, 1);

llvm/lib/Analysis/TargetTransformInfo.cpp

unsigned TargetTransformInfo::getInliningThresholdMultiplier() const {
   return TTIImpl->getInliningThresholdMultiplier();
 }
 
 unsigned
 TargetTransformInfo::adjustInliningThreshold(const CallBase *CB) const {
   return TTIImpl->adjustInliningThreshold(CB);
 }
 
+unsigned TargetTransformInfo::getCallerAllocaCost(const CallBase *CB,
+                                                  const AllocaInst *AI) const {
+  return TTIImpl->getCallerAllocaCost(CB, AI);
+}
+
 int TargetTransformInfo::getInlinerVectorBonusPercent() const {
   return TTIImpl->getInlinerVectorBonusPercent();
 }
 
 InstructionCost
 TargetTransformInfo::getGEPCost(Type *PointeeType, const Value *Ptr,
                                 ArrayRef<const Value *> Operands,
                                 TTI::TargetCostKind CostKind) const {

llvm/unittests/Analysis/InlineCostTest.cpp

 //===- InlineCostTest.cpp - test for InlineCost ---------------------------===//
 //
 // 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/Analysis/InlineCost.h"
 #include "llvm/Analysis/AssumptionCache.h"
+#include "llvm/Analysis/InlineModelFeatureMaps.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/SourceMgr.h"
 #include "gtest/gtest.h"
 
 namespace {
 
-// Tests that we can retrieve the CostFeatures without an error
-TEST(InlineCostTest, CostFeatures) {
-  using namespace llvm;
 
+CallBase *getCallInFunction(Function *F) {
+  for (auto &I : instructions(F)) {
+    if (auto *CB = dyn_cast<llvm::CallBase>(&I))
+      return CB;
+  }
+  return nullptr;
+}
+
+std::optional<InlineCostFeatures> getInliningCostFeaturesForCall(CallBase &CB) {
+  ModuleAnalysisManager MAM;
+  FunctionAnalysisManager FAM;
+  FAM.registerPass([&] { return TargetIRAnalysis(); });
+  FAM.registerPass([&] { return ModuleAnalysisManagerFunctionProxy(MAM); });
+  FAM.registerPass([&] { return AssumptionAnalysis(); });
+  MAM.registerPass([&] { return FunctionAnalysisManagerModuleProxy(FAM); });
+
+  MAM.registerPass([&] { return PassInstrumentationAnalysis(); });
+  FAM.registerPass([&] { return PassInstrumentationAnalysis(); });
+
+  ModulePassManager MPM;
+  MPM.run(*CB.getModule(), MAM);
+
+  auto GetAssumptionCache = [&](Function &F) -> AssumptionCache & {
+    return FAM.getResult<AssumptionAnalysis>(F);
+  };
+  auto &TIR = FAM.getResult<TargetIRAnalysis>(*CB.getFunction());
+
+  return getInliningCostFeatures(CB, TIR, GetAssumptionCache);
+}
+
+// Tests that we can retrieve the CostFeatures without an error
+TEST(InlineCostTest, CostFeatures) {
   const auto *const IR = R"IR(
 define i32 @f(i32) {
   ret i32 4
 }
 
 define i32 @g(i32) {
   %2 = call i32 @f(i32 0)
   ret i32 %2
 }
 )IR";
 
   LLVMContext C;
   SMDiagnostic Err;
   std::unique_ptr<Module> M = parseAssemblyString(IR, Err, C);
   ASSERT_TRUE(M);
 
   auto *G = M->getFunction("g");
   ASSERT_TRUE(G);
 
   // find the call to f in g
-  CallBase *CB = nullptr;
-  for (auto &BB : *G) {
-    for (auto &I : BB) {
-      if ((CB = dyn_cast<CallBase>(&I)))
-        break;
-    }
-  }
-  ASSERT_TRUE(CB);
+  CallBase *CB = getCallInFunction(G);
+  ASSERT_TRUE(CB);
+
+  const auto Features = getInliningCostFeaturesForCall(*CB);
+
+  // Check that the optional is not empty
+  ASSERT_TRUE(Features);
+}
+
+// Tests the calculated SROA cost
+TEST(InlineCostTest, SROACost) {
+  using namespace llvm;
+
+  const auto *const IR = R"IR(
+define void @f_savings(ptr %var) {
+  %load = load i32, ptr %var
+  %inc = add i32 %load, 1
+  store i32 %inc, ptr %var
+  ret void
+
+define void @g_savings(i32) {
+  %var = alloca i32
+  call void @f_savings(ptr %var)
+  ret void
 
-  ModuleAnalysisManager MAM;
-  FunctionAnalysisManager FAM;
-  FAM.registerPass([&] { return TargetIRAnalysis(); });
-  FAM.registerPass([&] { return ModuleAnalysisManagerFunctionProxy(MAM); });
-  FAM.registerPass([&] { return AssumptionAnalysis(); });
-  MAM.registerPass([&] { return FunctionAnalysisManagerModuleProxy(FAM); });
+define void @f_losses(ptr %var) {
+  %load = load i32, ptr %var
+  %inc = add i32 %load, 1
+  store i32 %inc, ptr %var
+  call void @prevent_sroa(ptr %var)
+  ret void
+}
 
-  MAM.registerPass([&] { return PassInstrumentationAnalysis(); });
-  FAM.registerPass([&] { return PassInstrumentationAnalysis(); });
+define void @g_losses(i32) {
+  %var = alloca i32
+  call void @f_losses(ptr %var)
+  ret void
+}
 
-  ModulePassManager MPM;
-  MPM.run(*M, MAM);
+declare void @prevent_sroa(ptr)
+)IR";
 
-  auto GetAssumptionCache = [&](Function &F) -> AssumptionCache & {
-    return FAM.getResult<AssumptionAnalysis>(F);
-  };
-  auto &TIR = FAM.getResult<TargetIRAnalysis>(*G);
+  LLVMContext C;
+  SMDiagnostic Err;
+  std::unique_ptr<Module> M = parseAssemblyString(IR, Err, C);
+  ASSERT_TRUE(M);
 
-  const auto Features =
-      llvm::getInliningCostFeatures(*CB, TIR, GetAssumptionCache);
+  const int DefaultInstCost = 5;
+  const int DefaultAllocaCost = 0;
 
-  // Check that the optional is not empty
+  const char *GName[] = {"g_savings", "g_losses", nullptr};
+  const int Savings[] = {2 * DefaultInstCost + DefaultAllocaCost, 0};
+  const int Losses[] = {0, 2 * DefaultInstCost + DefaultAllocaCost};
+
+  for (unsigned i = 0; GName[i]; ++i) {
+    auto *G = M->getFunction(GName[i]);
+    ASSERT_TRUE(G);
+
+    // find the call to f in g
+    CallBase *CB = getCallInFunction(G);
+    ASSERT_TRUE(CB);
+
+    const auto Features = getInliningCostFeaturesForCall(*CB);
     ASSERT_TRUE(Features);
+
+    // Check the predicted SROA cost
+    auto GetFeature = [&](InlineCostFeatureIndex I) {
+      return (*Features)[static_cast<size_t>(I)];
+    };
+    ASSERT_EQ(GetFeature(InlineCostFeatureIndex::sroa_savings), Savings[i]);
+    ASSERT_EQ(GetFeature(InlineCostFeatureIndex::sroa_losses), Losses[i]);
+  }
 }
 
 } // namespace