[SelectionDAG] Allow constant UREM decomposition without high multiply

[keshavvinayak01]

SelectionDAG may introduce a wide constant UREM when promoting funnel shifts on irregular integer types. On AMDGPU, an i65 fshl becomes an i128 remainder by 65; the existing decomposition is rejected because i64 MULHU and UMUL_LOHI are unavailable, after which legalization attempts the unsupported __umodti3 libcall.

Allow the existing decomposition when the wide UREM libcall is unavailable and half-width UDIVREM is legal or custom. This lets AMDGPU finish through its custom i64 UDIVREM lowering while preserving the existing libcall preference for targets with a usable wide libcall.
Fixes #197949

[keshavvinayak01]

GlobalISel support is left as follow-up work because irregular-width G_FSHL requires separate lowering and instruction-selection support.

[llvmorg-github-actions]

@llvm/pr-subscribers-llvm-selectiondag

@llvm/pr-subscribers-backend-amdgpu

Author: Keshav Vinayak Jha (keshavvinayak01)

Changes

AMDGPU SelectionDAG can produce an i128 UREM while expanding irregular-width funnel shifts such as i65. Generic legalization declines the constant decomposition and attempts the unsupported __umodti3 libcall.

Custom-lower constant i128 UREM in SITargetLowering using the existing half-width remainder decomposition. This covers types promoted to i128 while leaving generic SelectionDAG behavior and other targets unchanged.

Fixes #197949

Assisted-by: Codex

---

Full diff: https://github.com/llvm/llvm-project/pull/210232.diff

4 Files Affected:

• (modified) llvm/include/llvm/CodeGen/TargetLowering.h (+4-4)
• (modified) llvm/lib/Target/AMDGPU/SIISelLowering.cpp (+23)
• (added) llvm/test/CodeGen/AMDGPU/fshl-illegal-types.ll (+47)
• (added) llvm/test/CodeGen/AMDGPU/urem-constant-i128.ll (+42)

diff --git a/llvm/include/llvm/CodeGen/TargetLowering.h b/llvm/include/llvm/CodeGen/TargetLowering.h
index 75b1ca2d2745f..882999a9634f1 100644
--- a/llvm/include/llvm/CodeGen/TargetLowering.h
+++ b/llvm/include/llvm/CodeGen/TargetLowering.h
@@ -6032,6 +6032,10 @@ class LLVM_ABI TargetLowering : public TargetLoweringBase {
                                        SelectionDAG &DAG) const;
 
 protected:
+  bool expandUDIVREMByConstantViaUREMDecomposition(
+      SDNode *N, APInt Divisor, SmallVectorImpl<SDValue> &Result, EVT HiLoVT,
+      SelectionDAG &DAG, SDValue LL, SDValue LH) const;
+
   void setTypeIdForCallsiteInfo(const CallBase *CB, MachineFunction &MF,
                                 MachineFunction::CallSiteInfo &CSInfo) const;
 
@@ -6069,10 +6073,6 @@ class LLVM_ABI TargetLowering : public TargetLoweringBase {
                           ISD::CondCode Cond, DAGCombinerInfo &DCI,
                           const SDLoc &DL) const;
 
-  bool expandUDIVREMByConstantViaUREMDecomposition(
-      SDNode *N, APInt Divisor, SmallVectorImpl<SDValue> &Result, EVT HiLoVT,
-      SelectionDAG &DAG, SDValue LL, SDValue LH) const;
-
   bool expandUDIVREMByConstantViaUMulHiMagic(SDNode *N, const APInt &Divisor,
                                              SmallVectorImpl<SDValue> &Result,
                                              EVT HiLoVT, SelectionDAG &DAG,
diff --git a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
index 8a8a3b97e8e08..01dd4b8fecd9c 100644
--- a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
+++ b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
@@ -951,6 +951,10 @@ SITargetLowering::SITargetLowering(const TargetMachine &TM,
 
   setOperationAction({ISD::SMULO, ISD::UMULO}, MVT::i64, Custom);
 
+  // Try to decompose constant i128 remainders before falling back to an
+  // unsupported i128 libcall.
+  setOperationAction(ISD::UREM, MVT::i128, Custom);
+
   if (Subtarget->hasVMulU64Inst())
     setOperationAction(ISD::MUL, MVT::i64, Legal);
   else if (Subtarget->hasScalarSMulU64())
@@ -8229,6 +8233,25 @@ void SITargetLowering::ReplaceNodeResults(SDNode *N,
                                           SmallVectorImpl<SDValue> &Results,
                                           SelectionDAG &DAG) const {
   switch (N->getOpcode()) {
+  case ISD::UREM: {
+    if (N->getValueType(0) != MVT::i128)
+      break;
+
+    auto *Divisor = dyn_cast<ConstantSDNode>(N->getOperand(1));
+    if (!Divisor || Divisor->getAPIntValue().ule(1))
+      break;
+
+    SmallVector<SDValue, 2> Rem;
+    if (!expandUDIVREMByConstantViaUREMDecomposition(
+            N, Divisor->getAPIntValue(), Rem, MVT::i64, DAG, SDValue(),
+            SDValue()))
+      break;
+
+    assert(Rem.size() == 2 && "Expected low and high remainder parts");
+    Results.push_back(
+        DAG.getNode(ISD::BUILD_PAIR, SDLoc(N), MVT::i128, Rem[0], Rem[1]));
+    return;
+  }
   case ISD::INSERT_VECTOR_ELT: {
     if (SDValue Res = lowerINSERT_VECTOR_ELT(SDValue(N, 0), DAG))
       Results.push_back(Res);
diff --git a/llvm/test/CodeGen/AMDGPU/fshl-illegal-types.ll b/llvm/test/CodeGen/AMDGPU/fshl-illegal-types.ll
new file mode 100644
index 0000000000000..3ef063fcd7ee4
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/fshl-illegal-types.ll
@@ -0,0 +1,47 @@
+; RUN: llc -global-isel=0 -mtriple=amdgcn -mcpu=gfx600 -verify-machineinstrs < %s | FileCheck %s --implicit-check-not=__umodti3
+; RUN: llc -global-isel=0 -mtriple=amdgcn -mcpu=gfx1100 -verify-machineinstrs < %s | FileCheck %s --implicit-check-not=__umodti3
+
+; The generated remainder and funnel-shift expansions are long and subject to
+; frequent changes. Check that SelectionDAG can lower variable funnel shifts
+; whose irregular types promote to i128 without requiring an i128 urem libcall.
+
+declare i65 @llvm.fshl.i65(i65, i65, i65)
+declare i66 @llvm.fshl.i66(i66, i66, i66)
+declare i67 @llvm.fshl.i67(i67, i67, i67)
+declare i68 @llvm.fshl.i68(i68, i68, i68)
+declare i127 @llvm.fshl.i127(i127, i127, i127)
+
+define i65 @fshl_i65(i65 %amt) {
+; CHECK-LABEL: fshl_i65:
+; CHECK:       s_setpc_b64
+  %result = call i65 @llvm.fshl.i65(i65 1, i65 0, i65 %amt)
+  ret i65 %result
+}
+
+define i66 @fshl_i66(i66 %amt) {
+; CHECK-LABEL: fshl_i66:
+; CHECK:       s_setpc_b64
+  %result = call i66 @llvm.fshl.i66(i66 1, i66 0, i66 %amt)
+  ret i66 %result
+}
+
+define i67 @fshl_i67(i67 %amt) {
+; CHECK-LABEL: fshl_i67:
+; CHECK:       s_setpc_b64
+  %result = call i67 @llvm.fshl.i67(i67 1, i67 0, i67 %amt)
+  ret i67 %result
+}
+
+define i68 @fshl_i68(i68 %amt) {
+; CHECK-LABEL: fshl_i68:
+; CHECK:       s_setpc_b64
+  %result = call i68 @llvm.fshl.i68(i68 1, i68 0, i68 %amt)
+  ret i68 %result
+}
+
+define i127 @fshl_i127(i127 %amt) {
+; CHECK-LABEL: fshl_i127:
+; CHECK:       s_setpc_b64
+  %result = call i127 @llvm.fshl.i127(i127 1, i127 0, i127 %amt)
+  ret i127 %result
+}
diff --git a/llvm/test/CodeGen/AMDGPU/urem-constant-i128.ll b/llvm/test/CodeGen/AMDGPU/urem-constant-i128.ll
new file mode 100644
index 0000000000000..eab4234c6758b
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/urem-constant-i128.ll
@@ -0,0 +1,42 @@
+; RUN: llc -global-isel=0 -mtriple=amdgcn -mcpu=gfx600 -verify-machineinstrs < %s | FileCheck %s --implicit-check-not=__umodti3
+; RUN: llc -global-isel=0 -mtriple=amdgcn -mcpu=gfx1100 -verify-machineinstrs < %s | FileCheck %s --implicit-check-not=__umodti3
+
+; AMDGPU has custom i64 udivrem lowering, but no legal i64 high multiply and no
+; i128 urem libcall. Check that wide constant urem uses the existing half-width
+; urem decomposition for representative odd, even, alternate, full-half, and
+; upper-bound funnel-shift widths.
+
+define i128 @urem_i128_65(i128 %src) {
+; CHECK-LABEL: urem_i128_65:
+; CHECK:       s_setpc_b64
+  %result = urem i128 %src, 65
+  ret i128 %result
+}
+
+define i128 @urem_i128_66(i128 %src) {
+; CHECK-LABEL: urem_i128_66:
+; CHECK:       s_setpc_b64
+  %result = urem i128 %src, 66
+  ret i128 %result
+}
+
+define i128 @urem_i128_67(i128 %src) {
+; CHECK-LABEL: urem_i128_67:
+; CHECK:       s_setpc_b64
+  %result = urem i128 %src, 67
+  ret i128 %result
+}
+
+define i128 @urem_i128_68(i128 %src) {
+; CHECK-LABEL: urem_i128_68:
+; CHECK:       s_setpc_b64
+  %result = urem i128 %src, 68
+  ret i128 %result
+}
+
+define i128 @urem_i128_127(i128 %src) {
+; CHECK-LABEL: urem_i128_127:
+; CHECK:       s_setpc_b64
+  %result = urem i128 %src, 127
+  ret i128 %result
+}

[arsenm]

This should not require custom lowering. The legalizer code doing the wrong thing should be fixed

[keshavvinayak01]

@arsenm Fixed the legalization logic for UREM; please re-review!