AMDGPU: Reduce readlane for single demanded vector element

Author: makslevental

llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp

@@ -635,7 +635,8 @@ GCNTTIImpl::instCombineIntrinsic(InstCombiner &IC, IntrinsicInst &II) const {
       break;
 
     auto IID = SrcCI->getIntrinsicID();
-    // llvm.amdgcn.rcp(llvm.amdgcn.sqrt(x)) -> llvm.amdgcn.rsq(x) if contractable
+    // llvm.amdgcn.rcp(llvm.amdgcn.sqrt(x)) -> llvm.amdgcn.rsq(x) if
+    // contractable
     //
     // llvm.amdgcn.rcp(llvm.sqrt(x)) -> llvm.amdgcn.rsq(x) if contractable and
     // relaxed.
@@ -845,13 +846,13 @@ GCNTTIImpl::instCombineIntrinsic(InstCombiner &IC, IntrinsicInst &II) const {
     break;
   }
   case Intrinsic::amdgcn_cvt_off_f32_i4: {
-    Value* Arg = II.getArgOperand(0);
+    Value *Arg = II.getArgOperand(0);
     Type *Ty = II.getType();
 
     if (isa<PoisonValue>(Arg))
       return IC.replaceInstUsesWith(II, PoisonValue::get(Ty));
 
-    if(IC.getSimplifyQuery().isUndefValue(Arg))
+    if (IC.getSimplifyQuery().isUndefValue(Arg))
       return IC.replaceInstUsesWith(II, Constant::getNullValue(Ty));
 
     ConstantInt *CArg = dyn_cast<ConstantInt>(II.getArgOperand(0));
@@ -1629,18 +1630,18 @@ GCNTTIImpl::instCombineIntrinsic(InstCombiner &IC, IntrinsicInst &II) const {
   }
   }
   if (const AMDGPU::ImageDimIntrinsicInfo *ImageDimIntr =
-            AMDGPU::getImageDimIntrinsicInfo(II.getIntrinsicID())) {
+          AMDGPU::getImageDimIntrinsicInfo(II.getIntrinsicID())) {
     return simplifyAMDGCNImageIntrinsic(ST, ImageDimIntr, II, IC);
   }
   return std::nullopt;
 }
 
 /// Implement SimplifyDemandedVectorElts for amdgcn buffer and image intrinsics.
 ///
-/// The result of simplifying amdgcn image and buffer store intrinsics is updating
-/// definitions of the intrinsics vector argument, not Uses of the result like
-/// image and buffer loads.
-/// Note: This only supports non-TFE/LWE image intrinsic calls; those have
+/// The result of simplifying amdgcn image and buffer store intrinsics is
+/// updating definitions of the intrinsics vector argument, not Uses of the
+/// result like image and buffer loads. Note: This only supports non-TFE/LWE
+/// image intrinsic calls; those have
 ///       struct returns.
 static Value *simplifyAMDGCNMemoryIntrinsicDemanded(InstCombiner &IC,
                                                     IntrinsicInst &II,
@@ -1837,7 +1838,12 @@ Value *GCNTTIImpl::simplifyAMDGCNLaneIntrinsicDemanded(
     Value *Extract = IC.Builder.CreateExtractElement(Src, FirstElt);
 
     // TODO: Preserve callsite attributes?
-    CallInst *NewCall = IC.Builder.CreateCall(Remangled, {Extract}, OpBundles);
+    SmallVector<Value *> Args{Extract};
+    if (II.arg_size() > 1) {
+      for (int I = 1; I < II.arg_size(); ++I)
+        Args.push_back(II.getArgOperand(1));
+    }
+    CallInst *NewCall = IC.Builder.CreateCall(Remangled, Args, OpBundles);
 
     return IC.Builder.CreateInsertElement(PoisonValue::get(II.getType()),
                                           NewCall, FirstElt);
@@ -1872,6 +1878,7 @@ std::optional<Value *> GCNTTIImpl::simplifyDemandedVectorEltsIntrinsic(
         SimplifyAndSetOp) const {
   switch (II.getIntrinsicID()) {
   case Intrinsic::amdgcn_readfirstlane:
+  case Intrinsic::amdgcn_readlane:
     SimplifyAndSetOp(&II, 0, DemandedElts, UndefElts);
     return simplifyAMDGCNLaneIntrinsicDemanded(IC, II, DemandedElts, UndefElts);
   case Intrinsic::amdgcn_raw_buffer_load:

llvm/test/Transforms/InstCombine/AMDGPU/simplify-demanded-vector-elts-lane-intrinsics.ll

@@ -25,6 +25,42 @@ define i16 @extract_elt0_v1i16_readfirstlane(<1 x i16> %src) {
   ret i16 %elt
 }
 
+define half @extract_elt0_v2f16_readlane_imm_0(<2 x half> %src) {
+; CHECK-LABEL: define half @extract_elt0_v2f16_readlane_imm_0(
+; CHECK-SAME: <2 x half> [[SRC:%.*]]) #[[ATTR0]] {
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <2 x half> [[SRC]], i64 0
+; CHECK-NEXT:    [[TMP2:%.*]] = call half @llvm.amdgcn.readlane.f16(half [[TMP1]], i32 0)
+; CHECK-NEXT:    ret half [[TMP2]]
+;
+  %x = call <2 x half> @llvm.amdgcn.readlane.v2f16(<2 x half> %src, i32 0)
+  %elt = extractelement <2 x half> %x, i32 0
+  ret half %elt
+}
+
+define half @extract_elt0_v2f16_readlane_imm_1(<2 x half> %src) {
+; CHECK-LABEL: define half @extract_elt0_v2f16_readlane_imm_1(
+; CHECK-SAME: <2 x half> [[SRC:%.*]]) #[[ATTR0]] {
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <2 x half> [[SRC]], i64 1
+; CHECK-NEXT:    [[TMP2:%.*]] = call half @llvm.amdgcn.readlane.f16(half [[TMP1]], i32 1)
+; CHECK-NEXT:    ret half [[TMP2]]
+;
+  %x = call <2 x half> @llvm.amdgcn.readlane.v2f16(<2 x half> %src, i32 1)
+  %elt = extractelement <2 x half> %x, i32 1
+  ret half %elt
+}
+
+define half @extract_elt0_v2f16_readlane(<2 x half> %src, i32 %idx) {
+; CHECK-LABEL: define half @extract_elt0_v2f16_readlane(
+; CHECK-SAME: <2 x half> [[SRC:%.*]], i32 [[IDX:%.*]]) #[[ATTR0]] {
+; CHECK-NEXT:    [[X:%.*]] = call <2 x half> @llvm.amdgcn.readlane.v2f16(<2 x half> [[SRC]], i32 [[IDX]])
+; CHECK-NEXT:    [[ELT:%.*]] = extractelement <2 x half> [[X]], i32 [[IDX]]
+; CHECK-NEXT:    ret half [[ELT]]
+;
+  %x = call <2 x half> @llvm.amdgcn.readlane.v2f16(<2 x half> %src, i32 %idx)
+  %elt = extractelement <2 x half> %x, i32 %idx
+  ret half %elt
+}
+
 define i16 @extract_elt1_v2i16_readfirstlane(<2 x i16> %src) {
 ; CHECK-LABEL: define i16 @extract_elt1_v2i16_readfirstlane(
 ; CHECK-SAME: <2 x i16> [[SRC:%.*]]) #[[ATTR0]] {