[CodeGen] Add preliminary plumbing for samesign flag

antoniofrighetto · antoniofrighetto · commit 7c49142ae3cf · 2024-10-30T14:59:42.000+01:00
Extend recently-added poison-generating IR flag to codegen as well.
diff --git a/llvm/include/llvm/CodeGen/GlobalISel/GenericMachineInstrs.h b/llvm/include/llvm/CodeGen/GlobalISel/GenericMachineInstrs.h
@@ -28,7 +28,7 @@ namespace llvm {
 class GenericMachineInstr : public MachineInstr {
   constexpr static unsigned PoisonFlags = NoUWrap | NoSWrap | NoUSWrap |
                                           IsExact | Disjoint | NonNeg |
-                                          FmNoNans | FmNoInfs;
+                                          FmNoNans | FmNoInfs | SameSign;
 
 public:
   GenericMachineInstr() = delete;
diff --git a/llvm/include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h b/llvm/include/llvm/CodeGen/GlobalISel/MachineIRBuilder.h
@@ -1255,7 +1255,8 @@ class MachineIRBuilder {
   ///
   /// \return a MachineInstrBuilder for the newly created instruction.
   MachineInstrBuilder buildICmp(CmpInst::Predicate Pred, const DstOp &Res,
-                                const SrcOp &Op0, const SrcOp &Op1);
+                                const SrcOp &Op0, const SrcOp &Op1,
+                                std::optional<unsigned> Flags = std::nullopt);
 
   /// Build and insert a \p Res = G_FCMP \p Pred\p Op0, \p Op1
   ///
diff --git a/llvm/include/llvm/CodeGen/MachineInstr.h b/llvm/include/llvm/CodeGen/MachineInstr.h
@@ -119,6 +119,7 @@ class MachineInstr
     Disjoint = 1 << 19,      // Each bit is zero in at least one of the inputs.
     NoUSWrap = 1 << 20,      // Instruction supports geps
                              // no unsigned signed wrap.
+    SameSign = 1 << 21,      // Compare instruction has same sign.
   };
 
 private:
diff --git a/llvm/include/llvm/CodeGen/SelectionDAGNodes.h b/llvm/include/llvm/CodeGen/SelectionDAGNodes.h
@@ -410,12 +410,14 @@ struct SDNodeFlags {
     NoFPExcept = 1 << 12,
     // Instructions with attached 'unpredictable' metadata on IR level.
     Unpredictable = 1 << 13,
+    // Compare instructions which may carry the samesign flag.
+    SameSign = 1 << 14,
 
     // NOTE: Please update LargestValue in LLVM_DECLARE_ENUM_AS_BITMASK below
     // the class definition when adding new flags.
 
     PoisonGeneratingFlags = NoUnsignedWrap | NoSignedWrap | Exact | Disjoint |
-                            NonNeg | NoNaNs | NoInfs,
+                            NonNeg | NoNaNs | NoInfs | SameSign,
   };
 
   /// Default constructor turns off all optimization flags.
@@ -437,6 +439,7 @@ struct SDNodeFlags {
   void setNoSignedWrap(bool b) { setFlag<NoSignedWrap>(b); }
   void setExact(bool b) { setFlag<Exact>(b); }
   void setDisjoint(bool b) { setFlag<Disjoint>(b); }
+  void setSameSign(bool b) { setFlag<SameSign>(b); }
   void setNonNeg(bool b) { setFlag<NonNeg>(b); }
   void setNoNaNs(bool b) { setFlag<NoNaNs>(b); }
   void setNoInfs(bool b) { setFlag<NoInfs>(b); }
@@ -453,6 +456,7 @@ struct SDNodeFlags {
   bool hasNoSignedWrap() const { return Flags & NoSignedWrap; }
   bool hasExact() const { return Flags & Exact; }
   bool hasDisjoint() const { return Flags & Disjoint; }
+  bool hasSameSign() const { return Flags & SameSign; }
   bool hasNonNeg() const { return Flags & NonNeg; }
   bool hasNoNaNs() const { return Flags & NoNaNs; }
   bool hasNoInfs() const { return Flags & NoInfs; }
@@ -474,7 +478,7 @@ struct SDNodeFlags {
 };
 
 LLVM_DECLARE_ENUM_AS_BITMASK(decltype(SDNodeFlags::None),
-                             SDNodeFlags::Unpredictable);
+                             SDNodeFlags::SameSign);
 
 /// Represents one node in the SelectionDAG.
 ///
diff --git a/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp b/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
@@ -340,18 +340,18 @@ bool IRTranslator::translateCompare(const User &U,
   Register Op1 = getOrCreateVReg(*U.getOperand(1));
   Register Res = getOrCreateVReg(U);
   CmpInst::Predicate Pred = CI->getPredicate();
+  uint32_t Flags = 0;
+  if (CI)
+    Flags = MachineInstr::copyFlagsFromInstruction(*CI);
   if (CmpInst::isIntPredicate(Pred))
-    MIRBuilder.buildICmp(Pred, Res, Op0, Op1);
+    MIRBuilder.buildICmp(Pred, Res, Op0, Op1, Flags);
   else if (Pred == CmpInst::FCMP_FALSE)
     MIRBuilder.buildCopy(
         Res, getOrCreateVReg(*Constant::getNullValue(U.getType())));
   else if (Pred == CmpInst::FCMP_TRUE)
     MIRBuilder.buildCopy(
         Res, getOrCreateVReg(*Constant::getAllOnesValue(U.getType())));
   else {
-    uint32_t Flags = 0;
-    if (CI)
-      Flags = MachineInstr::copyFlagsFromInstruction(*CI);
     MIRBuilder.buildFCmp(Pred, Res, Op0, Op1, Flags);
   }
 
diff --git a/llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp b/llvm/lib/CodeGen/GlobalISel/MachineIRBuilder.cpp
@@ -898,8 +898,9 @@ MachineIRBuilder::buildFPTrunc(const DstOp &Res, const SrcOp &Op,
 MachineInstrBuilder MachineIRBuilder::buildICmp(CmpInst::Predicate Pred,
                                                 const DstOp &Res,
                                                 const SrcOp &Op0,
-                                                const SrcOp &Op1) {
-  return buildInstr(TargetOpcode::G_ICMP, Res, {Pred, Op0, Op1});
+                                                const SrcOp &Op1,
+                                                std::optional<unsigned> Flags) {
+  return buildInstr(TargetOpcode::G_ICMP, Res, {Pred, Op0, Op1}, Flags);
 }
 
 MachineInstrBuilder MachineIRBuilder::buildFCmp(CmpInst::Predicate Pred,
diff --git a/llvm/lib/CodeGen/MIRParser/MILexer.cpp b/llvm/lib/CodeGen/MIRParser/MILexer.cpp
@@ -216,6 +216,7 @@ static MIToken::TokenKind getIdentifierKind(StringRef Identifier) {
       .Case("exact", MIToken::kw_exact)
       .Case("nneg", MIToken::kw_nneg)
       .Case("disjoint", MIToken::kw_disjoint)
+      .Case("samesign", MIToken::kw_samesign)
       .Case("nofpexcept", MIToken::kw_nofpexcept)
       .Case("unpredictable", MIToken::kw_unpredictable)
       .Case("debug-location", MIToken::kw_debug_location)
diff --git a/llvm/lib/CodeGen/MIRParser/MILexer.h b/llvm/lib/CodeGen/MIRParser/MILexer.h
@@ -77,6 +77,7 @@ struct MIToken {
     kw_unpredictable,
     kw_nneg,
     kw_disjoint,
+    kw_samesign,
     kw_debug_location,
     kw_debug_instr_number,
     kw_dbg_instr_ref,
diff --git a/llvm/lib/CodeGen/MIRParser/MIParser.cpp b/llvm/lib/CodeGen/MIRParser/MIParser.cpp
@@ -1476,7 +1476,8 @@ bool MIParser::parseInstruction(unsigned &OpCode, unsigned &Flags) {
          Token.is(MIToken::kw_noconvergent) ||
          Token.is(MIToken::kw_unpredictable) ||
          Token.is(MIToken::kw_nneg) ||
-         Token.is(MIToken::kw_disjoint)) {
+         Token.is(MIToken::kw_disjoint) ||
+         Token.is(MIToken::kw_samesign)) {
     // clang-format on
     // Mine frame and fast math flags
     if (Token.is(MIToken::kw_frame_setup))
@@ -1513,6 +1514,8 @@ bool MIParser::parseInstruction(unsigned &OpCode, unsigned &Flags) {
       Flags |= MachineInstr::NonNeg;
     if (Token.is(MIToken::kw_disjoint))
       Flags |= MachineInstr::Disjoint;
+    if (Token.is(MIToken::kw_samesign))
+      Flags |= MachineInstr::SameSign;
 
     lex();
   }
diff --git a/llvm/lib/CodeGen/MIRPrinter.cpp b/llvm/lib/CodeGen/MIRPrinter.cpp
@@ -837,6 +837,8 @@ void MIPrinter::print(const MachineInstr &MI) {
     OS << "disjoint ";
   if (MI.getFlag(MachineInstr::NoUSWrap))
     OS << "nusw ";
+  if (MI.getFlag(MachineInstr::SameSign))
+    OS << "samesign ";
 
   OS << TII->getName(MI.getOpcode());
   if (I < E)
diff --git a/llvm/lib/CodeGen/MachineInstr.cpp b/llvm/lib/CodeGen/MachineInstr.cpp
@@ -620,6 +620,12 @@ uint32_t MachineInstr::copyFlagsFromInstruction(const Instruction &I) {
       MIFlags |= MachineInstr::MIFlag::FmReassoc;
   }
 
+  // Copy the samesign flag.
+  if (const auto *ICI = dyn_cast<ICmpInst>(&I)) {
+    if (ICI->hasSameSign())
+      MIFlags |= MachineInstr::MIFlag::SameSign;
+  }
+
   if (I.getMetadata(LLVMContext::MD_unpredictable))
     MIFlags |= MachineInstr::MIFlag::Unpredictable;
 
@@ -1770,6 +1776,8 @@ void MachineInstr::print(raw_ostream &OS, ModuleSlotTracker &MST,
     OS << "nneg ";
   if (getFlag(MachineInstr::Disjoint))
     OS << "disjoint ";
+  if (getFlag(MachineInstr::SameSign))
+    OS << "samesign ";
 
   // Print the opcode name.
   if (TII)
diff --git a/llvm/lib/CodeGen/SelectionDAG/InstrEmitter.cpp b/llvm/lib/CodeGen/SelectionDAG/InstrEmitter.cpp
@@ -1105,6 +1105,9 @@ EmitMachineNode(SDNode *Node, bool IsClone, bool IsCloned,
 
     if (Flags.hasDisjoint())
       MI->setFlag(MachineInstr::MIFlag::Disjoint);
+
+    if (Flags.hasSameSign())
+      MI->setFlag(MachineInstr::MIFlag::SameSign);
   }
 
   // Emit all of the actual operands of this instruction, adding them to the
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -3652,6 +3652,10 @@ void SelectionDAGBuilder::visitICmp(const ICmpInst &I) {
     Op2 = DAG.getPtrExtOrTrunc(Op2, getCurSDLoc(), MemVT);
   }
 
+  SDNodeFlags Flags;
+  Flags.setSameSign(I.hasSameSign());
+  SelectionDAG::FlagInserter FlagsInserter(DAG, Flags);
+
   EVT DestVT = DAG.getTargetLoweringInfo().getValueType(DAG.getDataLayout(),
                                                         I.getType());
   setValue(&I, DAG.getSetCC(getCurSDLoc(), DestVT, Op1, Op2, Opcode));
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGDumper.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGDumper.cpp
@@ -653,6 +653,9 @@ void SDNode::print_details(raw_ostream &OS, const SelectionDAG *G) const {
   if (getFlags().hasDisjoint())
     OS << " disjoint";
 
+  if (getFlags().hasSameSign())
+    OS << " samesign";
+
   if (getFlags().hasNonNeg())
     OS << " nneg";
 
diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/irtranslator-icmp-samesign.ll b/llvm/test/CodeGen/AArch64/GlobalISel/irtranslator-icmp-samesign.ll
@@ -0,0 +1,33 @@
+; NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 5
+; RUN: llc -mtriple=aarch64-linux-gnu -O0 -stop-after=irtranslator -global-isel -verify-machineinstrs %s -o - 2>&1 | FileCheck %s
+
+define i1 @icmp_samesign(i32 %a, i32 %b) {
+  ; CHECK-LABEL: name: icmp_samesign
+  ; CHECK: bb.1 (%ir-block.0):
+  ; CHECK-NEXT:   liveins: $w0, $w1
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:_(s32) = COPY $w0
+  ; CHECK-NEXT:   [[COPY1:%[0-9]+]]:_(s32) = COPY $w1
+  ; CHECK-NEXT:   %2:_(s1) = samesign G_ICMP intpred(ult), [[COPY]](s32), [[COPY1]]
+  ; CHECK-NEXT:   [[ZEXT:%[0-9]+]]:_(s8) = G_ZEXT %2(s1)
+  ; CHECK-NEXT:   [[ANYEXT:%[0-9]+]]:_(s32) = G_ANYEXT [[ZEXT]](s8)
+  ; CHECK-NEXT:   $w0 = COPY [[ANYEXT]](s32)
+  ; CHECK-NEXT:   RET_ReallyLR implicit $w0
+    %res = icmp samesign ult i32 %a, %b
+    ret i1 %res
+}
+
+define <2 x i1> @icmp_samesign2(<2 x i32> %a, <2 x i32> %b) {
+  ; CHECK-LABEL: name: icmp_samesign2
+  ; CHECK: bb.1 (%ir-block.0):
+  ; CHECK-NEXT:   liveins: $d0, $d1
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   [[COPY:%[0-9]+]]:_(<2 x s32>) = COPY $d0
+  ; CHECK-NEXT:   [[COPY1:%[0-9]+]]:_(<2 x s32>) = COPY $d1
+  ; CHECK-NEXT:   %2:_(<2 x s1>) = samesign G_ICMP intpred(ult), [[COPY]](<2 x s32>), [[COPY1]]
+  ; CHECK-NEXT:   [[ANYEXT:%[0-9]+]]:_(<2 x s32>) = G_ANYEXT %2(<2 x s1>)
+  ; CHECK-NEXT:   $d0 = COPY [[ANYEXT]](<2 x s32>)
+  ; CHECK-NEXT:   RET_ReallyLR implicit $d0
+    %res = icmp samesign ult <2 x i32> %a, %b
+    ret <2 x i1> %res
+}
diff --git a/llvm/test/CodeGen/MIR/X86/icmp-samesign.mir b/llvm/test/CodeGen/MIR/X86/icmp-samesign.mir
@@ -0,0 +1,40 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 5
+# RUN: llc -x mir -march=x86-64 -run-pass=none -verify-machineinstrs %s -o - | FileCheck %s
+
+--- |
+
+  define i1 @icmp_samesign(i32 %a, i32 %b) {
+  entry:
+    %res = icmp samesign ult i32 %a, %b
+    ret i1 %res
+  }
+
+...
+---
+name:            icmp_samesign
+tracksRegLiveness: true
+registers:
+  - { id: 0, class: gr32 }
+  - { id: 1, class: gr32 }
+liveins:
+  - { reg: '$edi', virtual-reg: '%0' }
+  - { reg: '$esi', virtual-reg: '%1' }
+body:             |
+  bb.0:
+    liveins: $edi, $esi
+
+    ; CHECK-LABEL: name: icmp_samesign
+    ; CHECK: liveins: $edi, $esi
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: %tmp0:_(s32) = COPY $edi
+    ; CHECK-NEXT: %tmp1:_(s32) = COPY $esi
+    ; CHECK-NEXT: %tmp4:_(s8) = samesign G_ICMP intpred(ult), %tmp0(s32), %tmp1
+    ; CHECK-NEXT: $al = COPY %tmp4(s8)
+    ; CHECK-NEXT: RET 0, implicit $al
+    %tmp0:_(s32) = COPY $edi
+    %tmp1:_(s32) = COPY $esi
+    %tmp4:_(s8) = samesign G_ICMP intpred(ult), %tmp0(s32), %tmp1
+    $al = COPY %tmp4(s8)
+    RET 0, implicit $al
+
+...
