[AMDGPU] Propagate constants for llvm.amdgcn.wave.reduce.umin/umax

[pravinjagtap]

Formatting

[sameerds]

Should the change description say "propagate constants for llvm.amdgcn.wave.reduce.umin/umax"? That seems more informative.

[foad]

llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp

1139 ↗

(On Diff #543381)

This should go in lib/Analysis/ConstantFolding.cpp instead.

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

6197 ↗

(On Diff #543381)

Tests can go in test/Transforms/InstSimplify/ConstProp/AMDGPU/.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp

1139 ↗

(On Diff #543381)

that's true of many folds we do here, but I thought we were supposed to avoid putting target intrinsics there. We do lose out of the benefits of being a pure analysis here though

1144 ↗

(On Diff #543381)

isa<Constant> covers isa<Poison>

[pravinjagtap]

Moved const fold impl from /lib/Target/AMDGPU/AMDGPUInstCombineIntrinsic.cpp to
lib/Analysis/ConstantFolding.cpp

[pravinjagtap]

llvm/test/CodeGen/AMDGPU/llvm.amdgcn.reduce.umax.ll

138

Hello @foad, I am not sure why const fold optimization is not being applied here in llc.

[foad]

llvm/test/CodeGen/AMDGPU/llvm.amdgcn.reduce.umax.ll

138

llc does not run IR optimization passes. Use opt for that.

[pravinjagtap]

llvm/test/CodeGen/AMDGPU/llvm.amdgcn.reduce.umax.ll

138

If you notice, poison_value test point is being modified after early-cse pass during llc invocation in this test because of the current change in ConstantFolding.cpp. I was expecting same for const_value test point here.

[foad]

llvm/test/CodeGen/AMDGPU/llvm.amdgcn.reduce.umax.ll

138

Why do you think "const fold optimization is not being applied here"? Looking at the ISA, it just stores a constant value. What did you expect to see?

[pravinjagtap]

llvm/test/CodeGen/AMDGPU/llvm.amdgcn.reduce.umax.ll

138

Agree. Sorry, I overlooked. The optimization is applied in this case also.

[arsenm]

llvm/lib/Analysis/ConstantFolding.cpp

2851

It doesn't need to be a literal int isa<Constant> should cover every case

llvm/test/Transforms/InstSimplify/ConstProp/AMDGPU/wave.reduce.ll

34

Also try something exotic, like ptrtoint of a global declaration

[pravinjagtap]

Addressed review comments

[pravinjagtap]

llvm/test/Transforms/InstSimplify/ConstProp/AMDGPU/wave.reduce.ll

65

This usecase is not being optimized (const-fold).

[arsenm]

LGTM with the spurious include dropped

llvm/lib/Analysis/ConstantFolding.cpp

63

Don't add new include

llvm/test/Transforms/InstSimplify/ConstProp/AMDGPU/wave.reduce.ll

65

Don't see why that would be but it's unimportant

[pravinjagtap]

Addressed review comments

llvm/lib/Analysis/ConstantFolding.cpp

 #include "llvm/Support/KnownBits.h"
 #include "llvm/Support/MathExtras.h"
 #include <cassert>
 #include <cerrno>
 #include <cfenv>
 #include <cmath>
 #include <cstdint>
 
 using namespace llvm;
 
 namespace {
 
 //===----------------------------------------------------------------------===//
 // Constant Folding internal helper functions
 //===----------------------------------------------------------------------===//
 
 static Constant *foldConstVectorToAPInt(APInt &Result, Type *DestTy,
bool llvm::canConstantFoldCallTo(const CallBase *Call, const Function *F) {
   case Intrinsic::vector_reduce_or:
   case Intrinsic::vector_reduce_xor:
   case Intrinsic::vector_reduce_smin:
   case Intrinsic::vector_reduce_smax:
   case Intrinsic::vector_reduce_umin:
   case Intrinsic::vector_reduce_umax:
   // Target intrinsics
   case Intrinsic::amdgcn_perm:
+  case Intrinsic::amdgcn_wave_reduce_umin:
+  case Intrinsic::amdgcn_wave_reduce_umax:
   case Intrinsic::arm_mve_vctp8:
   case Intrinsic::arm_mve_vctp16:
   case Intrinsic::arm_mve_vctp32:
   case Intrinsic::arm_mve_vctp64:
   case Intrinsic::aarch64_sve_convert_from_svbool:
   // WebAssembly float semantics are always known
   case Intrinsic::wasm_trunc_signed:
   case Intrinsic::wasm_trunc_unsigned:
case Intrinsic::abs:
       if (C1->isOne() && (!C0 || C0->isMinSignedValue()))
         return UndefValue::get(Ty);
 
       // Undef operand with no poison min --> 0 (sign bit must be clear)
       if (!C0)
         return Constant::getNullValue(Ty);
 
       return ConstantInt::get(Ty, C0->abs());
+    case Intrinsic::amdgcn_wave_reduce_umin:
+    case Intrinsic::amdgcn_wave_reduce_umax:
+      return dyn_cast<Constant>(Operands[0]);
     }
 
     return nullptr;
   }
 
   // Support ConstantVector in case we have an Undef in the top.
   if ((isa<ConstantVector>(Operands[0]) ||
        isa<ConstantDataVector>(Operands[0])) &&
       // Check for default rounding mode.
       // FIXME: Support other rounding modes?
       isa<ConstantInt>(Operands[1]) &&
       cast<ConstantInt>(Operands[1])->getValue() == 4) {
     auto *Op = cast<Constant>(Operands[0]);

llvm/test/CodeGen/AMDGPU/llvm.amdgcn.reduce.umax.ll

 ; GFX1132GISEL-NEXT:    s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)
 ; GFX1132GISEL-NEXT:    s_endpgm
 entry:
     %result = call i32 @llvm.amdgcn.wave.reduce.umax.i32(i32 %in, i32 1)
     store i32 %result, ptr addrspace(1) %out
     ret void
 }
 
 define amdgpu_kernel void @const_value(ptr addrspace(1) %out) {
 ; GFX8DAGISEL-LABEL: const_value:
 ; GFX8DAGISEL:       ; %bb.0: ; %entry
 ; GFX8DAGISEL-NEXT:    s_load_dwordx2 s[0:1], s[0:1], 0x24
 ; GFX8DAGISEL-NEXT:    v_mov_b32_e32 v2, 0x7b
 ; GFX8DAGISEL-NEXT:    s_waitcnt lgkmcnt(0)
 ; GFX8DAGISEL-NEXT:    v_mov_b32_e32 v0, s0
 ; GFX8DAGISEL-NEXT:    v_mov_b32_e32 v1, s1
 ; GFX8DAGISEL-NEXT:    flat_store_dword v[0:1], v2
entry:
     %result = call i32 @llvm.amdgcn.wave.reduce.umax.i32(i32 123, i32 1)
     store i32 %result, ptr addrspace(1) %out
     ret void
 }
 
 define amdgpu_kernel void @poison_value(ptr addrspace(1) %out, i32 %in) {
 ; GFX8DAGISEL-LABEL: poison_value:
 ; GFX8DAGISEL:       ; %bb.0: ; %entry
-; GFX8DAGISEL-NEXT:    s_load_dwordx2 s[0:1], s[0:1], 0x24
-; GFX8DAGISEL-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX8DAGISEL-NEXT:    v_mov_b32_e32 v0, s0
-; GFX8DAGISEL-NEXT:    v_mov_b32_e32 v1, s1
-; GFX8DAGISEL-NEXT:    flat_store_dword v[0:1], v0
 ; GFX8DAGISEL-NEXT:    s_endpgm
 ;
 ; GFX8GISEL-LABEL: poison_value:
 ; GFX8GISEL:       ; %bb.0: ; %entry
-; GFX8GISEL-NEXT:    s_load_dwordx2 s[0:1], s[0:1], 0x24
-; GFX8GISEL-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX8GISEL-NEXT:    v_mov_b32_e32 v0, s0
-; GFX8GISEL-NEXT:    v_mov_b32_e32 v1, s1
-; GFX8GISEL-NEXT:    flat_store_dword v[0:1], v0
 ; GFX8GISEL-NEXT:    s_endpgm
 ;
 ; GFX9DAGISEL-LABEL: poison_value:
 ; GFX9DAGISEL:       ; %bb.0: ; %entry
-; GFX9DAGISEL-NEXT:    s_load_dwordx2 s[0:1], s[0:1], 0x24
-; GFX9DAGISEL-NEXT:    v_mov_b32_e32 v0, 0
-; GFX9DAGISEL-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX9DAGISEL-NEXT:    global_store_dword v0, v0, s[0:1]
 ; GFX9DAGISEL-NEXT:    s_endpgm
 ;
 ; GFX9GISEL-LABEL: poison_value:
 ; GFX9GISEL:       ; %bb.0: ; %entry
-; GFX9GISEL-NEXT:    s_load_dwordx2 s[0:1], s[0:1], 0x24
-; GFX9GISEL-NEXT:    v_mov_b32_e32 v0, 0
-; GFX9GISEL-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX9GISEL-NEXT:    global_store_dword v0, v0, s[0:1]
 ; GFX9GISEL-NEXT:    s_endpgm
 ;
 ; GFX10DAGISEL-LABEL: poison_value:
 ; GFX10DAGISEL:       ; %bb.0: ; %entry
-; GFX10DAGISEL-NEXT:    s_load_dwordx2 s[0:1], s[0:1], 0x24
-; GFX10DAGISEL-NEXT:    v_mov_b32_e32 v0, 0
-; GFX10DAGISEL-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX10DAGISEL-NEXT:    global_store_dword v0, v0, s[0:1]
 ; GFX10DAGISEL-NEXT:    s_endpgm
 ;
 ; GFX10GISEL-LABEL: poison_value:
 ; GFX10GISEL:       ; %bb.0: ; %entry
-; GFX10GISEL-NEXT:    s_load_dwordx2 s[0:1], s[0:1], 0x24
-; GFX10GISEL-NEXT:    v_mov_b32_e32 v0, 0
-; GFX10GISEL-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX10GISEL-NEXT:    global_store_dword v0, v0, s[0:1]
 ; GFX10GISEL-NEXT:    s_endpgm
 ;
 ; GFX11DAGISEL-LABEL: poison_value:
 ; GFX11DAGISEL:       ; %bb.0: ; %entry
-; GFX11DAGISEL-NEXT:    s_load_b64 s[0:1], s[0:1], 0x24
-; GFX11DAGISEL-NEXT:    v_mov_b32_e32 v0, 0
-; GFX11DAGISEL-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX11DAGISEL-NEXT:    global_store_b32 v0, v0, s[0:1]
-; GFX11DAGISEL-NEXT:    s_nop 0
-; GFX11DAGISEL-NEXT:    s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)
 ; GFX11DAGISEL-NEXT:    s_endpgm
 ;
 ; GFX11GISEL-LABEL: poison_value:
 ; GFX11GISEL:       ; %bb.0: ; %entry
-; GFX11GISEL-NEXT:    s_load_b64 s[0:1], s[0:1], 0x24
-; GFX11GISEL-NEXT:    v_mov_b32_e32 v0, 0
-; GFX11GISEL-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX11GISEL-NEXT:    global_store_b32 v0, v0, s[0:1]
-; GFX11GISEL-NEXT:    s_nop 0
-; GFX11GISEL-NEXT:    s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)
 ; GFX11GISEL-NEXT:    s_endpgm
 entry:
     %result = call i32 @llvm.amdgcn.wave.reduce.umax.i32(i32 poison, i32 1)
     store i32 %result, ptr addrspace(1) %out
     ret void
 }
 
 define amdgpu_kernel void @divergent_value(ptr addrspace(1) %out, i32 %in) {

llvm/test/CodeGen/AMDGPU/llvm.amdgcn.reduce.umin.ll

 ; GFX1132GISEL-NEXT:    s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)
 ; GFX1132GISEL-NEXT:    s_endpgm
 entry:
     %result = call i32 @llvm.amdgcn.wave.reduce.umin.i32(i32 123, i32 1)
     store i32 %result, ptr addrspace(1) %out
     ret void
 }
 
-define amdgpu_kernel void @poison_value(ptr addrspace(1) %out, i32 %in) {
+define amdgpu_kernel void @poison_value(ptr addrspace(1) %out) {
 ; GFX8DAGISEL-LABEL: poison_value:
 ; GFX8DAGISEL:       ; %bb.0: ; %entry
-; GFX8DAGISEL-NEXT:    s_load_dwordx2 s[0:1], s[0:1], 0x24
-; GFX8DAGISEL-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX8DAGISEL-NEXT:    v_mov_b32_e32 v0, s0
-; GFX8DAGISEL-NEXT:    v_mov_b32_e32 v1, s1
-; GFX8DAGISEL-NEXT:    flat_store_dword v[0:1], v0
 ; GFX8DAGISEL-NEXT:    s_endpgm
 ;
 ; GFX8GISEL-LABEL: poison_value:
 ; GFX8GISEL:       ; %bb.0: ; %entry
-; GFX8GISEL-NEXT:    s_load_dwordx2 s[0:1], s[0:1], 0x24
-; GFX8GISEL-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX8GISEL-NEXT:    v_mov_b32_e32 v0, s0
-; GFX8GISEL-NEXT:    v_mov_b32_e32 v1, s1
-; GFX8GISEL-NEXT:    flat_store_dword v[0:1], v0
 ; GFX8GISEL-NEXT:    s_endpgm
 ;
 ; GFX9DAGISEL-LABEL: poison_value:
 ; GFX9DAGISEL:       ; %bb.0: ; %entry
-; GFX9DAGISEL-NEXT:    s_load_dwordx2 s[0:1], s[0:1], 0x24
-; GFX9DAGISEL-NEXT:    v_mov_b32_e32 v0, 0
-; GFX9DAGISEL-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX9DAGISEL-NEXT:    global_store_dword v0, v0, s[0:1]
 ; GFX9DAGISEL-NEXT:    s_endpgm
 ;
 ; GFX9GISEL-LABEL: poison_value:
 ; GFX9GISEL:       ; %bb.0: ; %entry
-; GFX9GISEL-NEXT:    s_load_dwordx2 s[0:1], s[0:1], 0x24
-; GFX9GISEL-NEXT:    v_mov_b32_e32 v0, 0
-; GFX9GISEL-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX9GISEL-NEXT:    global_store_dword v0, v0, s[0:1]
 ; GFX9GISEL-NEXT:    s_endpgm
 ;
 ; GFX10DAGISEL-LABEL: poison_value:
 ; GFX10DAGISEL:       ; %bb.0: ; %entry
-; GFX10DAGISEL-NEXT:    s_load_dwordx2 s[0:1], s[0:1], 0x24
-; GFX10DAGISEL-NEXT:    v_mov_b32_e32 v0, 0
-; GFX10DAGISEL-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX10DAGISEL-NEXT:    global_store_dword v0, v0, s[0:1]
 ; GFX10DAGISEL-NEXT:    s_endpgm
 ;
 ; GFX10GISEL-LABEL: poison_value:
 ; GFX10GISEL:       ; %bb.0: ; %entry
-; GFX10GISEL-NEXT:    s_load_dwordx2 s[0:1], s[0:1], 0x24
-; GFX10GISEL-NEXT:    v_mov_b32_e32 v0, 0
-; GFX10GISEL-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX10GISEL-NEXT:    global_store_dword v0, v0, s[0:1]
 ; GFX10GISEL-NEXT:    s_endpgm
 ;
 ; GFX11DAGISEL-LABEL: poison_value:
 ; GFX11DAGISEL:       ; %bb.0: ; %entry
-; GFX11DAGISEL-NEXT:    s_load_b64 s[0:1], s[0:1], 0x24
-; GFX11DAGISEL-NEXT:    v_mov_b32_e32 v0, 0
-; GFX11DAGISEL-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX11DAGISEL-NEXT:    global_store_b32 v0, v0, s[0:1]
-; GFX11DAGISEL-NEXT:    s_nop 0
-; GFX11DAGISEL-NEXT:    s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)
 ; GFX11DAGISEL-NEXT:    s_endpgm
 ;
 ; GFX11GISEL-LABEL: poison_value:
 ; GFX11GISEL:       ; %bb.0: ; %entry
-; GFX11GISEL-NEXT:    s_load_b64 s[0:1], s[0:1], 0x24
-; GFX11GISEL-NEXT:    v_mov_b32_e32 v0, 0
-; GFX11GISEL-NEXT:    s_waitcnt lgkmcnt(0)
-; GFX11GISEL-NEXT:    global_store_b32 v0, v0, s[0:1]
-; GFX11GISEL-NEXT:    s_nop 0
-; GFX11GISEL-NEXT:    s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)
 ; GFX11GISEL-NEXT:    s_endpgm
 entry:
     %result = call i32 @llvm.amdgcn.wave.reduce.umin.i32(i32 poison, i32 1)
     store i32 %result, ptr addrspace(1) %out
     ret void
 }
 
-define amdgpu_kernel void @divergent_value(ptr addrspace(1) %out, i32 %in) {
+define amdgpu_kernel void @divergent_value(ptr addrspace(1) %out) {
 ; GFX8DAGISEL-LABEL: divergent_value:
 ; GFX8DAGISEL:       ; %bb.0: ; %entry
 ; GFX8DAGISEL-NEXT:    s_load_dwordx2 s[0:1], s[0:1], 0x24
 ; GFX8DAGISEL-NEXT:    s_mov_b64 s[2:3], exec
 ; GFX8DAGISEL-NEXT:    s_mov_b32 s4, -1
 ; GFX8DAGISEL-NEXT:  .LBB3_1: ; =>This Inner Loop Header: Depth=1
 ; GFX8DAGISEL-NEXT:    s_ff1_i32_b64 s5, s[2:3]
 ; GFX8DAGISEL-NEXT:    v_readlane_b32 s6, v0, s5

llvm/test/Transforms/InstSimplify/ConstProp/AMDGPU/wave.reduce.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -passes=instsimplify -S | FileCheck %s
+
+; --------------------------------------------------------------------
+; llvm.amdgcn.wave.reduce.umin.i32
+; --------------------------------------------------------------------
+
+declare i32 @llvm.amdgcn.wave.reduce.umin.i32(i32, i32 immarg)
+
+define amdgpu_kernel void @test_constant_fold_wave_reduce_umin_poison(ptr addrspace(1) %out, i32 %in) {
+; CHECK-LABEL: @test_constant_fold_wave_reduce_umin_poison(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    store i32 poison, ptr addrspace(1) [[OUT:%.*]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %result = call i32 @llvm.amdgcn.wave.reduce.umin.i32(i32 poison, i32 1)
+  store i32 %result, ptr addrspace(1) %out
+  ret void
+}
+
+define amdgpu_kernel void @test_constant_fold_wave_reduce_umin_const(ptr addrspace(1) %out) {
+; CHECK-LABEL: @test_constant_fold_wave_reduce_umin_const(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    store i32 123, ptr addrspace(1) [[OUT:%.*]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %result = call i32 @llvm.amdgcn.wave.reduce.umin.i32(i32 123, i32 1)
+  store i32 %result, ptr addrspace(1) %out
+  ret void
+}
+
+; --------------------------------------------------------------------
+; llvm.amdgcn.wave.reduce.umin.i32
+; --------------------------------------------------------------------
+
+declare i32 @llvm.amdgcn.wave.reduce.umax.i32(i32, i32 immarg)
+
+define amdgpu_kernel void @test_constant_fold_wave_reduce_umax_poison(ptr addrspace(1) %out, i32 %in) {
+; CHECK-LABEL: @test_constant_fold_wave_reduce_umax_poison(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    store i32 poison, ptr addrspace(1) [[OUT:%.*]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %result = call i32 @llvm.amdgcn.wave.reduce.umax.i32(i32 poison, i32 1)
+  store i32 %result, ptr addrspace(1) %out
+  ret void
+}
+
+define amdgpu_kernel void @test_constant_fold_wave_reduce_umax_const(ptr addrspace(1) %out) {
+; CHECK-LABEL: @test_constant_fold_wave_reduce_umax_const(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    store i32 123, ptr addrspace(1) [[OUT:%.*]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %result = call i32 @llvm.amdgcn.wave.reduce.umax.i32(i32 123, i32 1)
+  store i32 %result, ptr addrspace(1) %out
+  ret void
+}
+
+@gv = constant i32 0
+define amdgpu_kernel void @test_constant_fold_wave_reduce_umax_gv(ptr addrspace(1) %out) {
+; CHECK-LABEL: @test_constant_fold_wave_reduce_umax_gv(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[RESULT:%.*]] = call i32 @llvm.amdgcn.wave.reduce.umax.i32(i32 ptrtoint (ptr @gv to i32), i32 1)
+; CHECK-NEXT:    store i32 [[RESULT]], ptr addrspace(1) [[OUT:%.*]], align 4
+; CHECK-NEXT:    ret void
+;
+entry:
+  %result = call i32 @llvm.amdgcn.wave.reduce.umax.i32(i32 ptrtoint (ptr @gv to i32), i32 1)
+  store i32 %result, ptr addrspace(1) %out
+  ret void
+}