[AMDGPU] gfx1010 wavefrontsize intrinsic folding

[arsenm]

I don't see codegen for this, so this looks like depending on an optimization pass for correctness

lib/Target/AMDGPU/AMDGPULibCalls.cpp

1392–1414

This doesn't really belong in this pass

test/CodeGen/AMDGPU/llvm.amdgcn.wavefrontsize.ll

7–8

Capitalized feature name is wrong? Also shouldn't be using opt -O3

[rampitec]

In D63206#1539987, @arsenm wrote:

I don't see codegen for this, so this looks like depending on an optimization pass for correctness

It is in the parent patch, inside ISelLowering. Optimization is only to fold it earlier.

[arsenm]

lib/Target/AMDGPU/AMDGPULibCalls.cpp

1392–1414

This might make more sense In AMDGPULowerIntrinsics?

[rampitec]

lib/Target/AMDGPU/AMDGPULibCalls.cpp

1392–1414

The point is to give folding a chance as early as possible. It also belongs to libcalls because library is the primary user.

[kzhuravl]

lgtm

[rampitec]

Rebased.

lib/Target/AMDGPU/AMDGPU.h

 FunctionPass *createSIWholeQuadModePass();
 FunctionPass *createSIFixControlFlowLiveIntervalsPass();
 FunctionPass *createSIOptimizeExecMaskingPreRAPass();
 FunctionPass *createSIFixSGPRCopiesPass();
 FunctionPass *createSIMemoryLegalizerPass();
 FunctionPass *createSIInsertWaitcntsPass();
 FunctionPass *createSIPreAllocateWWMRegsPass();
 FunctionPass *createSIFormMemoryClausesPass();
-FunctionPass *createAMDGPUSimplifyLibCallsPass(const TargetOptions &);
+FunctionPass *createAMDGPUSimplifyLibCallsPass(const TargetOptions &,
+                                               const TargetMachine *);
 FunctionPass *createAMDGPUUseNativeCallsPass();
 FunctionPass *createAMDGPUCodeGenPreparePass();
 FunctionPass *createAMDGPUMachineCFGStructurizerPass();
 FunctionPass *createAMDGPURewriteOutArgumentsPass();
 FunctionPass *createSIModeRegisterPass();
 
 void initializeAMDGPUDAGToDAGISelPass(PassRegistry&);
 

lib/Target/AMDGPU/AMDGPULibCalls.cpp

 /// This file does AMD library function optimizations.
 //
 //===----------------------------------------------------------------------===//
 
 #define DEBUG_TYPE "amdgpu-simplifylib"
 
 #include "AMDGPU.h"
 #include "AMDGPULibFunc.h"
+#include "AMDGPUSubtarget.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/Loads.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/ValueSymbolTable.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetOptions.h"
 #include <vector>
 #include <cmath>
 
 using namespace llvm;
 
 static cl::opt<bool> EnablePreLink("amdgpu-prelink",
   cl::desc("Enable pre-link mode optimizations"),

 
 namespace llvm {
 
 class AMDGPULibCalls {
 private:
 
   typedef llvm::AMDGPULibFunc FuncInfo;
 
+  const TargetMachine *TM;
+
   // -fuse-native.
   bool AllNative = false;
 
   bool useNativeFunc(const StringRef F) const;
 
   // Return a pointer (pointer expr) to the function if function defintion with
   // "FuncName" exists. It may create a new function prototype in pre-link mode.
   FunctionCallee getFunction(Module *M, const FuncInfo &fInfo);
private:
   bool fold_sqrt(CallInst *CI, IRBuilder<> &B, const FuncInfo &FInfo);
 
   // sin/cos
   bool fold_sincos(CallInst * CI, IRBuilder<> &B, AliasAnalysis * AA);
 
   // __read_pipe/__write_pipe
   bool fold_read_write_pipe(CallInst *CI, IRBuilder<> &B, FuncInfo &FInfo);
 
+  // llvm.amdgcn.wavefrontsize
+  bool fold_wavefrontsize(CallInst *CI, IRBuilder<> &B);
+
   // Get insertion point at entry.
   BasicBlock::iterator getEntryIns(CallInst * UI);
   // Insert an Alloc instruction.
   AllocaInst* insertAlloca(CallInst * UI, IRBuilder<> &B, const char *prefix);
   // Get a scalar native builtin signle argument FP function
   FunctionCallee getNativeFunction(Module *M, const FuncInfo &FInfo);
 
 protected:
   CallInst *CI;
 
   bool isUnsafeMath(const CallInst *CI) const;
 
   void replaceCall(Value *With) {
     CI->replaceAllUsesWith(With);
     CI->eraseFromParent();
   }
 
 public:
+  AMDGPULibCalls(const TargetMachine *TM_ = nullptr) : TM(TM_) {}
+
   bool fold(CallInst *CI, AliasAnalysis *AA = nullptr);
 
   void initNativeFuncs();
 
   // Replace a normal math function call with that native version
   bool useNative(CallInst *CI);
 };
 
 } // end llvm namespace
 
 namespace {
 
   class AMDGPUSimplifyLibCalls : public FunctionPass {
 
-  AMDGPULibCalls Simplifier;
-
   const TargetOptions Options;
 
+  AMDGPULibCalls Simplifier;
+
   public:
     static char ID; // Pass identification
 
-    AMDGPUSimplifyLibCalls(const TargetOptions &Opt = TargetOptions())
-      : FunctionPass(ID), Options(Opt) {
+    AMDGPUSimplifyLibCalls(const TargetOptions &Opt = TargetOptions(),
+                           const TargetMachine *TM = nullptr)
+      : FunctionPass(ID), Options(Opt), Simplifier(TM) {
       initializeAMDGPUSimplifyLibCallsPass(*PassRegistry::getPassRegistry());
     }
 
     void getAnalysisUsage(AnalysisUsage &AU) const override {
       AU.addRequired<AAResultsWrapperPass>();
     }
 
     bool runOnFunction(Function &M) override;

 // This function returns false if no change; return true otherwise.
 bool AMDGPULibCalls::fold(CallInst *CI, AliasAnalysis *AA) {
   this->CI = CI;
   Function *Callee = CI->getCalledFunction();
 
   // Ignore indirect calls.
   if (Callee == 0) return false;
 
-  FuncInfo FInfo;
-  if (!parseFunctionName(Callee->getName(), &FInfo))
-    return false;
-
-  // Further check the number of arguments to see if they match.
-  if (CI->getNumArgOperands() != FInfo.getNumArgs())
-    return false;
-
   BasicBlock *BB = CI->getParent();
   LLVMContext &Context = CI->getParent()->getContext();
   IRBuilder<> B(Context);
 
   // Set the builder to the instruction after the call.
   B.SetInsertPoint(BB, CI->getIterator());
 
   // Copy fast flags from the original call.
   if (const FPMathOperator *FPOp = dyn_cast<const FPMathOperator>(CI))
     B.setFastMathFlags(FPOp->getFastMathFlags());
 
+  switch (Callee->getIntrinsicID()) {
+  default:
+    break;
+  case Intrinsic::amdgcn_wavefrontsize:
+    return !EnablePreLink && fold_wavefrontsize(CI, B);
+  }
+
+  FuncInfo FInfo;
+  if (!parseFunctionName(Callee->getName(), &FInfo))
+    return false;
+
+  // Further check the number of arguments to see if they match.
+  if (CI->getNumArgOperands() != FInfo.getNumArgs())
+    return false;
+
   if (TDOFold(CI, FInfo))
     return true;
 
   // Under unsafe-math, evaluate calls if possible.
   // According to Brian Sumner, we can do this for all f32 function calls
   // using host's double function calls.
   if (isUnsafeMath(CI) && evaluateCall(CI, FInfo))
     return true;
if (!isSin) { // CI->cos, UI->sin
     UI->replaceAllUsesWith(Reload);
     CI->replaceAllUsesWith(Call);
     UI->eraseFromParent();
     CI->eraseFromParent();
   }
   return true;
 }
 
+bool AMDGPULibCalls::fold_wavefrontsize(CallInst *CI, IRBuilder<> &B) {
+  if (!TM)
+    return false;
+
+  StringRef CPU = TM->getTargetCPU();
+  StringRef Features = TM->getTargetFeatureString();
+  if ((CPU.empty() || CPU.equals_lower("generic")) &&
+      (Features.empty() ||
+       Features.find_lower("wavefrontsize") == StringRef::npos))
+    return false;
+
+  Function *F = CI->getParent()->getParent();
+  const GCNSubtarget &ST = TM->getSubtarget<GCNSubtarget>(*F);
+  unsigned N = ST.getWavefrontSize();
+
+  LLVM_DEBUG(errs() << "AMDIC: fold_wavefrontsize (" << *CI << ") with "
+               << N << "\n");
+
+  CI->replaceAllUsesWith(ConstantInt::get(B.getInt32Ty(), N));
+  CI->eraseFromParent();
+  return true;
+}
+
 // Get insertion point at entry.
 BasicBlock::iterator AMDGPULibCalls::getEntryIns(CallInst * UI) {
   Function * Func = UI->getParent()->getParent();
   BasicBlock * BB = &Func->getEntryBlock();
   assert(BB && "Entry block not found!");
   BasicBlock::iterator ItNew = BB->begin();
   return ItNew;
 }
if (hasTwoResults) {
     new StoreInst(nval1, aCI->getArgOperand(1), aCI);
   }
 
   replaceCall(nval0);
   return true;
 }
 
 // Public interface to the Simplify LibCalls pass.
-FunctionPass *llvm::createAMDGPUSimplifyLibCallsPass(const TargetOptions &Opt) {
-  return new AMDGPUSimplifyLibCalls(Opt);
+FunctionPass *llvm::createAMDGPUSimplifyLibCallsPass(const TargetOptions &Opt,
+                                                     const TargetMachine *TM) {
+  return new AMDGPUSimplifyLibCalls(Opt, TM);
 }
 
 FunctionPass *llvm::createAMDGPUUseNativeCallsPass() {
   return new AMDGPUUseNativeCalls();
 }
 
 static bool setFastFlags(Function &F, const TargetOptions &Options) {
   AttrBuilder B;

lib/Target/AMDGPU/AMDGPUTargetMachine.cpp

[Internalize, EarlyInline, AMDGPUAA](const PassManagerBuilder &,
       }
       if (EarlyInline)
         PM.add(createAMDGPUAlwaysInlinePass(false));
   });
 
   const auto &Opt = Options;
   Builder.addExtension(
     PassManagerBuilder::EP_EarlyAsPossible,
-    [AMDGPUAA, LibCallSimplify, &Opt](const PassManagerBuilder &,
+    [AMDGPUAA, LibCallSimplify, &Opt, this](const PassManagerBuilder &,
                                             legacy::PassManagerBase &PM) {
       if (AMDGPUAA) {
         PM.add(createAMDGPUAAWrapperPass());
         PM.add(createAMDGPUExternalAAWrapperPass());
       }
       PM.add(llvm::createAMDGPUUseNativeCallsPass());
       if (LibCallSimplify)
-        PM.add(llvm::createAMDGPUSimplifyLibCallsPass(Opt));
+        PM.add(llvm::createAMDGPUSimplifyLibCallsPass(Opt, this));
   });
 
   Builder.addExtension(
     PassManagerBuilder::EP_CGSCCOptimizerLate,
     [](const PassManagerBuilder &, legacy::PassManagerBase &PM) {
       // Add infer address spaces pass to the opt pipeline after inlining
       // but before SROA to increase SROA opportunities.
       PM.add(createInferAddressSpacesPass());

test/CodeGen/AMDGPU/llvm.amdgcn.wavefrontsize.ll

+; RUN: llc -march=amdgcn -mcpu=fiji -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,W64 %s
+; RUN: llc -march=amdgcn -mcpu=gfx1010 -mattr=+wavefrontsize32,-wavefrontsize64 -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,W32 %s
+; RUN: llc -march=amdgcn -mcpu=gfx1010 -mattr=-wavefrontsize32,+wavefrontsize64 -verify-machineinstrs < %s | FileCheck -check-prefixes=GCN,W64 %s
+
+; RUN: opt -O3 -S < %s | FileCheck -check-prefixes=OPT,OPT-WXX %s
+; RUN: opt -mtriple=amdgcn-- -O3 -S < %s | FileCheck -check-prefixes=OPT,OPT-WXX %s
+; RUN: opt -mtriple=amdgcn-- -O3 -mattr=+WavefrontSize32 -S < %s | FileCheck -check-prefixes=OPT,OPT-W32 %s
+; RUN: opt -mtriple=amdgcn-- -O3 -mattr=+WavefrontSize64 -S < %s | FileCheck -check-prefixes=OPT,OPT-W64 %s
+; RUN: opt -mtriple=amdgcn-- -mcpu=tonga -O3 -S < %s | FileCheck -check-prefixes=OPT,OPT-W64 %s
+; RUN: opt -mtriple=amdgcn-- -mcpu=gfx1010 -O3 -mattr=+wavefrontsize32,-wavefrontsize64 -S < %s | FileCheck -check-prefixes=OPT,OPT-W32 %s
+; RUN: opt -mtriple=amdgcn-- -mcpu=gfx1010 -O3 -mattr=-wavefrontsize32,+wavefrontsize64 -S < %s | FileCheck -check-prefixes=OPT,OPT-W64 %s
+
+; GCN-LABEL: {{^}}fold_wavefrontsize:
+; OPT-LABEL: define amdgpu_kernel void @fold_wavefrontsize(
+
+; W32:       v_mov_b32_e32 [[V:v[0-9]+]], 32
+; W64:       v_mov_b32_e32 [[V:v[0-9]+]], 64
+; GCN:       store_dword v[{{[0-9:]+}}], [[V]]
+
+; OPT-W32:   store i32 32, i32 addrspace(1)* %arg, align 4
+; OPT-W64:   store i32 64, i32 addrspace(1)* %arg, align 4
+; OPT-WXX:   %tmp = tail call i32 @llvm.amdgcn.wavefrontsize()
+; OPT-WXX:   store i32 %tmp, i32 addrspace(1)* %arg, align 4
+; OPT-NEXT:  ret void
+
+define amdgpu_kernel void @fold_wavefrontsize(i32 addrspace(1)* nocapture %arg) {
+bb:
+  %tmp = tail call i32 @llvm.amdgcn.wavefrontsize() #0
+  store i32 %tmp, i32 addrspace(1)* %arg, align 4
+  ret void
+}
+
+; GCN-LABEL: {{^}}fold_and_optimize_wavefrontsize:
+; OPT-LABEL: define amdgpu_kernel void @fold_and_optimize_wavefrontsize(
+
+; W32:       v_mov_b32_e32 [[V:v[0-9]+]], 1{{$}}
+; W64:       v_mov_b32_e32 [[V:v[0-9]+]], 2{{$}}
+; GCN-NOT:   cndmask
+; GCN:       store_dword v[{{[0-9:]+}}], [[V]]
+
+; OPT-W32:   store i32 1, i32 addrspace(1)* %arg, align 4
+; OPT-W64:   store i32 2, i32 addrspace(1)* %arg, align 4
+; OPT-WXX:   %tmp = tail call i32 @llvm.amdgcn.wavefrontsize()
+; OPT-WXX:   %tmp1 = icmp ugt i32 %tmp, 32
+; OPT-WXX:   %tmp2 = select i1 %tmp1, i32 2, i32 1
+; OPT-WXX:   store i32 %tmp2, i32 addrspace(1)* %arg
+; OPT-NEXT:  ret void
+
+define amdgpu_kernel void @fold_and_optimize_wavefrontsize(i32 addrspace(1)* nocapture %arg) {
+bb:
+  %tmp = tail call i32 @llvm.amdgcn.wavefrontsize() #0
+  %tmp1 = icmp ugt i32 %tmp, 32
+  %tmp2 = select i1 %tmp1, i32 2, i32 1
+  store i32 %tmp2, i32 addrspace(1)* %arg
+  ret void
+}
+
+; GCN-LABEL: {{^}}fold_and_optimize_if_wavefrontsize:
+; OPT-LABEL: define amdgpu_kernel void @fold_and_optimize_if_wavefrontsize(
+
+; OPT:       bb:
+; OPT-WXX:   %tmp = tail call i32 @llvm.amdgcn.wavefrontsize()
+; OPT-WXX:   %tmp1 = icmp ugt i32 %tmp, 32
+; OPT-WXX:   bb3:
+; OPT-W64:   store i32 1, i32 addrspace(1)* %arg, align 4
+; OPT-NEXT:  ret void
+
+define amdgpu_kernel void @fold_and_optimize_if_wavefrontsize(i32 addrspace(1)* nocapture %arg) {
+bb:
+  %tmp = tail call i32 @llvm.amdgcn.wavefrontsize() #0
+  %tmp1 = icmp ugt i32 %tmp, 32
+  br i1 %tmp1, label %bb2, label %bb3
+
+bb2:                                              ; preds = %bb
+  store i32 1, i32 addrspace(1)* %arg, align 4
+  br label %bb3
+
+bb3:                                              ; preds = %bb2, %bb
+  ret void
+}
+
+declare i32 @llvm.amdgcn.wavefrontsize() #0
+
+attributes #0 = { nounwind readnone speculatable }