[AMDGPU] Do not generate mul with 1 in AMDGPU Atomic Optimizer

[mbrkusanin]

Summary

Check if operand of mul is constant value of one for certain atomic
instructions in order to avoid making unnecessary instructions when
-amdgpu-atomic-optimizer is present.

[arsenm]

llvm/test/CodeGen/AMDGPU/atomic_optimizations_mul_one.mir

1 ↗

(On Diff #294336)

It's weird to use MIR and -run-pass for an IR pass. Should have regular ll test, with an end to end llc line and an opt line for just the pass

[foad]

llvm/lib/Target/AMDGPU/AMDGPUAtomicOptimizer.cpp

407

It would be neater to have a BuildMul helper function that checks for and optimizes the mul-by-1 case.

[mbrkusanin]

• Changed test to opt + llc global-isel (SDag was already removing these mul instructions, just later on). Should I split the file into two tests?

[mbrkusanin]

llvm/lib/Target/AMDGPU/AMDGPUAtomicOptimizer.cpp

407

A new BuildMul or update the current one (looks like some unittests need to be updated in this case)?

[foad]

llvm/lib/Target/AMDGPU/AMDGPUAtomicOptimizer.cpp

407

I was suggesting a new BuildMul helper function in this source file, not changing IRBuilder at all.

[mbrkusanin]

• Added buildMul that checks for mul with 1.

[arsenm]

llvm/test/CodeGen/AMDGPU/GlobalISel/atomic_optimizations_mul_one.ll

8–11

These negative checks are extremely fragile. These should use positive checks

[mbrkusanin]

• Updated llc test check lines.

llvm/test/CodeGen/AMDGPU/GlobalISel/atomic_optimizations_mul_one.ll

8–11

Just llc or opt check lines as well?

[arsenm]

llvm/test/CodeGen/AMDGPU/GlobalISel/atomic_optimizations_mul_one.ll

8–11

Both. Negative checks are just too easy to get wrong

[mbrkusanin]

• Updated opt test check lines.

llvm/lib/Target/AMDGPU/AMDGPUAtomicOptimizer.cpp

case AtomicRMWInst::UMin:
     return APInt::getMaxValue(BitWidth);
   case AtomicRMWInst::Max:
     return APInt::getSignedMinValue(BitWidth);
   case AtomicRMWInst::Min:
     return APInt::getSignedMaxValue(BitWidth);
   }
 }
 
+static Value *buildMul(IRBuilder<> &B, Value *LHS, Value *RHS) {
+  const ConstantInt *CI = dyn_cast<ConstantInt>(LHS);
+  return (CI && CI->isOne()) ? RHS : B.CreateMul(LHS, RHS);
+}
+
 void AMDGPUAtomicOptimizer::optimizeAtomic(Instruction &I,
                                            AtomicRMWInst::BinOp Op,
                                            unsigned ValIdx,
                                            bool ValDivergent) const {
   // Start building just before the instruction.
   IRBuilder<> B(&I);
 
   // If we are in a pixel shader, because of how we have to mask out helper
default:
       llvm_unreachable("Unhandled atomic op");
 
     case AtomicRMWInst::Add:
     case AtomicRMWInst::Sub: {
       // The new value we will be contributing to the atomic operation is the
       // old value times the number of active lanes.
       Value *const Ctpop = B.CreateIntCast(
           B.CreateUnaryIntrinsic(Intrinsic::ctpop, Ballot), Ty, false);
-      NewV = B.CreateMul(V, Ctpop);
+      NewV = buildMul(B, V, Ctpop);
       break;
     }
 
     case AtomicRMWInst::And:
     case AtomicRMWInst::Or:
     case AtomicRMWInst::Max:
     case AtomicRMWInst::Min:
     case AtomicRMWInst::UMax:
     case AtomicRMWInst::UMin:
       // These operations with a uniform value are idempotent: doing the atomic
       // operation multiple times has the same effect as doing it once.
       NewV = V;
       break;
 
     case AtomicRMWInst::Xor:
       // The new value we will be contributing to the atomic operation is the
       // old value times the parity of the number of active lanes.
       Value *const Ctpop = B.CreateIntCast(
           B.CreateUnaryIntrinsic(Intrinsic::ctpop, Ballot), Ty, false);
-      NewV = B.CreateMul(V, B.CreateAnd(Ctpop, 1));
+      NewV = buildMul(B, V, B.CreateAnd(Ctpop, 1));
       break;
     }
   }
 
   // We only want a single lane to enter our new control flow, and we do this
   // by checking if there are any active lanes below us. Only one lane will
   // have 0 active lanes below us, so that will be the only one to progress.
   Value *const Cond = B.CreateICmpEQ(Mbcnt, B.getIntN(TyBitWidth, 0));
if (NeedResult) {
     if (ValDivergent) {
       LaneOffset = B.CreateIntrinsic(Intrinsic::amdgcn_wwm, Ty, ExclScan);
     } else {
       switch (Op) {
       default:
         llvm_unreachable("Unhandled atomic op");
       case AtomicRMWInst::Add:
       case AtomicRMWInst::Sub:
-        LaneOffset = B.CreateMul(V, Mbcnt);
+        LaneOffset = buildMul(B, V, Mbcnt);
         break;
       case AtomicRMWInst::And:
       case AtomicRMWInst::Or:
       case AtomicRMWInst::Max:
       case AtomicRMWInst::Min:
       case AtomicRMWInst::UMax:
       case AtomicRMWInst::UMin:
         LaneOffset = B.CreateSelect(Cond, Identity, V);
         break;
       case AtomicRMWInst::Xor:
-        LaneOffset = B.CreateMul(V, B.CreateAnd(Mbcnt, 1));
+        LaneOffset = buildMul(B, V, B.CreateAnd(Mbcnt, 1));
         break;
       }
     }
     Value *const Result = buildNonAtomicBinOp(B, Op, BroadcastI, LaneOffset);
 
     if (IsPixelShader) {
       // Need a final PHI to reconverge to above the helper lane branch mask.
       B.SetInsertPoint(PixelExitBB->getFirstNonPHI());

llvm/test/CodeGen/AMDGPU/GlobalISel/atomic_optimizations_mul_one.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: opt -S -mtriple=amdgcn-- -amdgpu-atomic-optimizer -verify-machineinstrs %s | FileCheck -check-prefix=IR %s
+; RUN: llc -global-isel -mtriple=amdgcn-- -amdgpu-atomic-optimizations -verify-machineinstrs < %s | FileCheck -check-prefix=GCN %s
+
+; IR-NOT: mul i32 1, %[[REG:[0-9]+]]
+; IR-NOT: mul i32 %[[REG:[0-9]+]], 1
+
+declare i32 @llvm.amdgcn.struct.buffer.atomic.add.i32(i32, <4 x i32>, i32, i32, i32, i32 immarg)
+declare i32 @llvm.amdgcn.struct.buffer.atomic.sub.i32(i32, <4 x i32>, i32, i32, i32, i32 immarg)
+declare i32 @llvm.amdgcn.struct.buffer.atomic.xor.i32(i32, <4 x i32>, i32, i32, i32, i32 immarg)
+declare void @llvm.amdgcn.struct.buffer.store.format.v4i32(<4 x i32>, <4 x i32>, i32, i32, i32, i32 immarg)
+
+define amdgpu_cs void @atomic_add(<4 x i32> inreg %arg)  {
+; GCN-LABEL: atomic_add:
+; GCN:       ; %bb.0: ; %.entry
+; GCN-NEXT:    s_mov_b64 s[4:5], exec
+; GCN-NEXT:    v_mbcnt_lo_u32_b32_e64 v0, s4, 0
+; GCN-NEXT:    v_mbcnt_hi_u32_b32_e32 v0, s5, v0
+; GCN-NEXT:    v_cmp_eq_u32_e32 vcc, 0, v0
+; GCN-NEXT:    s_and_saveexec_b64 s[6:7], vcc
+; GCN-NEXT:    s_cbranch_execz BB0_2
+; GCN-NEXT:  ; %bb.1:
+; GCN-NEXT:    s_bcnt1_i32_b64 s4, s[4:5]
+; GCN-NEXT:    v_mov_b32_e32 v0, s4
+; GCN-NEXT:    v_mov_b32_e32 v1, 0
+; GCN-NEXT:    buffer_atomic_add v0, v1, s[0:3], 0 idxen glc
+; GCN-NEXT:  BB0_2:
+; GCN-NEXT:    s_endpgm
+.entry:
+  call i32 @llvm.amdgcn.struct.buffer.atomic.add.i32(i32 1, <4 x i32> %arg, i32 0, i32 0, i32 0, i32 0)
+  ret void
+}
+
+define amdgpu_cs void @atomic_add_and_format(<4 x i32> inreg %arg) {
+; GCN-LABEL: atomic_add_and_format:
+; GCN:       ; %bb.0: ; %.entry
+; GCN-NEXT:    s_mov_b64 s[6:7], exec
+; GCN-NEXT:    v_mbcnt_lo_u32_b32_e64 v0, s6, 0
+; GCN-NEXT:    v_mbcnt_hi_u32_b32_e32 v0, s7, v0
+; GCN-NEXT:    v_cmp_eq_u32_e32 vcc, 0, v0
+; GCN-NEXT:    ; implicit-def: $vgpr1
+; GCN-NEXT:    s_and_saveexec_b64 s[4:5], vcc
+; GCN-NEXT:    s_cbranch_execz BB1_2
+; GCN-NEXT:  ; %bb.1:
+; GCN-NEXT:    s_bcnt1_i32_b64 s6, s[6:7]
+; GCN-NEXT:    v_mov_b32_e32 v1, s6
+; GCN-NEXT:    v_mov_b32_e32 v2, 0
+; GCN-NEXT:    buffer_atomic_add v1, v2, s[0:3], 0 idxen glc
+; GCN-NEXT:  BB1_2:
+; GCN-NEXT:    s_or_b64 exec, exec, s[4:5]
+; GCN-NEXT:    s_waitcnt vmcnt(0)
+; GCN-NEXT:    v_readfirstlane_b32 s4, v1
+; GCN-NEXT:    v_add_i32_e32 v4, vcc, s4, v0
+; GCN-NEXT:    s_waitcnt expcnt(0)
+; GCN-NEXT:    v_mov_b32_e32 v0, s0
+; GCN-NEXT:    v_mov_b32_e32 v1, s1
+; GCN-NEXT:    v_mov_b32_e32 v2, s2
+; GCN-NEXT:    v_mov_b32_e32 v3, s3
+; GCN-NEXT:    buffer_store_format_xyzw v[0:3], v4, s[0:3], 0 idxen
+; GCN-NEXT:    s_endpgm
+.entry:
+  %a = call i32 @llvm.amdgcn.struct.buffer.atomic.add.i32(i32 1, <4 x i32> %arg, i32 0, i32 0, i32 0, i32 0)
+  call void @llvm.amdgcn.struct.buffer.store.format.v4i32(<4 x i32> %arg, <4 x i32> %arg, i32 %a, i32 0, i32 0, i32 0)
+  ret void
+}
+
+define amdgpu_cs void @atomic_sub(<4 x i32> inreg %arg)  {
+; GCN-LABEL: atomic_sub:
+; GCN:       ; %bb.0: ; %.entry
+; GCN-NEXT:    s_mov_b64 s[4:5], exec
+; GCN-NEXT:    v_mbcnt_lo_u32_b32_e64 v0, s4, 0
+; GCN-NEXT:    v_mbcnt_hi_u32_b32_e32 v0, s5, v0
+; GCN-NEXT:    v_cmp_eq_u32_e32 vcc, 0, v0
+; GCN-NEXT:    s_and_saveexec_b64 s[6:7], vcc
+; GCN-NEXT:    s_cbranch_execz BB2_2
+; GCN-NEXT:  ; %bb.1:
+; GCN-NEXT:    s_bcnt1_i32_b64 s4, s[4:5]
+; GCN-NEXT:    v_mov_b32_e32 v0, s4
+; GCN-NEXT:    v_mov_b32_e32 v1, 0
+; GCN-NEXT:    buffer_atomic_sub v0, v1, s[0:3], 0 idxen glc
+; GCN-NEXT:  BB2_2:
+; GCN-NEXT:    s_endpgm
+.entry:
+  call i32 @llvm.amdgcn.struct.buffer.atomic.sub.i32(i32 1, <4 x i32> %arg, i32 0, i32 0, i32 0, i32 0)
+  ret void
+}
+
+define amdgpu_cs void @atomic_sub_and_format(<4 x i32> inreg %arg) {
+; GCN-LABEL: atomic_sub_and_format:
+; GCN:       ; %bb.0: ; %.entry
+; GCN-NEXT:    s_mov_b64 s[6:7], exec
+; GCN-NEXT:    v_mbcnt_lo_u32_b32_e64 v0, s6, 0
+; GCN-NEXT:    v_mbcnt_hi_u32_b32_e32 v0, s7, v0
+; GCN-NEXT:    v_cmp_eq_u32_e32 vcc, 0, v0
+; GCN-NEXT:    ; implicit-def: $vgpr1
+; GCN-NEXT:    s_and_saveexec_b64 s[4:5], vcc
+; GCN-NEXT:    s_cbranch_execz BB3_2
+; GCN-NEXT:  ; %bb.1:
+; GCN-NEXT:    s_bcnt1_i32_b64 s6, s[6:7]
+; GCN-NEXT:    v_mov_b32_e32 v1, s6
+; GCN-NEXT:    v_mov_b32_e32 v2, 0
+; GCN-NEXT:    buffer_atomic_sub v1, v2, s[0:3], 0 idxen glc
+; GCN-NEXT:  BB3_2:
+; GCN-NEXT:    s_or_b64 exec, exec, s[4:5]
+; GCN-NEXT:    s_waitcnt vmcnt(0)
+; GCN-NEXT:    v_readfirstlane_b32 s4, v1
+; GCN-NEXT:    v_sub_i32_e32 v4, vcc, s4, v0
+; GCN-NEXT:    s_waitcnt expcnt(0)
+; GCN-NEXT:    v_mov_b32_e32 v0, s0
+; GCN-NEXT:    v_mov_b32_e32 v1, s1
+; GCN-NEXT:    v_mov_b32_e32 v2, s2
+; GCN-NEXT:    v_mov_b32_e32 v3, s3
+; GCN-NEXT:    buffer_store_format_xyzw v[0:3], v4, s[0:3], 0 idxen
+; GCN-NEXT:    s_endpgm
+.entry:
+  %a = call i32 @llvm.amdgcn.struct.buffer.atomic.sub.i32(i32 1, <4 x i32> %arg, i32 0, i32 0, i32 0, i32 0)
+  call void @llvm.amdgcn.struct.buffer.store.format.v4i32(<4 x i32> %arg, <4 x i32> %arg, i32 %a, i32 0, i32 0, i32 0)
+  ret void
+}
+
+define amdgpu_cs void @atomic_xor(<4 x i32> inreg %arg)  {
+; GCN-LABEL: atomic_xor:
+; GCN:       ; %bb.0: ; %.entry
+; GCN-NEXT:    s_mov_b64 s[4:5], exec
+; GCN-NEXT:    v_mbcnt_lo_u32_b32_e64 v0, s4, 0
+; GCN-NEXT:    v_mbcnt_hi_u32_b32_e32 v0, s5, v0
+; GCN-NEXT:    v_cmp_eq_u32_e32 vcc, 0, v0
+; GCN-NEXT:    s_and_saveexec_b64 s[6:7], vcc
+; GCN-NEXT:    s_cbranch_execz BB4_2
+; GCN-NEXT:  ; %bb.1:
+; GCN-NEXT:    s_bcnt1_i32_b64 s4, s[4:5]
+; GCN-NEXT:    s_and_b32 s4, s4, 1
+; GCN-NEXT:    v_mov_b32_e32 v0, s4
+; GCN-NEXT:    v_mov_b32_e32 v1, 0
+; GCN-NEXT:    buffer_atomic_xor v0, v1, s[0:3], 0 idxen glc
+; GCN-NEXT:  BB4_2:
+; GCN-NEXT:    s_endpgm
+.entry:
+  call i32 @llvm.amdgcn.struct.buffer.atomic.xor.i32(i32 1, <4 x i32> %arg, i32 0, i32 0, i32 0, i32 0)
+  ret void
+}
+
+define amdgpu_cs void @atomic_xor_and_format(<4 x i32> inreg %arg) {
+; GCN-LABEL: atomic_xor_and_format:
+; GCN:       ; %bb.0: ; %.entry
+; GCN-NEXT:    s_mov_b64 s[6:7], exec
+; GCN-NEXT:    v_mbcnt_lo_u32_b32_e64 v0, s6, 0
+; GCN-NEXT:    v_mbcnt_hi_u32_b32_e32 v0, s7, v0
+; GCN-NEXT:    v_cmp_eq_u32_e32 vcc, 0, v0
+; GCN-NEXT:    ; implicit-def: $vgpr1
+; GCN-NEXT:    s_and_saveexec_b64 s[4:5], vcc
+; GCN-NEXT:    s_cbranch_execz BB5_2
+; GCN-NEXT:  ; %bb.1:
+; GCN-NEXT:    s_bcnt1_i32_b64 s6, s[6:7]
+; GCN-NEXT:    s_and_b32 s6, s6, 1
+; GCN-NEXT:    v_mov_b32_e32 v1, s6
+; GCN-NEXT:    v_mov_b32_e32 v2, 0
+; GCN-NEXT:    buffer_atomic_xor v1, v2, s[0:3], 0 idxen glc
+; GCN-NEXT:  BB5_2:
+; GCN-NEXT:    s_or_b64 exec, exec, s[4:5]
+; GCN-NEXT:    s_waitcnt vmcnt(0)
+; GCN-NEXT:    v_readfirstlane_b32 s4, v1
+; GCN-NEXT:    v_and_b32_e32 v0, 1, v0
+; GCN-NEXT:    v_xor_b32_e32 v4, s4, v0
+; GCN-NEXT:    s_waitcnt expcnt(0)
+; GCN-NEXT:    v_mov_b32_e32 v0, s0
+; GCN-NEXT:    v_mov_b32_e32 v1, s1
+; GCN-NEXT:    v_mov_b32_e32 v2, s2
+; GCN-NEXT:    v_mov_b32_e32 v3, s3
+; GCN-NEXT:    buffer_store_format_xyzw v[0:3], v4, s[0:3], 0 idxen
+; GCN-NEXT:    s_endpgm
+.entry:
+  %a = call i32 @llvm.amdgcn.struct.buffer.atomic.xor.i32(i32 1, <4 x i32> %arg, i32 0, i32 0, i32 0, i32 0)
+  call void @llvm.amdgcn.struct.buffer.store.format.v4i32(<4 x i32> %arg, <4 x i32> %arg, i32 %a, i32 0, i32 0, i32 0)
+  ret void
+}