[AMDGPU] gfx10 atomic optimizer changes.

Summary:
Add support for gfx10, where all DPP operations are confined to work
within a single row of 16 lanes, and wave32.

Reviewers: arsenm, sheredom, critson, rampitec

Subscribers: kzhuravl, jvesely, wdng, nhaehnle, yaxunl, t-tye, hiraditya, jfb, dstuttard, tpr, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D65644

llvm-svn: 369745

1 files changed, 143 insertions, 57 deletions

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUAtomicOptimizer.cpp b/llvm/lib/Target/AMDGPU/AMDGPUAtomicOptimizer.cpp
index 92d94ee894a..ba8343142c6 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUAtomicOptimizer.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUAtomicOptimizer.cpp
@@ -15,6 +15,7 @@
 
 #include "AMDGPU.h"
 #include "AMDGPUSubtarget.h"
+#include "SIDefines.h"
 #include "llvm/Analysis/LegacyDivergenceAnalysis.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/IR/IRBuilder.h"
@@ -24,20 +25,10 @@
 #define DEBUG_TYPE "amdgpu-atomic-optimizer"
 
 using namespace llvm;
+using namespace llvm::AMDGPU;
 
 namespace {
 
-enum DPP_CTRL {
-  DPP_ROW_SR1 = 0x111,
-  DPP_ROW_SR2 = 0x112,
-  DPP_ROW_SR3 = 0x113,
-  DPP_ROW_SR4 = 0x114,
-  DPP_ROW_SR8 = 0x118,
-  DPP_WF_SR1 = 0x138,
-  DPP_ROW_BCAST15 = 0x142,
-  DPP_ROW_BCAST31 = 0x143
-};
-
 struct ReplacementInfo {
   Instruction *I;
   AtomicRMWInst::BinOp Op;
@@ -52,9 +43,12 @@ private:
   const LegacyDivergenceAnalysis *DA;
   const DataLayout *DL;
   DominatorTree *DT;
-  bool HasDPP;
+  const GCNSubtarget *ST;
   bool IsPixelShader;
 
+  Value *buildScan(IRBuilder<> &B, AtomicRMWInst::BinOp Op, Value *V,
+                   Value *const Identity) const;
+  Value *buildShiftRight(IRBuilder<> &B, Value *V, Value *const Identity) const;
   void optimizeAtomic(Instruction &I, AtomicRMWInst::BinOp Op, unsigned ValIdx,
                       bool ValDivergent) const;
 
@@ -93,8 +87,7 @@ bool AMDGPUAtomicOptimizer::runOnFunction(Function &F) {
   DT = DTW ? &DTW->getDomTree() : nullptr;
   const TargetPassConfig &TPC = getAnalysis<TargetPassConfig>();
   const TargetMachine &TM = TPC.getTM<TargetMachine>();
-  const GCNSubtarget &ST = TM.getSubtarget<GCNSubtarget>(F);
-  HasDPP = ST.hasDPP();
+  ST = &TM.getSubtarget<GCNSubtarget>(F);
   IsPixelShader = F.getCallingConv() == CallingConv::AMDGPU_PS;
 
   visit(F);
@@ -152,7 +145,8 @@ void AMDGPUAtomicOptimizer::visitAtomicRMWInst(AtomicRMWInst &I) {
   // value to the atomic calculation. We can only optimize divergent values if
   // we have DPP available on our subtarget, and the atomic operation is 32
   // bits.
-  if (ValDivergent && (!HasDPP || (DL->getTypeSizeInBits(I.getType()) != 32))) {
+  if (ValDivergent &&
+      (!ST->hasDPP() || DL->getTypeSizeInBits(I.getType()) != 32)) {
     return;
   }
 
@@ -225,7 +219,8 @@ void AMDGPUAtomicOptimizer::visitIntrinsicInst(IntrinsicInst &I) {
   // value to the atomic calculation. We can only optimize divergent values if
   // we have DPP available on our subtarget, and the atomic operation is 32
   // bits.
-  if (ValDivergent && (!HasDPP || (DL->getTypeSizeInBits(I.getType()) != 32))) {
+  if (ValDivergent &&
+      (!ST->hasDPP() || DL->getTypeSizeInBits(I.getType()) != 32)) {
     return;
   }
 
@@ -282,6 +277,111 @@ static Value *buildNonAtomicBinOp(IRBuilder<> &B, AtomicRMWInst::BinOp Op,
   return B.CreateSelect(Cond, LHS, RHS);
 }
 
+// Use the builder to create an inclusive scan of V across the wavefront, with
+// all lanes active.
+Value *AMDGPUAtomicOptimizer::buildScan(IRBuilder<> &B, AtomicRMWInst::BinOp Op,
+                                        Value *V, Value *const Identity) const {
+  Type *const Ty = V->getType();
+  Module *M = B.GetInsertBlock()->getModule();
+  Function *UpdateDPP =
+      Intrinsic::getDeclaration(M, Intrinsic::amdgcn_update_dpp, Ty);
+  Function *PermLaneX16 =
+      Intrinsic::getDeclaration(M, Intrinsic::amdgcn_permlanex16, {});
+  Function *ReadLane =
+      Intrinsic::getDeclaration(M, Intrinsic::amdgcn_readlane, {});
+
+  for (unsigned Idx = 0; Idx < 4; Idx++) {
+    V = buildNonAtomicBinOp(
+        B, Op, V,
+        B.CreateCall(UpdateDPP,
+                     {Identity, V, B.getInt32(DPP::ROW_SHR0 | 1 << Idx),
+                      B.getInt32(0xf), B.getInt32(0xf), B.getFalse()}));
+  }
+  if (ST->hasDPPBroadcasts()) {
+    // GFX9 has DPP row broadcast operations.
+    V = buildNonAtomicBinOp(
+        B, Op, V,
+        B.CreateCall(UpdateDPP,
+                     {Identity, V, B.getInt32(DPP::BCAST15), B.getInt32(0xa),
+                      B.getInt32(0xf), B.getFalse()}));
+    V = buildNonAtomicBinOp(
+        B, Op, V,
+        B.CreateCall(UpdateDPP,
+                     {Identity, V, B.getInt32(DPP::BCAST31), B.getInt32(0xc),
+                      B.getInt32(0xf), B.getFalse()}));
+  } else {
+    // On GFX10 all DPP operations are confined to a single row. To get cross-
+    // row operations we have to use permlane or readlane.
+
+    // Combine lane 15 into lanes 16..31 (and, for wave 64, lane 47 into lanes
+    // 48..63).
+    Value *const PermX =
+        B.CreateCall(PermLaneX16, {V, V, B.getInt32(-1), B.getInt32(-1),
+                                   B.getFalse(), B.getFalse()});
+    V = buildNonAtomicBinOp(
+        B, Op, V,
+        B.CreateCall(UpdateDPP,
+                     {Identity, PermX, B.getInt32(DPP::QUAD_PERM_ID),
+                      B.getInt32(0xa), B.getInt32(0xf), B.getFalse()}));
+    if (!ST->isWave32()) {
+      // Combine lane 31 into lanes 32..63.
+      Value *const Lane31 = B.CreateCall(ReadLane, {V, B.getInt32(31)});
+      V = buildNonAtomicBinOp(
+          B, Op, V,
+          B.CreateCall(UpdateDPP,
+                       {Identity, Lane31, B.getInt32(DPP::QUAD_PERM_ID),
+                        B.getInt32(0xc), B.getInt32(0xf), B.getFalse()}));
+    }
+  }
+  return V;
+}
+
+// Use the builder to create a shift right of V across the wavefront, with all
+// lanes active, to turn an inclusive scan into an exclusive scan.
+Value *AMDGPUAtomicOptimizer::buildShiftRight(IRBuilder<> &B, Value *V,
+                                              Value *const Identity) const {
+  Type *const Ty = V->getType();
+  Module *M = B.GetInsertBlock()->getModule();
+  Function *UpdateDPP =
+      Intrinsic::getDeclaration(M, Intrinsic::amdgcn_update_dpp, Ty);
+  Function *ReadLane =
+      Intrinsic::getDeclaration(M, Intrinsic::amdgcn_readlane, {});
+  Function *WriteLane =
+      Intrinsic::getDeclaration(M, Intrinsic::amdgcn_writelane, {});
+
+  if (ST->hasDPPWavefrontShifts()) {
+    // GFX9 has DPP wavefront shift operations.
+    V = B.CreateCall(UpdateDPP,
+                     {Identity, V, B.getInt32(DPP::WAVE_SHR1), B.getInt32(0xf),
+                      B.getInt32(0xf), B.getFalse()});
+  } else {
+    // On GFX10 all DPP operations are confined to a single row. To get cross-
+    // row operations we have to use permlane or readlane.
+    Value *Old = V;
+    V = B.CreateCall(UpdateDPP,
+                     {Identity, V, B.getInt32(DPP::ROW_SHR0 + 1),
+                      B.getInt32(0xf), B.getInt32(0xf), B.getFalse()});
+
+    // Copy the old lane 15 to the new lane 16.
+    V = B.CreateCall(WriteLane, {B.CreateCall(ReadLane, {Old, B.getInt32(15)}),
+                                 B.getInt32(16), V});
+
+    if (!ST->isWave32()) {
+      // Copy the old lane 31 to the new lane 32.
+      V = B.CreateCall(
+          WriteLane,
+          {B.CreateCall(ReadLane, {Old, B.getInt32(31)}), B.getInt32(32), V});
+
+      // Copy the old lane 47 to the new lane 48.
+      V = B.CreateCall(
+          WriteLane,
+          {B.CreateCall(ReadLane, {Old, B.getInt32(47)}), B.getInt32(48), V});
+    }
+  }
+
+  return V;
+}
+
 static APInt getIdentityValueForAtomicOp(AtomicRMWInst::BinOp Op,
                                          unsigned BitWidth) {
   switch (Op) {
@@ -345,23 +445,29 @@ void AMDGPUAtomicOptimizer::optimizeAtomic(Instruction &I,
 
   // We need to know how many lanes are active within the wavefront, and we do
   // this by doing a ballot of active lanes.
+  Type *const WaveTy = B.getIntNTy(ST->getWavefrontSize());
   CallInst *const Ballot = B.CreateIntrinsic(
-      Intrinsic::amdgcn_icmp, {B.getInt64Ty(), B.getInt32Ty()},
+      Intrinsic::amdgcn_icmp, {WaveTy, B.getInt32Ty()},
       {B.getInt32(1), B.getInt32(0), B.getInt32(CmpInst::ICMP_NE)});
 
   // We need to know how many lanes are active within the wavefront that are
   // below us. If we counted each lane linearly starting from 0, a lane is
   // below us only if its associated index was less than ours. We do this by
   // using the mbcnt intrinsic.
-  Value *const BitCast = B.CreateBitCast(Ballot, VecTy);
-  Value *const ExtractLo = B.CreateExtractElement(BitCast, B.getInt32(0));
-  Value *const ExtractHi = B.CreateExtractElement(BitCast, B.getInt32(1));
-  CallInst *const PartialMbcnt = B.CreateIntrinsic(
-      Intrinsic::amdgcn_mbcnt_lo, {}, {ExtractLo, B.getInt32(0)});
-  Value *const Mbcnt =
-      B.CreateIntCast(B.CreateIntrinsic(Intrinsic::amdgcn_mbcnt_hi, {},
-                                        {ExtractHi, PartialMbcnt}),
-                      Ty, false);
+  Value *Mbcnt;
+  if (ST->isWave32()) {
+    Mbcnt = B.CreateIntrinsic(Intrinsic::amdgcn_mbcnt_lo, {},
+                              {Ballot, B.getInt32(0)});
+  } else {
+    Value *const BitCast = B.CreateBitCast(Ballot, VecTy);
+    Value *const ExtractLo = B.CreateExtractElement(BitCast, B.getInt32(0));
+    Value *const ExtractHi = B.CreateExtractElement(BitCast, B.getInt32(1));
+    Mbcnt = B.CreateIntrinsic(Intrinsic::amdgcn_mbcnt_lo, {},
+                              {ExtractLo, B.getInt32(0)});
+    Mbcnt =
+        B.CreateIntrinsic(Intrinsic::amdgcn_mbcnt_hi, {}, {ExtractHi, Mbcnt});
+  }
+  Mbcnt = B.CreateIntCast(Mbcnt, Ty, false);
 
   Value *const Identity = B.getInt(getIdentityValueForAtomicOp(Op, TyBitWidth));
 
@@ -373,45 +479,25 @@ void AMDGPUAtomicOptimizer::optimizeAtomic(Instruction &I,
   if (ValDivergent) {
     // First we need to set all inactive invocations to the identity value, so
     // that they can correctly contribute to the final result.
-    CallInst *const SetInactive =
-        B.CreateIntrinsic(Intrinsic::amdgcn_set_inactive, Ty, {V, Identity});
-
-    ExclScan =
-        B.CreateIntrinsic(Intrinsic::amdgcn_update_dpp, Ty,
-                          {Identity, SetInactive, B.getInt32(DPP_WF_SR1),
-                           B.getInt32(0xf), B.getInt32(0xf), B.getFalse()});
-
-    const unsigned Iters = 6;
-    const unsigned DPPCtrl[Iters] = {DPP_ROW_SR1,     DPP_ROW_SR2,
-                                     DPP_ROW_SR4,     DPP_ROW_SR8,
-                                     DPP_ROW_BCAST15, DPP_ROW_BCAST31};
-    const unsigned RowMask[Iters] = {0xf, 0xf, 0xf, 0xf, 0xa, 0xc};
-    const unsigned BankMask[Iters] = {0xf, 0xf, 0xe, 0xc, 0xf, 0xf};
-
-    // This loop performs an exclusive scan across the wavefront, with all lanes
-    // active (by using the WWM intrinsic).
-    for (unsigned Idx = 0; Idx < Iters; Idx++) {
-      CallInst *const DPP = B.CreateIntrinsic(
-          Intrinsic::amdgcn_update_dpp, Ty,
-          {Identity, ExclScan, B.getInt32(DPPCtrl[Idx]),
-           B.getInt32(RowMask[Idx]), B.getInt32(BankMask[Idx]), B.getFalse()});
-
-      ExclScan = buildNonAtomicBinOp(B, Op, ExclScan, DPP);
-    }
+    NewV = B.CreateIntrinsic(Intrinsic::amdgcn_set_inactive, Ty, {V, Identity});
 
-    NewV = buildNonAtomicBinOp(B, Op, SetInactive, ExclScan);
+    const AtomicRMWInst::BinOp ScanOp =
+        Op == AtomicRMWInst::Sub ? AtomicRMWInst::Add : Op;
+    NewV = buildScan(B, ScanOp, NewV, Identity);
+    ExclScan = buildShiftRight(B, NewV, Identity);
 
     // Read the value from the last lane, which has accumlated the values of
     // each active lane in the wavefront. This will be our new value which we
     // will provide to the atomic operation.
+    Value *const LastLaneIdx = B.getInt32(ST->getWavefrontSize() - 1);
     if (TyBitWidth == 64) {
       Value *const ExtractLo = B.CreateTrunc(NewV, B.getInt32Ty());
       Value *const ExtractHi =
-          B.CreateTrunc(B.CreateLShr(NewV, B.getInt64(32)), B.getInt32Ty());
+          B.CreateTrunc(B.CreateLShr(NewV, 32), B.getInt32Ty());
       CallInst *const ReadLaneLo = B.CreateIntrinsic(
-          Intrinsic::amdgcn_readlane, {}, {ExtractLo, B.getInt32(63)});
+          Intrinsic::amdgcn_readlane, {}, {ExtractLo, LastLaneIdx});
       CallInst *const ReadLaneHi = B.CreateIntrinsic(
-          Intrinsic::amdgcn_readlane, {}, {ExtractHi, B.getInt32(63)});
+          Intrinsic::amdgcn_readlane, {}, {ExtractHi, LastLaneIdx});
       Value *const PartialInsert = B.CreateInsertElement(
           UndefValue::get(VecTy), ReadLaneLo, B.getInt32(0));
       Value *const Insert =
@@ -419,7 +505,7 @@ void AMDGPUAtomicOptimizer::optimizeAtomic(Instruction &I,
       NewV = B.CreateBitCast(Insert, Ty);
     } else if (TyBitWidth == 32) {
       NewV = B.CreateIntrinsic(Intrinsic::amdgcn_readlane, {},
-                               {NewV, B.getInt32(63)});
+                               {NewV, LastLaneIdx});
     } else {
       llvm_unreachable("Unhandled atomic bit width");
     }
@@ -506,7 +592,7 @@ void AMDGPUAtomicOptimizer::optimizeAtomic(Instruction &I,
     if (TyBitWidth == 64) {
       Value *const ExtractLo = B.CreateTrunc(PHI, B.getInt32Ty());
       Value *const ExtractHi =
-          B.CreateTrunc(B.CreateLShr(PHI, B.getInt64(32)), B.getInt32Ty());
+          B.CreateTrunc(B.CreateLShr(PHI, 32), B.getInt32Ty());
       CallInst *const ReadFirstLaneLo =
           B.CreateIntrinsic(Intrinsic::amdgcn_readfirstlane, {}, ExtractLo);
       CallInst *const ReadFirstLaneHi =