[AMDGPU] Use AssumptionCacheTracker in the divrem32 expansion

[rampitec]

Summary

This allows to use a short div/rem expansion given a kernel with __builtin_assume() call:

kernel void divrem24_uint(global int *a, uint x) {
  uint id = get_local_id(0);
  __builtin_assume(x < 42);
  a[id / x] = 0;
}

[artem.tamazov]

LGTM

test/CodeGen/AMDGPU/divrem24-assume.ll

31

The usage of the Answer to the Ultimate Question of Life, the Universe, and Everything to resolve the divrem problem is definitely overkill ))

lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp

 /// selection.
 //
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPU.h"
 #include "AMDGPUSubtarget.h"
 #include "AMDGPUTargetMachine.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/DivergenceAnalysis.h"
 #include "llvm/Analysis/Loads.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Constants.h"
static cl::opt<bool> WidenLoads(
   "amdgpu-codegenprepare-widen-constant-loads",
   cl::desc("Widen sub-dword constant address space loads in AMDGPUCodeGenPrepare"),
   cl::ReallyHidden,
   cl::init(true));
 
 class AMDGPUCodeGenPrepare : public FunctionPass,
                              public InstVisitor<AMDGPUCodeGenPrepare, bool> {
   const GCNSubtarget *ST = nullptr;
+  AssumptionCache *AC = nullptr;
   DivergenceAnalysis *DA = nullptr;
   Module *Mod = nullptr;
   bool HasUnsafeFPMath = false;
   AMDGPUAS AMDGPUASI;
 
   /// Copies exact/nsw/nuw flags (if any) from binary operation \p I to
   /// binary operation \p V.
   ///
class AMDGPUCodeGenPrepare : public FunctionPass,
   /// result of 32 bit 'bitreverse' intrinsic to the right with zero fill (the
   /// shift amount is 32 minus \p I's base element bit width), and truncating
   /// the result of the shift operation back to \p I's original type.
   ///
   /// \returns True.
   bool promoteUniformBitreverseToI32(IntrinsicInst &I) const;
 
   /// Expands 24 bit div or rem.
-  Value* expandDivRem24(IRBuilder<> &Builder, Value *Num, Value *Den,
+  Value* expandDivRem24(IRBuilder<> &Builder, BinaryOperator &I,
+                        Value *Num, Value *Den,
                         bool IsDiv, bool IsSigned) const;
 
   /// Expands 32 bit div or rem.
-  Value* expandDivRem32(IRBuilder<> &Builder, Instruction::BinaryOps Opc,
+  Value* expandDivRem32(IRBuilder<> &Builder, BinaryOperator &I,
                         Value *Num, Value *Den) const;
 
   /// Widen a scalar load.
   ///
   /// \details \p Widen scalar load for uniform, small type loads from constant
   //  memory / to a full 32-bits and then truncate the input to allow a scalar
   //  load instead of a vector load.
   //
public:
   bool visitBitreverseIntrinsicInst(IntrinsicInst &I);
 
   bool doInitialization(Module &M) override;
   bool runOnFunction(Function &F) override;
 
   StringRef getPassName() const override { return "AMDGPU IR optimizations"; }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.addRequired<AssumptionCacheTracker>();
     AU.addRequired<DivergenceAnalysis>();
     AU.setPreservesAll();
  }
 };
 
 } // end anonymous namespace
 
 unsigned AMDGPUCodeGenPrepare::getBaseElementBitWidth(const Type *T) const {

 
 static Value* getMulHu(IRBuilder<> &Builder, Value *LHS, Value *RHS) {
   return getMul64(Builder, LHS, RHS).second;
 }
 
 // The fractional part of a float is enough to accurately represent up to
 // a 24-bit signed integer.
 Value* AMDGPUCodeGenPrepare::expandDivRem24(IRBuilder<> &Builder,
+                                            BinaryOperator &I,
                                             Value *Num, Value *Den,
                                             bool IsDiv, bool IsSigned) const {
   assert(Num->getType()->isIntegerTy(32));
 
   const DataLayout &DL = Mod->getDataLayout();
-  unsigned LHSSignBits = ComputeNumSignBits(Num, DL);
+  unsigned LHSSignBits = ComputeNumSignBits(Num, DL, 0, AC, &I);
   if (LHSSignBits < 9)
     return nullptr;
 
-  unsigned RHSSignBits = ComputeNumSignBits(Den, DL);
+  unsigned RHSSignBits = ComputeNumSignBits(Den, DL, 0, AC, &I);
   if (RHSSignBits < 9)
     return nullptr;
 
 
   unsigned SignBits = std::min(LHSSignBits, RHSSignBits);
   unsigned DivBits = 32 - SignBits;
   if (IsSigned)
     ++DivBits;
if (IsSigned) {
     ConstantInt *TruncMask = Builder.getInt32((UINT64_C(1) << DivBits) - 1);
     Res = Builder.CreateAnd(Res, TruncMask);
   }
 
   return Res;
 }
 
 Value* AMDGPUCodeGenPrepare::expandDivRem32(IRBuilder<> &Builder,
-                                            Instruction::BinaryOps Opc,
+                                            BinaryOperator &I,
                                             Value *Num, Value *Den) const {
+  Instruction::BinaryOps Opc = I.getOpcode();
   assert(Opc == Instruction::URem || Opc == Instruction::UDiv ||
          Opc == Instruction::SRem || Opc == Instruction::SDiv);
 
   FastMathFlags FMF;
   FMF.setFast();
   Builder.setFastMathFlags(FMF);
 
   if (isa<Constant>(Den))
if (IsSigned) {
       Num = Builder.CreateSExt(Num, I32Ty);
       Den = Builder.CreateSExt(Den, I32Ty);
     } else {
       Num = Builder.CreateZExt(Num, I32Ty);
       Den = Builder.CreateZExt(Den, I32Ty);
     }
   }
 
-  if (Value *Res = expandDivRem24(Builder, Num, Den, IsDiv, IsSigned)) {
+  if (Value *Res = expandDivRem24(Builder, I, Num, Den, IsDiv, IsSigned)) {
     Res = Builder.CreateTrunc(Res, Ty);
     return Res;
   }
 
   ConstantInt *Zero = Builder.getInt32(0);
   ConstantInt *One = Builder.getInt32(1);
   ConstantInt *MinusOne = Builder.getInt32(~0);
 
if ((Opc == Instruction::URem || Opc == Instruction::UDiv ||
     Value *Num = I.getOperand(0);
     Value *Den = I.getOperand(1);
     IRBuilder<> Builder(&I);
     Builder.SetCurrentDebugLocation(I.getDebugLoc());
 
     if (VectorType *VT = dyn_cast<VectorType>(Ty)) {
       NewDiv = UndefValue::get(VT);
 
-      for (unsigned I = 0, E = VT->getNumElements(); I != E; ++I) {
-        Value *NumEltI = Builder.CreateExtractElement(Num, I);
-        Value *DenEltI = Builder.CreateExtractElement(Den, I);
-        Value *NewElt = expandDivRem32(Builder, Opc, NumEltI, DenEltI);
+      for (unsigned N = 0, E = VT->getNumElements(); N != E; ++N) {
+        Value *NumEltN = Builder.CreateExtractElement(Num, N);
+        Value *DenEltN = Builder.CreateExtractElement(Den, N);
+        Value *NewElt = expandDivRem32(Builder, I, NumEltN, DenEltN);
         if (!NewElt)
-          NewElt = Builder.CreateBinOp(Opc, NumEltI, DenEltI);
-        NewDiv = Builder.CreateInsertElement(NewDiv, NewElt, I);
+          NewElt = Builder.CreateBinOp(Opc, NumEltN, DenEltN);
+        NewDiv = Builder.CreateInsertElement(NewDiv, NewElt, N);
       }
     } else {
-      NewDiv = expandDivRem32(Builder, Opc, Num, Den);
+      NewDiv = expandDivRem32(Builder, I, Num, Den);
     }
 
     if (NewDiv) {
       I.replaceAllUsesWith(NewDiv);
       I.eraseFromParent();
       Changed = true;
     }
   }
if (skipFunction(F))
     return false;
 
   auto *TPC = getAnalysisIfAvailable<TargetPassConfig>();
   if (!TPC)
     return false;
 
   const AMDGPUTargetMachine &TM = TPC->getTM<AMDGPUTargetMachine>();
   ST = &TM.getSubtarget<GCNSubtarget>(F);
+  AC = &getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
   DA = &getAnalysis<DivergenceAnalysis>();
   HasUnsafeFPMath = hasUnsafeFPMath(F);
   AMDGPUASI = TM.getAMDGPUAS();
 
   bool MadeChange = false;
 
   for (BasicBlock &BB : F) {
     BasicBlock::iterator Next;
     for (BasicBlock::iterator I = BB.begin(), E = BB.end(); I != E; I = Next) {
       Next = std::next(I);
       MadeChange |= visit(*I);
     }
   }
 
   return MadeChange;
 }
 
 INITIALIZE_PASS_BEGIN(AMDGPUCodeGenPrepare, DEBUG_TYPE,
                       "AMDGPU IR optimizations", false, false)
+INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_DEPENDENCY(DivergenceAnalysis)
 INITIALIZE_PASS_END(AMDGPUCodeGenPrepare, DEBUG_TYPE, "AMDGPU IR optimizations",
                     false, false)
 
 char AMDGPUCodeGenPrepare::ID = 0;
 
 FunctionPass *llvm::createAMDGPUCodeGenPreparePass() {
   return new AMDGPUCodeGenPrepare();
 }

test/CodeGen/AMDGPU/divrem24-assume.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -mtriple=amdgcn-- -amdgpu-codegenprepare %s | FileCheck %s
+
+define amdgpu_kernel void @divrem24_assume(i32 addrspace(1)* %arg, i32 %arg1) {
+; CHECK-LABEL: @divrem24_assume(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    [[TMP:%.*]] = tail call i32 @llvm.amdgcn.workitem.id.x(), !range !0
+; CHECK-NEXT:    [[TMP2:%.*]] = icmp ult i32 [[ARG1:%.*]], 42
+; CHECK-NEXT:    tail call void @llvm.assume(i1 [[TMP2]])
+; CHECK-NEXT:    [[TMP0:%.*]] = uitofp i32 [[TMP]] to float
+; CHECK-NEXT:    [[TMP1:%.*]] = uitofp i32 [[ARG1]] to float
+; CHECK-NEXT:    [[TMP2:%.*]] = fdiv fast float 1.000000e+00, [[TMP1]]
+; CHECK-NEXT:    [[TMP3:%.*]] = fmul fast float [[TMP0]], [[TMP2]]
+; CHECK-NEXT:    [[TMP4:%.*]] = call fast float @llvm.trunc.f32(float [[TMP3]])
+; CHECK-NEXT:    [[TMP5:%.*]] = fsub fast float -0.000000e+00, [[TMP4]]
+; CHECK-NEXT:    [[TMP6:%.*]] = call fast float @llvm.amdgcn.fmad.ftz.f32(float [[TMP5]], float [[TMP1]], float [[TMP0]])
+; CHECK-NEXT:    [[TMP7:%.*]] = fptoui float [[TMP4]] to i32
+; CHECK-NEXT:    [[TMP8:%.*]] = call fast float @llvm.fabs.f32(float [[TMP6]])
+; CHECK-NEXT:    [[TMP9:%.*]] = call fast float @llvm.fabs.f32(float [[TMP1]])
+; CHECK-NEXT:    [[TMP10:%.*]] = fcmp fast oge float [[TMP8]], [[TMP9]]
+; CHECK-NEXT:    [[TMP11:%.*]] = select i1 [[TMP10]], i32 1, i32 0
+; CHECK-NEXT:    [[TMP12:%.*]] = add i32 [[TMP7]], [[TMP11]]
+; CHECK-NEXT:    [[TMP13:%.*]] = and i32 [[TMP12]], 1023
+; CHECK-NEXT:    [[TMP4:%.*]] = zext i32 [[TMP13]] to i64
+; CHECK-NEXT:    [[TMP5:%.*]] = getelementptr inbounds i32, i32 addrspace(1)* [[ARG:%.*]], i64 [[TMP4]]
+; CHECK-NEXT:    store i32 0, i32 addrspace(1)* [[TMP5]], align 4
+; CHECK-NEXT:    ret void
+;
+bb:
+  %tmp = tail call i32 @llvm.amdgcn.workitem.id.x(), !range !0
+  %tmp2 = icmp ult i32 %arg1, 42
+  tail call void @llvm.assume(i1 %tmp2)
+  %tmp3 = udiv i32 %tmp, %arg1
+  %tmp4 = zext i32 %tmp3 to i64
+  %tmp5 = getelementptr inbounds i32, i32 addrspace(1)* %arg, i64 %tmp4
+  store i32 0, i32 addrspace(1)* %tmp5, align 4
+  ret void
+}
+
+declare void @llvm.assume(i1)
+declare i32 @llvm.amdgcn.workitem.id.x()
+
+!0 = !{i32 0, i32 1024}