AMDGPU: Fold readlane/readfirstlane calls

[arsenm]

Fix extra whitespace change

[nhaehnle]

Good idea, but needs a fix I think.

lib/Transforms/InstCombine/InstCombineCalls.cpp

3787–3798

I believe all of these need to be guarded by being in the same basic block, since the user could be observing a read(first)lane inside a divergent loop.

[arsenm]

Check parent

[arsenm]

Rebased into wrong patch

[nhaehnle]

LGTM

[arsenm]

r363587

lib/Transforms/InstCombine/InstCombineCalls.cpp

case Intrinsic::amdgcn_update_dpp: {
     II->setOperand(0, UndefValue::get(Old->getType()));
     return II;
   }
   case Intrinsic::amdgcn_readfirstlane:
   case Intrinsic::amdgcn_readlane: {
     // A constant value is trivially uniform.
     if (Constant *C = dyn_cast<Constant>(II->getArgOperand(0)))
       return replaceInstUsesWith(*II, C);
+
+    // The rest of these may not be safe if the exec may not be the same between
+    // the def and use.
+    Value *Src = II->getArgOperand(0);
+    Instruction *SrcInst = dyn_cast<Instruction>(Src);
+    if (SrcInst && SrcInst->getParent() != II->getParent())
+      break;
+
+    // readfirstlane (readfirstlane x) -> readfirstlane x
+    // readlane (readfirstlane x), y -> readfirstlane x
+    if (match(Src, m_Intrinsic<Intrinsic::amdgcn_readfirstlane>()))
+      return replaceInstUsesWith(*II, Src);
+
+    if (IID == Intrinsic::amdgcn_readfirstlane) {
+      // readfirstlane (readlane x, y) -> readlane x, y
+      if (match(Src, m_Intrinsic<Intrinsic::amdgcn_readlane>()))
+        return replaceInstUsesWith(*II, Src);
+    } else {
+      // readlane (readlane x, y), y -> readlane x, y
+      if (match(Src, m_Intrinsic<Intrinsic::amdgcn_readlane>(
+                  m_Value(), m_Specific(II->getArgOperand(1)))))
+        return replaceInstUsesWith(*II, Src);
+    }
+
     break;
   }
   case Intrinsic::stackrestore: {
     // If the save is right next to the restore, remove the restore.  This can
     // happen when variable allocas are DCE'd.
     if (IntrinsicInst *SS = dyn_cast<IntrinsicInst>(II->getArgOperand(0))) {
       if (SS->getIntrinsicID() == Intrinsic::stacksave) {
         // Skip over debug info.

test/Transforms/InstCombine/AMDGPU/amdgcn-intrinsics.ll

;
   store volatile i32 %var, i32* undef
   store volatile i32 %zero, i32* undef
   store volatile i32 %imm, i32* undef
   store volatile i32 %constexpr, i32* undef
   store volatile i32 %undef, i32* undef
   ret void
 }
 
+define i32 @readfirstlane_idempotent(i32 %arg) {
+; CHECK-LABEL: @readfirstlane_idempotent(
+; CHECK-NEXT:    [[READ0:%.*]] = call i32 @llvm.amdgcn.readfirstlane(i32 [[ARG:%.*]])
+; CHECK-NEXT:    ret i32 [[READ0]]
+;
+  %read0 = call i32 @llvm.amdgcn.readfirstlane(i32 %arg)
+  %read1 = call i32 @llvm.amdgcn.readfirstlane(i32 %read0)
+  %read2 = call i32 @llvm.amdgcn.readfirstlane(i32 %read1)
+  ret i32 %read2
+}
+
+define i32 @readfirstlane_readlane(i32 %arg) {
+; CHECK-LABEL: @readfirstlane_readlane(
+; CHECK-NEXT:    [[READ0:%.*]] = call i32 @llvm.amdgcn.readfirstlane(i32 [[ARG:%.*]])
+; CHECK-NEXT:    ret i32 [[READ0]]
+;
+  %read0 = call i32 @llvm.amdgcn.readfirstlane(i32 %arg)
+  %read1 = call i32 @llvm.amdgcn.readlane(i32 %read0, i32 0)
+  ret i32 %read1
+}
+
+define i32 @readfirstlane_readfirstlane_different_block(i32 %arg) {
+; CHECK-LABEL: @readfirstlane_readfirstlane_different_block(
+; CHECK-NEXT:  bb0:
+; CHECK-NEXT:    [[READ0:%.*]] = call i32 @llvm.amdgcn.readfirstlane(i32 [[ARG:%.*]])
+; CHECK-NEXT:    br label [[BB1:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    [[READ1:%.*]] = call i32 @llvm.amdgcn.readfirstlane(i32 [[READ0]])
+; CHECK-NEXT:    ret i32 [[READ1]]
+;
+bb0:
+  %read0 = call i32 @llvm.amdgcn.readfirstlane(i32 %arg)
+  br label %bb1
+
+bb1:
+  %read1 = call i32 @llvm.amdgcn.readfirstlane(i32 %read0)
+  ret i32 %read1
+}
+
+define i32 @readfirstlane_readlane_different_block(i32 %arg) {
+; CHECK-LABEL: @readfirstlane_readlane_different_block(
+; CHECK-NEXT:  bb0:
+; CHECK-NEXT:    [[READ0:%.*]] = call i32 @llvm.amdgcn.readlane(i32 [[ARG:%.*]], i32 0)
+; CHECK-NEXT:    br label [[BB1:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    [[READ1:%.*]] = call i32 @llvm.amdgcn.readfirstlane(i32 [[READ0]])
+; CHECK-NEXT:    ret i32 [[READ1]]
+;
+bb0:
+  %read0 = call i32 @llvm.amdgcn.readlane(i32 %arg, i32 0)
+  br label %bb1
+
+bb1:
+  %read1 = call i32 @llvm.amdgcn.readfirstlane(i32 %read0)
+  ret i32 %read1
+}
+
 ; --------------------------------------------------------------------
 ; llvm.amdgcn.readlane
 ; --------------------------------------------------------------------
 
 declare i32 @llvm.amdgcn.readlane(i32, i32)
 
 define amdgpu_kernel void @readlane_constant(i32 %arg, i32 %lane) {
 ; CHECK-LABEL: @readlane_constant(
;
   store volatile i32 %var, i32* undef
   store volatile i32 %zero, i32* undef
   store volatile i32 %imm, i32* undef
   store volatile i32 %constexpr, i32* undef
   store volatile i32 %undef, i32* undef
   ret void
 }
 
+define i32 @readlane_idempotent(i32 %arg, i32 %lane) {
+; CHECK-LABEL: @readlane_idempotent(
+; CHECK-NEXT:    [[READ0:%.*]] = call i32 @llvm.amdgcn.readlane(i32 [[ARG:%.*]], i32 [[LANE:%.*]])
+; CHECK-NEXT:    ret i32 [[READ0]]
+;
+  %read0 = call i32 @llvm.amdgcn.readlane(i32 %arg, i32 %lane)
+  %read1 = call i32 @llvm.amdgcn.readlane(i32 %read0, i32 %lane)
+  ret i32 %read1
+}
+
+define i32 @readlane_idempotent_different_lanes(i32 %arg, i32 %lane0, i32 %lane1) {
+; CHECK-LABEL: @readlane_idempotent_different_lanes(
+; CHECK-NEXT:    [[READ0:%.*]] = call i32 @llvm.amdgcn.readlane(i32 [[ARG:%.*]], i32 [[LANE0:%.*]])
+; CHECK-NEXT:    [[READ1:%.*]] = call i32 @llvm.amdgcn.readlane(i32 [[READ0]], i32 [[LANE1:%.*]])
+; CHECK-NEXT:    ret i32 [[READ1]]
+;
+  %read0 = call i32 @llvm.amdgcn.readlane(i32 %arg, i32 %lane0)
+  %read1 = call i32 @llvm.amdgcn.readlane(i32 %read0, i32 %lane1)
+  ret i32 %read1
+}
+
+define i32 @readlane_readfirstlane(i32 %arg) {
+; CHECK-LABEL: @readlane_readfirstlane(
+; CHECK-NEXT:    [[READ0:%.*]] = call i32 @llvm.amdgcn.readfirstlane(i32 [[ARG:%.*]])
+; CHECK-NEXT:    ret i32 [[READ0]]
+;
+  %read0 = call i32 @llvm.amdgcn.readfirstlane(i32 %arg)
+  %read1 = call i32 @llvm.amdgcn.readlane(i32 %read0, i32 0)
+  ret i32 %read1
+}
+
+define i32 @readlane_idempotent_different_block(i32 %arg, i32 %lane) {
+; CHECK-LABEL: @readlane_idempotent_different_block(
+; CHECK-NEXT:  bb0:
+; CHECK-NEXT:    [[READ0:%.*]] = call i32 @llvm.amdgcn.readlane(i32 [[ARG:%.*]], i32 [[LANE:%.*]])
+; CHECK-NEXT:    br label [[BB1:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    [[READ1:%.*]] = call i32 @llvm.amdgcn.readlane(i32 [[READ0]], i32 [[LANE]])
+; CHECK-NEXT:    ret i32 [[READ1]]
+;
+bb0:
+  %read0 = call i32 @llvm.amdgcn.readlane(i32 %arg, i32 %lane)
+  br label %bb1
+
+bb1:
+  %read1 = call i32 @llvm.amdgcn.readlane(i32 %read0, i32 %lane)
+  ret i32 %read1
+}
+
+
+define i32 @readlane_readfirstlane_different_block(i32 %arg) {
+; CHECK-LABEL: @readlane_readfirstlane_different_block(
+; CHECK-NEXT:  bb0:
+; CHECK-NEXT:    [[READ0:%.*]] = call i32 @llvm.amdgcn.readfirstlane(i32 [[ARG:%.*]])
+; CHECK-NEXT:    br label [[BB1:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    [[READ1:%.*]] = call i32 @llvm.amdgcn.readlane(i32 [[READ0]], i32 0)
+; CHECK-NEXT:    ret i32 [[READ1]]
+;
+bb0:
+  %read0 = call i32 @llvm.amdgcn.readfirstlane(i32 %arg)
+  br label %bb1
+
+bb1:
+  %read1 = call i32 @llvm.amdgcn.readlane(i32 %read0, i32 0)
+  ret i32 %read1
+}
+
 ; --------------------------------------------------------------------
 ; llvm.amdgcn.update.dpp.i32
 ; --------------------------------------------------------------------
 
 declare i32 @llvm.amdgcn.update.dpp.i32(i32, i32, i32, i32, i32, i1)
 
 define amdgpu_kernel void @update_dpp_no_combine(i32 addrspace(1)* %out, i32 %in1, i32 %in2) {
 ; CHECK-LABEL: @update_dpp_no_combine(