[AMDGPU] Propagate amdgpu-waves-per-eu to callees

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp

68

Typo Attrbutes. Also since this isn't strictly for attributes, FunctionProperties?

llvm/test/CodeGen/AMDGPU/propagate-attributes-clone.ll

141

Test where the caller and callee already have the same attribute?

[rampitec]

Addressed review comments.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp

173

This should probably be an SCC pass, and use the Attributor

[rampitec]

llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp

173

Ideally we should move it to LTO, but then we would need to run it always.

llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp

 using namespace llvm;
 
 namespace llvm {
 extern const SubtargetFeatureKV AMDGPUFeatureKV[AMDGPU::NumSubtargetFeatures-1];
 }
 
 namespace {
 
-class AMDGPUPropagateAttributes {
-  const FeatureBitset TargetFeatures = {
-    AMDGPU::FeatureWavefrontSize16,
-    AMDGPU::FeatureWavefrontSize32,
-    AMDGPU::FeatureWavefrontSize64
-  };
+// Target features to propagate.
+static constexpr const FeatureBitset TargetFeatures = {
 
-  class Clone{
+// Attributes to propagate.
+static constexpr const char* AttributeNames[] = {
+  "amdgpu-waves-per-eu"
+};
+
+static constexpr unsigned NumAttr =
+  sizeof(AttributeNames) / sizeof(AttributeNames[0]);
+
+class AMDGPUPropagateAttributes {
+
+  class FnAttrbutes {
+  private:
+    explicit FnAttrbutes(const FeatureBitset &&FB) : Features(FB) {}
+
+  public:
+    explicit FnAttrbutes(const TargetMachine &TM, const Function &F) {
+      Features = TM.getSubtargetImpl(F)->getFeatureBits();
+
+      for (unsigned I = 0; I < NumAttr; ++I)
+        if (F.hasFnAttribute(AttributeNames[I]))
+          Attributes[I] = F.getFnAttribute(AttributeNames[I]);
+    }
+
+    bool operator == (const FnAttrbutes &Other) const {
+      if ((Features & TargetFeatures) != (Other.Features & TargetFeatures))
+        return false;
+      for (unsigned I = 0; I < NumAttr; ++I)
+        if (Attributes[I] != Other.Attributes[I])
+          return false;
+      return true;
+    }
+
+    FnAttrbutes adjustToCaller(const FnAttrbutes &CallerAttrs) const {
+      FnAttrbutes New((Features & ~TargetFeatures) | CallerAttrs.Features);
+      for (unsigned I = 0; I < NumAttr; ++I)
+        New.Attributes[I] = CallerAttrs.Attributes[I];
+      return New;
+    }
+
+    FeatureBitset Features;
+    Optional<Attribute> Attributes[NumAttr];
+  };
+
+  class Clone {
   public:
-    Clone(FeatureBitset FeatureMask, Function *OrigF, Function *NewF) :
-      FeatureMask(FeatureMask), OrigF(OrigF), NewF(NewF) {}
+    Clone(const FnAttrbutes &Attrs, Function *OrigF, Function *NewF) :
+      Attributes(Attrs), OrigF(OrigF), NewF(NewF) {}
 
-    FeatureBitset FeatureMask;
+    FnAttrbutes Attributes;
     Function *OrigF;
     Function *NewF;
   };
 
   const TargetMachine *TM;
 
   // Clone functions as needed or just set attributes.
   bool AllowClone;
 
   // Option propagation roots.
   SmallSet<Function *, 32> Roots;
 
   // Clones of functions with their attributes.
   SmallVector<Clone, 32> Clones;
 
   // Find a clone with required features.
-  Function *findFunction(const FeatureBitset &FeaturesNeeded,
+  Function *findFunction(const FnAttrbutes &AttrsNeeded,
                          Function *OrigF);
 
-  // Clone function F and set NewFeatures on the clone.
+  // Clone function \p F and set \p NewAttrs on the clone.
   // Cole takes the name of original function.
-  Function *cloneWithFeatures(Function &F,
-                              const FeatureBitset &NewFeatures);
+  Function *cloneWithAttributes(Function &F, const FnAttrbutes &NewAttrs);
 
   // Set new function's features in place.
   void setFeatures(Function &F, const FeatureBitset &NewFeatures);
 
+  // Set new function's attributes in place.
+  void setAttributes(Function &F, const ArrayRef<Optional<Attribute>> NewAttrs);
+
   std::string getFeatureString(const FeatureBitset &Features) const;
 
   // Propagate attributes from Roots.
   bool process();
 
 public:
   AMDGPUPropagateAttributes(const TargetMachine *TM, bool AllowClone) :
     TM(TM), AllowClone(AllowClone) {}
initializeAMDGPUPropagateAttributesEarlyPass(
       *PassRegistry::getPassRegistry());
   }
 
   bool runOnFunction(Function &F) override;
 };
 
 // Allows to propagate attributes with clonning but does that late in the
 // pipeline.
 class AMDGPUPropagateAttributesLate : public ModulePass {
   const TargetMachine *TM;
 
 public:
   static char ID; // Pass identification
 
   AMDGPUPropagateAttributesLate(const TargetMachine *TM = nullptr) :
     ModulePass(ID), TM(TM) {
     initializeAMDGPUPropagateAttributesLatePass(
INITIALIZE_PASS(AMDGPUPropagateAttributesEarly,
                 "Early propagate attributes from kernels to functions",
                 false, false)
 INITIALIZE_PASS(AMDGPUPropagateAttributesLate,
                 "amdgpu-propagate-attributes-late",
                 "Late propagate attributes from kernels to functions",
                 false, false)
 
 Function *
-AMDGPUPropagateAttributes::findFunction(const FeatureBitset &FeaturesNeeded,
+AMDGPUPropagateAttributes::findFunction(const FnAttrbutes &AttrsNeeded,
                                         Function *OrigF) {
   // TODO: search for clone's clones.
   for (Clone &C : Clones)
-    if (C.OrigF == OrigF && FeaturesNeeded == C.FeatureMask)
+    if (C.OrigF == OrigF && AttrsNeeded == C.Attributes)
       return C.NewF;
 
   return nullptr;
 }
 
 bool AMDGPUPropagateAttributes::process(Module &M) {
   for (auto &F : M.functions())
     if (AMDGPU::isEntryFunctionCC(F.getCallingConv()))
bool AMDGPUPropagateAttributes::process() {
   do {
     Roots.insert(NewRoots.begin(), NewRoots.end());
     NewRoots.clear();
 
     for (auto &F : M.functions()) {
       if (F.isDeclaration())
         continue;
 
-      const FeatureBitset &CalleeBits =
-        TM->getSubtargetImpl(F)->getFeatureBits();
+      const FnAttrbutes CalleeAttrs(*TM, F);
       SmallVector<std::pair<CallBase *, Function *>, 32> ToReplace;
       SmallSet<CallBase *, 32> Visited;
 
       for (User *U : F.users()) {
         Instruction *I = dyn_cast<Instruction>(U);
         if (!I)
           continue;
         CallBase *CI = dyn_cast<CallBase>(I);
         if (!CI)
           continue;
         Function *Caller = CI->getCaller();
         if (!Caller || !Visited.insert(CI).second)
           continue;
         if (!Roots.count(Caller) && !NewRoots.count(Caller))
           continue;
 
-        const FeatureBitset &CallerBits =
-          TM->getSubtargetImpl(*Caller)->getFeatureBits() & TargetFeatures;
+        const FnAttrbutes CallerAttrs(*TM, *Caller);
 
-        if (CallerBits == (CalleeBits  & TargetFeatures)) {
+        if (CalleeAttrs == CallerAttrs) {
           if (!Roots.count(&F))
             NewRoots.insert(&F);
           continue;
         }
 
-        Function *NewF = findFunction(CallerBits, &F);
+        Function *NewF = findFunction(CallerAttrs, &F);
         if (!NewF) {
-          FeatureBitset NewFeatures((CalleeBits & ~TargetFeatures) |
-                                    CallerBits);
+          const FnAttrbutes NewAttrs = CalleeAttrs.adjustToCaller(CallerAttrs);
           if (!AllowClone) {
             // This may set different features on different iteartions if
             // there is a contradiction in callers' attributes. In this case
             // we rely on a second pass running on Module, which is allowed
             // to clone.
-            setFeatures(F, NewFeatures);
+            setFeatures(F, NewAttrs.Features);
+            setAttributes(F, NewAttrs.Attributes);
             NewRoots.insert(&F);
             Changed = true;
             break;
           }
 
-          NewF = cloneWithFeatures(F, NewFeatures);
-          Clones.push_back(Clone(CallerBits, &F, NewF));
+          NewF = cloneWithAttributes(F, NewAttrs);
+          Clones.push_back(Clone(CallerAttrs, &F, NewF));
           NewRoots.insert(NewF);
         }
 
         ToReplace.push_back(std::make_pair(CI, NewF));
         Replaced.insert(&F);
 
         Changed = true;
       }
bool AMDGPUPropagateAttributes::process() {
 
   Roots.clear();
   Clones.clear();
 
   return Changed;
 }
 
 Function *
-AMDGPUPropagateAttributes::cloneWithFeatures(Function &F,
-                                             const FeatureBitset &NewFeatures) {
+AMDGPUPropagateAttributes::cloneWithAttributes(Function &F,
+                                               const FnAttrbutes &NewAttrs) {
   LLVM_DEBUG(dbgs() << "Cloning " << F.getName() << '\n');
 
   ValueToValueMapTy dummy;
   Function *NewF = CloneFunction(&F, dummy);
-  setFeatures(*NewF, NewFeatures);
+  setFeatures(*NewF, NewAttrs.Features);
+  setAttributes(*NewF, NewAttrs.Attributes);
   NewF->setVisibility(GlobalValue::DefaultVisibility);
   NewF->setLinkage(GlobalValue::InternalLinkage);
 
   // Swap names. If that is the only clone it will retain the name of now
   // dead value. Preserve original name for externally visible functions.
   if (F.hasName() && F.hasLocalLinkage()) {
     std::string NewName = std::string(NewF->getName());
     NewF->takeName(&F);
void AMDGPUPropagateAttributes::setFeatures(Function &F,
   LLVM_DEBUG(dbgs() << "Set features "
                     << getFeatureString(NewFeatures & TargetFeatures)
                     << " on " << F.getName() << '\n');
 
   F.removeFnAttr("target-features");
   F.addFnAttr("target-features", NewFeatureStr);
 }
 
+void AMDGPUPropagateAttributes::setAttributes(Function &F,
+    const ArrayRef<Optional<Attribute>> NewAttrs) {
+  LLVM_DEBUG(dbgs() << "Set attributes on " << F.getName() << ":\n");
+  for (unsigned I = 0; I < NumAttr; ++I) {
+    F.removeFnAttr(AttributeNames[I]);
+    if (NewAttrs[I]) {
+      LLVM_DEBUG(dbgs() << '\t' << NewAttrs[I]->getAsString() << '\n');
+      F.addFnAttr(*NewAttrs[I]);
+    }
+  }
+}
+
 std::string
 AMDGPUPropagateAttributes::getFeatureString(const FeatureBitset &Features) const
 {
   std::string Ret;
   for (const SubtargetFeatureKV &KV : AMDGPUFeatureKV) {
     if (Features[KV.Value])
       Ret += (StringRef("+") + KV.Key + ",").str();
     else if (TargetFeatures[KV.Value])

llvm/test/CodeGen/AMDGPU/propagate-attributes-clone.ll

 ; OPT-INT:  tail call fastcc void @foo1()
 ; OPT-INT:  tail call fastcc void @1()
 ; OPT:      ret void
 ; OPT: define internal fastcc void @1() unnamed_addr #4
 ; OPT-EXT: define internal fastcc void @foo3.2() unnamed_addr #4
 ; OPT-INT: define internal fastcc void @foo3() unnamed_addr #4
 ; OPT-EXT: define internal fastcc void @foo2.3() unnamed_addr #4
 ; OPT-INT: define internal fastcc void @foo2() unnamed_addr #4
-; OPT: attributes #0 = { {{.*}} "target-features"="+wavefrontsize64" }
+; OPT: attributes #0 = { {{.*}} "amdgpu-waves-per-eu"="1,1" "target-features"="+wavefrontsize64" }
 ; OPT: attributes #1 = { {{.*}} "target-features"="{{.*}},-wavefrontsize16,-wavefrontsize32,+wavefrontsize64{{.*}}" }
-; OPT: attributes #2 = { {{.*}} "target-features"="+wavefrontsize32" }
+; OPT: attributes #2 = { {{.*}} "amdgpu-waves-per-eu"="2,4" "target-features"="+wavefrontsize32" }
 ; OPT: attributes #3 = { {{.*}} "target-features"="+wavefrontsize64" }
-; OPT: attributes #4 = { {{.*}} "target-features"="{{.*}},-wavefrontsize16,+wavefrontsize32,-wavefrontsize64{{.*}}" }
+; OPT: attributes #4 = { {{.*}} "amdgpu-waves-per-eu"="2,4" "target-features"="{{.*}},-wavefrontsize16,+wavefrontsize32,-wavefrontsize64{{.*}}" }
 
 ; LLC: foo3:
 ; LLC: sample asm
 ; LLC: foo2:
 ; LLC: sample asm
 ; LLC: foo1:
 ; LLC: foo4@gotpcrel32@lo+4
 ; LLC: foo4@gotpcrel32@hi+4

 }
 
 define amdgpu_kernel void @kernel3() #3 {
 entry:
   tail call void @foo1()
   ret void
 }
 
-attributes #0 = { nounwind "target-features"="+wavefrontsize32" }
-attributes #1 = { noinline nounwind "target-features"="+wavefrontsize64" }
+attributes #0 = { nounwind "target-features"="+wavefrontsize32" "amdgpu-waves-per-eu"="2,4" }
+attributes #1 = { noinline nounwind "target-features"="+wavefrontsize64" "amdgpu-waves-per-eu"="1,1" }
 attributes #2 = { nounwind "target-features"="+wavefrontsize64" }
 attributes #3 = { nounwind "target-features"="+wavefrontsize64" }