[llvm] [SPIRV] Sign-extend operands of sign-sensitive ops on sub-pow2 widths (PR #203661)

[Juan Manuel Martinez Caamaño via llvm-commits]

================
@@ -466,6 +467,121 @@ void processInstr(MachineInstr &MI, MachineIRBuilder &MIB,
 }
 } // namespace llvm
 
+// Sign-sensitive integer ops: their result depends on the value of the input
+// sign bit at position (width-1). On sub-pow2 widths the general widening
+// loop is a pure LLT relabel, which leaves the sign bit at the *original*
+// position instead of the widened MSB. These ops therefore need an explicit
+// G_SEXT_INREG on each value operand to move the sign bit up.
+//
+// Signed-vs-unsigned G_ICMP is distinguished by its predicate operand.
+//
+// TODO: follow-up PRs will add the remaining sign-sensitive opcodes
+// (e.g. G_SMIN/G_SMAX, G_SADDSAT/G_SSUBSAT, signed overflow ops).
+static bool isSignSensitiveOp(const MachineInstr &MI) {
+  switch (MI.getOpcode()) {
+  case TargetOpcode::G_ASHR:
+  case TargetOpcode::G_SDIV:
+  case TargetOpcode::G_SREM:
+    return true;
+  case TargetOpcode::G_ICMP:
+    return CmpInst::isSigned(
+        static_cast<CmpInst::Predicate>(MI.getOperand(1).getPredicate()));
+  default:
+    return false;
+  }
+}
+
+// Record each sign-sensitive op's scalar value-operand widths before any
+// widening runs. Once later passes retype those vregs to their pow2 LLTs, the
+// original narrow width is gone; widenSignSensitiveOps needs it to know how
+// many low bits to sign-extend.
+static DenseMap<Register, unsigned>
+recordSignSensitiveOperandWidths(MachineFunction &MF,
+                                 MachineRegisterInfo &MRI) {
+  DenseMap<Register, unsigned> OrigWidth;
+  for (MachineBasicBlock &MBB : MF) {
+    for (MachineInstr &MI : MBB) {
+      if (!isSignSensitiveOp(MI))
+        continue;
+      // Value operands are the trailing two, past any def or predicate.
+      unsigned N = MI.getNumOperands();
+      for (unsigned I : {N - 2, N - 1}) {
+        const MachineOperand &MOP = MI.getOperand(I);
+        if (!MOP.isReg())
+          continue;
+        Register Reg = MOP.getReg();
+        if (!MRI.getType(Reg).isScalar())
+          continue;
+        OrigWidth.try_emplace(Reg, MRI.getType(Reg).getScalarSizeInBits());
+      }
+    }
+  }
+  return OrigWidth;
+}
+
+// For every sign-sensitive op, insert G_SEXT_INREG on each value operand
+// whose original width (from OrigWidth, captured before any retyping) is
+// narrower than the widened pow2 width. The rewritten operand's vreg LLT
+// is retyped in place to the widened width.
+//
+// OrigWidth must have been populated by recordSignSensitiveOperandWidths
+// before other passes retyped the vregs; otherwise the narrow widths
+// needed here are lost.
+//
+// TODO: handle vector operands.
+static void
+widenSignSensitiveOps(MachineFunction &MF, SPIRVGlobalRegistry *GR,
+                      MachineIRBuilder &MIB, MachineRegisterInfo &MRI,
+                      const DenseMap<Register, unsigned> &OrigWidth) {
+  // Emit G_SEXT_INREG from Reg's recorded narrow width; retypes Reg to the
+  // widened width. Returns an invalid Register if no widening is needed.
+  auto GetSignExtendedReg = [&](Register Reg, MachineInstr &MI) -> Register {
+    auto OWIt = OrigWidth.find(Reg);
+    assert(OWIt != OrigWidth.end() &&
+           "Sign-sensitive operand width was not recorded");
+    unsigned OldW = OWIt->second;
+    unsigned NewW = widenBitWidthToNextPow2(OldW);
+    if (NewW == OldW)
+      return Register();
+    LLT NewLLT = LLT::scalar(NewW);
+    SPIRVTypeInst SpvTy = GR->getOrCreateSPIRVIntegerType(NewW, MIB);
+    Register SExted = MRI.createGenericVirtualRegister(NewLLT);
+    GR->assignSPIRVTypeToVReg(SpvTy, SExted, MF);
+    MRI.setRegClass(SExted, GR->getRegClass(SpvTy));
+    MRI.setType(Reg, NewLLT);
+    MIB.setInsertPt(*MI.getParent(), MI.getIterator());
+    MIB.buildSExtInReg(SExted, Reg, OldW);
+    return SExted;
+  };
+
+  for (MachineBasicBlock &MBB : MF) {
+    for (MachineInstr &MI : llvm::make_early_inc_range(MBB)) {
+      if (!isSignSensitiveOp(MI))
+        continue;
----------------
jmmartinez wrote:

What if `recordSignSensitiveOperandWidths` instead of returning simply a DenseMap / returned a class that also contained the sign sensitive instructions?

Then we would not need to scan the whole Machine function all over again.

https://github.com/llvm/llvm-project/pull/203661