[AMDGPU] Check if type transforms to i16 (VI+) when getting AMDGPUISD::FFBH_U32

[kzhuravl]

Summary

This will prevent following regression when enabling i16 support (D18049):

test/CodeGen/AMDGPU/ctlz.ll
test/CodeGen/AMDGPU/ctlz_zero_undef.ll

[• tstellarAMD]

lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp

482

Shouldn't we only be doing this for uniform ops? Otherwise the SelectionDAG should be able to do the promotion.

[kzhuravl]

We do not need to promote it. Fixed during combining.

[arsenm]

lib/Target/AMDGPU/AMDGPUISelLowering.cpp

2461 ↗

(On Diff #75500)

Subtarget->hasI16

2468 ↗

(On Diff #75500)

Variable name probably shouldn't match the node type name

lib/Target/AMDGPU/AMDGPUISelLowering.h

33 ↗

(On Diff #75500)

I think SDLoc is supposed to be passed by const reference now

[kzhuravl]

lib/Target/AMDGPU/AMDGPUISelLowering.cpp

2461 ↗

(On Diff #75500)

This will involve moving has16BitInsts() from SISubtarget to AMDGPUSubtarget. Are there any objections?

[kzhuravl]

Address review feedback

[kzhuravl]

Ping

[• tstellarAMD]

LGTM.

lib/Target/AMDGPU/AMDGPUCodeGenPrepare.cpp

class AMDGPUCodeGenPrepare : public FunctionPass,
   /// bits, replacing \p I with 32 bit 'bitreverse' intrinsic, shifting the
   /// 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;
 
+  /// \brief Promotes 'ctlz' intrinsic \p I to 32 bit 'ctlz' intrinsic.
+  ///
+  /// \details \p I's base element bit width must be greater than 1 and less
+  /// than or equal 16. Promotion is done by zero extending the first operand to
+  /// 32 bits, replacing \p I with 32 bit 'ctlz' intrinsic, and truncating the
+  /// result of 32 bit 'ctlz' intrinsic back to \p I's original type.
+  ///
+  /// \returns True.
+  bool promoteCtlzToI32(IntrinsicInst &I) const;
+
 public:
   static char ID;
   AMDGPUCodeGenPrepare(const TargetMachine *TM = nullptr) :
     FunctionPass(ID),
     TM(static_cast<const GCNTargetMachine *>(TM)),
     ST(nullptr),
     DA(nullptr),
     Mod(nullptr),
     HasUnsafeFPMath(false) { }
 
   bool visitFDiv(BinaryOperator &I);
 
   bool visitInstruction(Instruction &I) { return false; }
   bool visitBinaryOperator(BinaryOperator &I);
   bool visitICmpInst(ICmpInst &I);
   bool visitSelectInst(SelectInst &I);
 
   bool visitIntrinsicInst(IntrinsicInst &I);
   bool visitBitreverseIntrinsicInst(IntrinsicInst &I);
+  bool visitCtlzIntrinsicInst(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<DivergenceAnalysis>();
Value *TruncRes =
       Builder.CreateTrunc(LShrOp, I.getType());
 
   I.replaceAllUsesWith(TruncRes);
   I.eraseFromParent();
 
   return true;
 }
 
+bool AMDGPUCodeGenPrepare::promoteCtlzToI32(IntrinsicInst &I) const {
+  assert(I.getIntrinsicID() == Intrinsic::ctlz &&
+         "I must be ctlz intrinsic");
+  assert(needsPromotionToI32(I.getType()) &&
+         "I does not need promotion to i32");
+
+  IRBuilder<> Builder(&I);
+  Builder.SetCurrentDebugLocation(I.getDebugLoc());
+
+  Type *I32Ty = getI32Ty(Builder, I.getType());
+  Function *I32 =
+      Intrinsic::getDeclaration(Mod, Intrinsic::ctlz, { I32Ty });
+  Value *ExtOp = Builder.CreateZExt(I.getOperand(0), I32Ty);
+  Value *ExtRes = Builder.CreateCall(I32, { ExtOp, I.getOperand(1) });
+  Value *TruncRes = Builder.CreateTrunc(ExtRes, I.getType());
+
+  I.replaceAllUsesWith(TruncRes);
+  I.eraseFromParent();
+
+  return true;
+}
+
 static bool shouldKeepFDivF32(Value *Num, bool UnsafeDiv) {
   const ConstantFP *CNum = dyn_cast<ConstantFP>(Num);
   if (!CNum)
     return false;
 
   // Reciprocal f32 is handled separately without denormals.
   return UnsafeDiv || CNum->isExactlyValue(+1.0);
 }
bool AMDGPUCodeGenPrepare::visitSelectInst(SelectInst &I) {
 
   return Changed;
 }
 
 bool AMDGPUCodeGenPrepare::visitIntrinsicInst(IntrinsicInst &I) {
   switch (I.getIntrinsicID()) {
   case Intrinsic::bitreverse:
     return visitBitreverseIntrinsicInst(I);
+  case Intrinsic::ctlz:
+    return visitCtlzIntrinsicInst(I);
   default:
     return false;
   }
 }
 
 bool AMDGPUCodeGenPrepare::visitBitreverseIntrinsicInst(IntrinsicInst &I) {
   bool Changed = false;
 
   if (ST->has16BitInsts() && needsPromotionToI32(I.getType()) &&
       DA->isUniform(&I))
     Changed |= promoteUniformBitreverseToI32(I);
 
   return Changed;
 }
 
+bool AMDGPUCodeGenPrepare::visitCtlzIntrinsicInst(IntrinsicInst &I) {
+  bool Changed = false;
+
+  if (needsPromotionToI32(I.getType()))
+    Changed |= promoteCtlzToI32(I);
+
+  return Changed;
+}
+
 bool AMDGPUCodeGenPrepare::doInitialization(Module &M) {
   Mod = &M;
   return false;
 }
 
 bool AMDGPUCodeGenPrepare::runOnFunction(Function &F) {
   if (!TM || skipFunction(F))
     return false;

test/CodeGen/AMDGPU/amdgpu-codegenprepare-i16-to-i32.ll

 ; VI-NEXT: %[[S_32:[0-9]+]] = lshr i32 %[[R_32]], 16
 ; VI-NEXT: %[[R_16:[0-9]+]] = trunc i32 %[[S_32]] to i16
 ; VI-NEXT: ret i16 %[[R_16]]
 define i16 @bitreverse_i16(i16 %a) {
   %brev = call i16 @llvm.bitreverse.i16(i16 %a)
   ret i16 %brev
 }
 
+declare i16 @llvm.ctlz.i16(i16, i1)
+; GCN-LABEL: @ctlz_i16_false(
+; GCN: %[[A_32:[0-9]+]] = zext i16 %a to i32
+; GCN-NEXT: %[[R_32:[0-9]+]] = call i32 @llvm.ctlz.i32(i32 %[[A_32]], i1 false)
+; GCN-NEXT: %[[R_16:[0-9]+]] = trunc i32 %[[R_32]] to i16
+; GCN-NEXT: ret i16 %[[R_16]]
+define i16 @ctlz_i16_false(i16 %a) {
+  %ctlz = call i16 @llvm.ctlz.i16(i16 %a, i1 false)
+  ret i16 %ctlz
+}
+
+; GCN-LABEL: @ctlz_i16_true(
+; GCN: %[[A_32:[0-9]+]] = zext i16 %a to i32
+; GCN-NEXT: %[[R_32:[0-9]+]] = call i32 @llvm.ctlz.i32(i32 %[[A_32]], i1 true)
+; GCN-NEXT: %[[R_16:[0-9]+]] = trunc i32 %[[R_32]] to i16
+; GCN-NEXT: ret i16 %[[R_16]]
+define i16 @ctlz_i16_true(i16 %a) {
+  %ctlz = call i16 @llvm.ctlz.i16(i16 %a, i1 true)
+  ret i16 %ctlz
+}
+
 ; GCN-LABEL: @add_3xi15(
 ; SI: %r = add <3 x i15> %a, %b
 ; SI-NEXT: ret <3 x i15> %r
 ; VI: %[[A_32:[0-9]+]] = zext <3 x i15> %a to <3 x i32>
 ; VI-NEXT: %[[B_32:[0-9]+]] = zext <3 x i15> %b to <3 x i32>
 ; VI-NEXT: %[[R_32:[0-9]+]] = add <3 x i32> %[[A_32]], %[[B_32]]
 ; VI-NEXT: %[[R_15:[0-9]+]] = trunc <3 x i32> %[[R_32]] to <3 x i15>
 ; VI-NEXT: ret <3 x i15> %[[R_15]]

 ; VI-NEXT: %[[R_32:[0-9]+]] = call <3 x i32> @llvm.bitreverse.v3i32(<3 x i32> %[[A_32]])
 ; VI-NEXT: %[[S_32:[0-9]+]] = lshr <3 x i32> %[[R_32]], <i32 16, i32 16, i32 16>
 ; VI-NEXT: %[[R_16:[0-9]+]] = trunc <3 x i32> %[[S_32]] to <3 x i16>
 ; VI-NEXT: ret <3 x i16> %[[R_16]]
 define <3 x i16> @bitreverse_3xi16(<3 x i16> %a) {
   %brev = call <3 x i16> @llvm.bitreverse.v3i16(<3 x i16> %a)
   ret <3 x i16> %brev
 }
+
+declare <3 x i16> @llvm.ctlz.v3i16(<3 x i16>, i1)
+; GCN-LABEL: @ctlz_v3i16_false(
+; GCN: %[[A_32:[0-9]+]] = zext <3 x i16> %a to <3 x i32>
+; GCN-NEXT: %[[R_32:[0-9]+]] = call <3 x i32> @llvm.ctlz.v3i32(<3 x i32> %[[A_32]], i1 false)
+; GCN-NEXT: %[[R_16:[0-9]+]] = trunc <3 x i32> %[[R_32]] to <3 x i16>
+; GCN-NEXT: ret <3 x i16> %[[R_16]]
+define <3 x i16> @ctlz_v3i16_false(<3 x i16> %a) {
+  %ctlz = call <3 x i16> @llvm.ctlz.v3i16(<3 x i16> %a, i1 false)
+  ret <3 x i16> %ctlz
+}
+
+; GCN-LABEL: @ctlz_v3i16_true(
+; GCN: %[[A_32:[0-9]+]] = zext <3 x i16> %a to <3 x i32>
+; GCN-NEXT: %[[R_32:[0-9]+]] = call <3 x i32> @llvm.ctlz.v3i32(<3 x i32> %[[A_32]], i1 true)
+; GCN-NEXT: %[[R_16:[0-9]+]] = trunc <3 x i32> %[[R_32]] to <3 x i16>
+; GCN-NEXT: ret <3 x i16> %[[R_16]]
+define <3 x i16> @ctlz_v3i16_true(<3 x i16> %a) {
+  %ctlz = call <3 x i16> @llvm.ctlz.v3i16(<3 x i16> %a, i1 true)
+  ret <3 x i16> %ctlz
+}

test/CodeGen/AMDGPU/ctlz.ll

 define void @v_ctlz_v4i32(<4 x i32> addrspace(1)* noalias %out, <4 x i32> addrspace(1)* noalias %valptr) nounwind {
   %val = load <4 x i32>, <4 x i32> addrspace(1)* %valptr, align 16
   %ctlz = call <4 x i32> @llvm.ctlz.v4i32(<4 x i32> %val, i1 false) nounwind readnone
   store <4 x i32> %ctlz, <4 x i32> addrspace(1)* %out, align 16
   ret void
 }
 
 ; FUNC-LABEL: {{^}}v_ctlz_i8:
-; GCN: buffer_load_ubyte [[VAL:v[0-9]+]],
-; GCN-DAG: v_ffbh_u32_e32 [[FFBH:v[0-9]+]], [[VAL]]
-; GCN-DAG: v_cmp_eq_u32_e32 vcc, 0, [[CTLZ]]
-; GCN-DAG: v_cndmask_b32_e64 [[CORRECTED_FFBH:v[0-9]+]], [[FFBH]], 32, vcc
-; GCN: v_add_i32_e32 [[RESULT:v[0-9]+]], vcc, 0xffffffe8, [[CORRECTED_FFBH]]
-; GCN: buffer_store_byte [[RESULT]],
+; GCN: {{buffer|flat}}_load_ubyte [[VAL:v[0-9]+]],
+; GCN: v_ffbh_u32_e32 [[FFBH:v[0-9]+]], [[VAL]]
+; GCN: {{buffer|flat}}_store_byte [[FFBH]],
 define void @v_ctlz_i8(i8 addrspace(1)* noalias %out, i8 addrspace(1)* noalias %valptr) nounwind {
   %val = load i8, i8 addrspace(1)* %valptr
   %ctlz = call i8 @llvm.ctlz.i8(i8 %val, i1 false) nounwind readnone
   store i8 %ctlz, i8 addrspace(1)* %out
   ret void
 }
 
 ; FUNC-LABEL: {{^}}s_ctlz_i64:
define void @v_ctlz_i16_sel_eq_neg1(i16 addrspace(1)* noalias %out, i16 addrspace(1)* noalias %valptr) nounwind {
   %val = load i16, i16 addrspace(1)* %valptr
   %ctlz = call i16 @llvm.ctlz.i16(i16 %val, i1 false) nounwind readnone
   %cmp = icmp eq i16 %val, 0
   %sel = select i1 %cmp, i16 -1, i16 %ctlz
   store i16 %sel, i16 addrspace(1)* %out
   ret void
 }
 
-; FIXME: Need to handle non-uniform case for function below (load without gep).
 ; FUNC-LABEL: {{^}}v_ctlz_i7_sel_eq_neg1:
 ; GCN: {{buffer|flat}}_load_ubyte [[VAL:v[0-9]+]],
 ; GCN: v_ffbh_u32_e32 [[FFBH:v[0-9]+]], [[VAL]]
-; GCN: v_and_b32_e32 [[TRUNC:v[0-9]+]], 0x7f, [[FFBH]]
-; GCN: {{buffer|flat}}_store_byte [[TRUNC]],
+; GCN: {{buffer|flat}}_store_byte [[FFBH]],
 define void @v_ctlz_i7_sel_eq_neg1(i7 addrspace(1)* noalias %out, i7 addrspace(1)* noalias %valptr) nounwind {
   %tid = call i32 @llvm.r600.read.tidig.x()
   %valptr.gep = getelementptr i7, i7 addrspace(1)* %valptr, i32 %tid
   %val = load i7, i7 addrspace(1)* %valptr.gep
   %ctlz = call i7 @llvm.ctlz.i7(i7 %val, i1 false) nounwind readnone
   %cmp = icmp eq i7 %val, 0
   %sel = select i1 %cmp, i7 -1, i7 %ctlz
   store i7 %sel, i7 addrspace(1)* %out
   ret void
 }

test/CodeGen/AMDGPU/ctlz_zero_undef.ll

 define void @v_ctlz_zero_undef_v4i32(<4 x i32> addrspace(1)* noalias %out, <4 x i32> addrspace(1)* noalias %valptr) nounwind {
   %val = load <4 x i32>, <4 x i32> addrspace(1)* %valptr, align 16
   %ctlz = call <4 x i32> @llvm.ctlz.v4i32(<4 x i32> %val, i1 true) nounwind readnone
   store <4 x i32> %ctlz, <4 x i32> addrspace(1)* %out, align 16
   ret void
 }
 
 ; FUNC-LABEL: {{^}}v_ctlz_zero_undef_i8:
-; SI: buffer_load_ubyte [[VAL:v[0-9]+]],
+; SI: {{buffer|flat}}_load_ubyte [[VAL:v[0-9]+]],
 ; SI: v_ffbh_u32_e32 [[FFBH:v[0-9]+]], [[VAL]]
-; SI: v_add_i32_e32 [[RESULT:v[0-9]+]], vcc, 0xffffffe8, [[FFBH]]
-; SI: buffer_store_byte [[RESULT]],
+; SI: {{buffer|flat}}_store_byte [[FFBH]],
 define void @v_ctlz_zero_undef_i8(i8 addrspace(1)* noalias %out, i8 addrspace(1)* noalias %valptr) nounwind {
   %val = load i8, i8 addrspace(1)* %valptr
   %ctlz = call i8 @llvm.ctlz.i8(i8 %val, i1 true) nounwind readnone
   store i8 %ctlz, i8 addrspace(1)* %out
   ret void
 }
 
 ; FUNC-LABEL: {{^}}s_ctlz_zero_undef_i64: