AMDGPU: Propagate amdgpu-flat-work-group-size attributes

[arsenm]

Summary

Fixes being overly conservative with the register counts in called
functions. This should try to do a conservative range merge, but for
now just clone.

Also fix not being able to functionally run the pass standalone.

[rampitec]

LGTM.

Cloning might be a good thing anyway. One call stack will use conservative slow version and some other a faster version.

[arsenm]

53c43431bc6a01ad1e29b9351450ac18d5270ab3

llvm/lib/Target/AMDGPU/AMDGPUPropagateAttributes.cpp

 //===----------------------------------------------------------------------===//
 
 #include "AMDGPU.h"
 #include "AMDGPUSubtarget.h"
 #include "MCTargetDesc/AMDGPUMCTargetDesc.h"
 #include "Utils/AMDGPUBaseInfo.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include <string>
 
 #define DEBUG_TYPE "amdgpu-propagate-attributes"
 
 using namespace llvm;
 
 namespace llvm {
 extern const SubtargetFeatureKV AMDGPUFeatureKV[AMDGPU::NumSubtargetFeatures-1];
 }
 
 namespace {
 
 // Target features to propagate.
 static constexpr const FeatureBitset TargetFeatures = {
   AMDGPU::FeatureWavefrontSize16,
   AMDGPU::FeatureWavefrontSize32,
   AMDGPU::FeatureWavefrontSize64
 };
 
 // Attributes to propagate.
+// TODO: Support conservative min/max merging instead of cloning.
 static constexpr const char* AttributeNames[] = {
-  "amdgpu-waves-per-eu"
+  "amdgpu-waves-per-eu",
+  "amdgpu-flat-work-group-size"
 };
 
 static constexpr unsigned NumAttr =
   sizeof(AttributeNames) / sizeof(AttributeNames[0]);
 
 class AMDGPUPropagateAttributes {
 
   class FnProperties {
for (const SubtargetFeatureKV &KV : AMDGPUFeatureKV) {
     else if (TargetFeatures[KV.Value])
       Ret += (StringRef("-") + KV.Key + ",").str();
   }
   Ret.pop_back(); // Remove last comma.
   return Ret;
 }
 
 bool AMDGPUPropagateAttributesEarly::runOnFunction(Function &F) {
-  if (!TM || !AMDGPU::isEntryFunctionCC(F.getCallingConv()))
+  if (!TM) {
+    auto *TPC = getAnalysisIfAvailable<TargetPassConfig>();
+    if (!TPC)
+      return false;
+
+    TM = &TPC->getTM<TargetMachine>();
+  }
+
+  if (!AMDGPU::isEntryFunctionCC(F.getCallingConv()))
     return false;
 
   return AMDGPUPropagateAttributes(TM, false).process(F);
 }
 
 bool AMDGPUPropagateAttributesLate::runOnModule(Module &M) {
-  if (!TM)
+  if (!TM) {
+    auto *TPC = getAnalysisIfAvailable<TargetPassConfig>();
+    if (!TPC)
       return false;
 
+    TM = &TPC->getTM<TargetMachine>();
+  }
+
   return AMDGPUPropagateAttributes(TM, true).process(M);
 }
 
 FunctionPass
 *llvm::createAMDGPUPropagateAttributesEarlyPass(const TargetMachine *TM) {
   return new AMDGPUPropagateAttributesEarly(TM);
 }
 
 ModulePass
 *llvm::createAMDGPUPropagateAttributesLatePass(const TargetMachine *TM) {
   return new AMDGPUPropagateAttributesLate(TM);
 }

llvm/test/CodeGen/AMDGPU/propagate-attributes-flat-work-group-size.ll

+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -amdgpu-propagate-attributes-late %s | FileCheck %s
+
+; CHECK: define internal void @max_flat_1_1024() #0 {
+define internal void @max_flat_1_1024() #0 {
+  ret void
+}
+
+; CHECK: define internal void @max_flat_1_256() #1 {
+define internal void @max_flat_1_256() #1 {
+  ret void
+}
+
+; CHECK: define amdgpu_kernel void @kernel_1_256_call_default() #1 {
+define amdgpu_kernel void @kernel_1_256_call_default() #1 {
+  call void @default()
+  ret void
+}
+
+; CHECK: define amdgpu_kernel void @kernel_1_256_call_1_256() #1 {
+define amdgpu_kernel void @kernel_1_256_call_1_256() #1 {
+  call void @max_flat_1_256()
+  ret void
+}
+
+; CHECK: define amdgpu_kernel void @kernel_1_256_call_64_64() #1 {
+define amdgpu_kernel void @kernel_1_256_call_64_64() #1 {
+  call void @max_flat_64_64()
+  ret void
+}
+
+; CHECK: define internal void @max_flat_64_64() #2 {
+define internal void @max_flat_64_64() #2 {
+  ret void
+}
+
+; CHECK: define internal void @default() #2 {
+define internal void @default() #3 {
+  ret void
+}
+
+attributes #0 = { noinline "amdgpu-flat-work-group-size"="1,1024" }
+attributes #1 = { noinline "amdgpu-flat-work-group-size"="1,256" }
+attributes #2 = { noinline "amdgpu-flat-work-group-size"="64,64" }
+attributes #3 = { noinline }
+
+; CHECK: attributes #0 = { noinline "amdgpu-flat-work-group-size"="1,1024"
+; CHECK-NEXT: attributes #1 = { noinline "amdgpu-flat-work-group-size"="1,256"
+; CHECK-NEXT: attributes #2 = { noinline "amdgpu-flat-work-group-size"="1,256"