[llvm] [IR] Add disjoint flag for Or instructions. (PR #72583)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-ir

Author: Craig Topper (topperc)

<details>
<summary>Changes</summary>

This flag indicates that every bit is known to be zero in at least one of the inputs. This allows the Or to be treated as an Add since there is no possibility of a carry from any bit.

If the flag is present and this property does not hold, the result is poison.

This makes it easier to reverse the InstCombine transform that turns Add into Or.

This is inspired by a comment here https://github.com/llvm/llvm-project/pull/71955#discussion_r1391614578

I will start a discourse thread as well.

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


14 Files Affected:

- (modified) llvm/docs/LangRef.rst (+7) 
- (modified) llvm/include/llvm/AsmParser/LLToken.h (+1) 
- (modified) llvm/include/llvm/Bitcode/LLVMBitCodes.h (+4) 
- (modified) llvm/include/llvm/IR/InstrTypes.h (+17-1) 
- (modified) llvm/include/llvm/IR/Instruction.h (+8) 
- (modified) llvm/lib/AsmParser/LLLexer.cpp (+1) 
- (modified) llvm/lib/AsmParser/LLParser.cpp (+7-1) 
- (modified) llvm/lib/Bitcode/Reader/BitcodeReader.cpp (+3) 
- (modified) llvm/lib/Bitcode/Writer/BitcodeWriter.cpp (+3) 
- (modified) llvm/lib/IR/AsmWriter.cpp (+4) 
- (modified) llvm/lib/IR/Instruction.cpp (+23) 
- (modified) llvm/test/Assembler/flags.ll (+5) 
- (modified) llvm/test/Bitcode/flags.ll (+4) 
- (modified) llvm/test/Transforms/InstCombine/freeze.ll (+11) 


``````````diff
diff --git a/llvm/docs/LangRef.rst b/llvm/docs/LangRef.rst
index bc1eab1e0b7a07f..a4ea477870448fa 100644
--- a/llvm/docs/LangRef.rst
+++ b/llvm/docs/LangRef.rst
@@ -9981,6 +9981,7 @@ Syntax:
 ::
 
       <result> = or <ty> <op1>, <op2>   ; yields ty:result
+      <result> = or disjoint <ty> <op1>, <op2>   ; yields ty:result
 
 Overview:
 """""""""
@@ -10012,6 +10013,12 @@ The truth table used for the '``or``' instruction is:
 |   1 |   1 |   1 |
 +-----+-----+-----+
 
+``disjoint`` means every bit is known to be zero in at least one of the inputs.
+This allows the Or to be treated as an Add since no carry can occur from any
+bit. If the disjoint keyword is present, the result value of the ``or`` is a
+:ref:`poison value <poisonvalues>` if both inputs have a one in any bit
+position. For vectors, only the element containing the bit is poison.
+
 Example:
 """"""""
 
diff --git a/llvm/include/llvm/AsmParser/LLToken.h b/llvm/include/llvm/AsmParser/LLToken.h
index c9dcd29b31955dc..f4b12938590fe18 100644
--- a/llvm/include/llvm/AsmParser/LLToken.h
+++ b/llvm/include/llvm/AsmParser/LLToken.h
@@ -109,6 +109,7 @@ enum Kind {
   kw_nuw,
   kw_nsw,
   kw_exact,
+  kw_disjoint,
   kw_inbounds,
   kw_nneg,
   kw_inrange,
diff --git a/llvm/include/llvm/Bitcode/LLVMBitCodes.h b/llvm/include/llvm/Bitcode/LLVMBitCodes.h
index 9fa70c0671ef340..99a41fa107d0811 100644
--- a/llvm/include/llvm/Bitcode/LLVMBitCodes.h
+++ b/llvm/include/llvm/Bitcode/LLVMBitCodes.h
@@ -512,6 +512,10 @@ enum PossiblyNonNegInstOptionalFlags { PNNI_NON_NEG = 0 };
 /// PossiblyExactOperator's SubclassOptionalData contents.
 enum PossiblyExactOperatorOptionalFlags { PEO_EXACT = 0 };
 
+/// PossiblyDisjointInstOptionalFlags - Flags for serializing
+/// PossiblyDisjointInst's SubclassOptionalData contents.
+enum PossiblyDisjointInstOptionalFlags { PDI_DISJOINT = 0 };
+
 /// Encoded AtomicOrdering values.
 enum AtomicOrderingCodes {
   ORDERING_NOTATOMIC = 0,
diff --git a/llvm/include/llvm/IR/InstrTypes.h b/llvm/include/llvm/IR/InstrTypes.h
index fc5e228168a058b..99145ab9acd7fd3 100644
--- a/llvm/include/llvm/IR/InstrTypes.h
+++ b/llvm/include/llvm/IR/InstrTypes.h
@@ -415,6 +415,22 @@ struct OperandTraits<BinaryOperator> :
 
 DEFINE_TRANSPARENT_OPERAND_ACCESSORS(BinaryOperator, Value)
 
+/// A or instruction, which can be marked as "disjoint", indicating that the
+/// inputs don't have a 1 in the same bit position. Meaning this instruction
+/// can also be treated as an add.
+class PossiblyDisjointInst : public BinaryOperator {
+public:
+  enum { IsDisjoint = (1 << 0) };
+
+  static bool classof(const Instruction *I) {
+    return I->getOpcode() == Instruction::Or;
+  }
+
+  static bool classof(const Value *V) {
+    return isa<Instruction>(V) && classof(cast<Instruction>(V));
+  }
+};
+
 //===----------------------------------------------------------------------===//
 //                               CastInst Class
 //===----------------------------------------------------------------------===//
@@ -1085,7 +1101,7 @@ class CmpInst : public Instruction {
   }
 };
 
-// FIXME: these are redundant if CmpInst < BinaryOperator
+// FIXME: these are redundant if CmpInst < ninaryOperator
 template <>
 struct OperandTraits<CmpInst> : public FixedNumOperandTraits<CmpInst, 2> {
 };
diff --git a/llvm/include/llvm/IR/Instruction.h b/llvm/include/llvm/IR/Instruction.h
index 58fc32237367d93..ba5fc35d0d408d5 100644
--- a/llvm/include/llvm/IR/Instruction.h
+++ b/llvm/include/llvm/IR/Instruction.h
@@ -448,6 +448,11 @@ class Instruction : public User,
   /// which supports this flag. See LangRef.html for the meaning of this flag.
   void setIsExact(bool b = true);
 
+  /// Set or clear the disjoint flag on this instruction, which must be an
+  /// operator which supports this flag. See LangRef.html for the meaning of
+  /// this flag.
+  void setIsDisjoint(bool b = true);
+
   /// Set or clear the nneg flag on this instruction, which must be a zext
   /// instruction.
   void setNonNeg(bool b = true);
@@ -500,6 +505,9 @@ class Instruction : public User,
   /// Determine whether the exact flag is set.
   bool isExact() const LLVM_READONLY;
 
+  /// Determine whether the disjoint flag is set.
+  bool isDisjoint() const LLVM_READONLY;
+
   /// Set or clear all fast-math-flags on this instruction, which must be an
   /// operator which supports this flag. See LangRef.html for the meaning of
   /// this flag.
diff --git a/llvm/lib/AsmParser/LLLexer.cpp b/llvm/lib/AsmParser/LLLexer.cpp
index da9e9f4a3c9833b..854aa9cca2c5e37 100644
--- a/llvm/lib/AsmParser/LLLexer.cpp
+++ b/llvm/lib/AsmParser/LLLexer.cpp
@@ -564,6 +564,7 @@ lltok::Kind LLLexer::LexIdentifier() {
   KEYWORD(nuw);
   KEYWORD(nsw);
   KEYWORD(exact);
+  KEYWORD(disjoint);
   KEYWORD(inbounds);
   KEYWORD(nneg);
   KEYWORD(inrange);
diff --git a/llvm/lib/AsmParser/LLParser.cpp b/llvm/lib/AsmParser/LLParser.cpp
index f9df70fb6fc0996..0c170d8da9b73eb 100644
--- a/llvm/lib/AsmParser/LLParser.cpp
+++ b/llvm/lib/AsmParser/LLParser.cpp
@@ -6368,8 +6368,14 @@ int LLParser::parseInstruction(Instruction *&Inst, BasicBlock *BB,
   case lltok::kw_srem:
     return parseArithmetic(Inst, PFS, KeywordVal,
                            /*IsFP*/ false);
+  case lltok::kw_or: {
+    bool Disjoint = EatIfPresent(lltok::kw_disjoint);
+    if (parseLogical(Inst, PFS, KeywordVal))
+      return true;
+    if (Disjoint) cast<PossiblyDisjointInst>(Inst)->setIsDisjoint(true);
+    return false;
+  }
   case lltok::kw_and:
-  case lltok::kw_or:
   case lltok::kw_xor:
     return parseLogical(Inst, PFS, KeywordVal);
   case lltok::kw_icmp:
diff --git a/llvm/lib/Bitcode/Reader/BitcodeReader.cpp b/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
index 76431e883b8d96d..e5aaa56f575c3ab 100644
--- a/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
+++ b/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
@@ -4870,6 +4870,9 @@ Error BitcodeReader::parseFunctionBody(Function *F) {
                    Opc == Instruction::AShr) {
           if (Record[OpNum] & (1 << bitc::PEO_EXACT))
             cast<BinaryOperator>(I)->setIsExact(true);
+        } else if (Opc == Instruction::Or) {
+          if (Record[OpNum] & (1 << bitc::PDI_DISJOINT))
+            cast<BinaryOperator>(I)->setIsDisjoint(true);
         } else if (isa<FPMathOperator>(I)) {
           FastMathFlags FMF = getDecodedFastMathFlags(Record[OpNum]);
           if (FMF.any())
diff --git a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
index d16b5c7781c2413..135801a5c61c434 100644
--- a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
+++ b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
@@ -1541,6 +1541,9 @@ static uint64_t getOptimizationFlags(const Value *V) {
   } else if (const auto *PEO = dyn_cast<PossiblyExactOperator>(V)) {
     if (PEO->isExact())
       Flags |= 1 << bitc::PEO_EXACT;
+  } else if (const auto *PDI = dyn_cast<PossiblyDisjointInst>(V)) {
+    if (PDI->isDisjoint())
+      Flags |= 1 << bitc::PDI_DISJOINT;
   } else if (const auto *FPMO = dyn_cast<FPMathOperator>(V)) {
     if (FPMO->hasAllowReassoc())
       Flags |= bitc::AllowReassoc;
diff --git a/llvm/lib/IR/AsmWriter.cpp b/llvm/lib/IR/AsmWriter.cpp
index 6d66b34423949fb..688f1d7e078eaec 100644
--- a/llvm/lib/IR/AsmWriter.cpp
+++ b/llvm/lib/IR/AsmWriter.cpp
@@ -1355,6 +1355,10 @@ static void WriteOptimizationInfo(raw_ostream &Out, const User *U) {
                dyn_cast<PossiblyExactOperator>(U)) {
     if (Div->isExact())
       Out << " exact";
+  } else if (const PossiblyDisjointInst *PDI =
+               dyn_cast<PossiblyDisjointInst>(U)) {
+    if (PDI->isDisjoint())
+      Out << " disjoint";
   } else if (const GEPOperator *GEP = dyn_cast<GEPOperator>(U)) {
     if (GEP->isInBounds())
       Out << " inbounds";
diff --git a/llvm/lib/IR/Instruction.cpp b/llvm/lib/IR/Instruction.cpp
index 7449692f05d7bf9..fcf79f6cd8c4615 100644
--- a/llvm/lib/IR/Instruction.cpp
+++ b/llvm/lib/IR/Instruction.cpp
@@ -317,6 +317,12 @@ void Instruction::setIsExact(bool b) {
   cast<PossiblyExactOperator>(this)->setIsExact(b);
 }
 
+void Instruction::setIsDisjoint(bool b) {
+  assert(isa<PossiblyDisjointInst>(this) && "Must be or");
+  SubclassOptionalData = (SubclassOptionalData & ~PossiblyDisjointInst::IsDisjoint) |
+                         (b * PossiblyDisjointInst::IsDisjoint);
+}
+
 void Instruction::setNonNeg(bool b) {
   assert(isa<PossiblyNonNegInst>(this) && "Must be zext");
   SubclassOptionalData = (SubclassOptionalData & ~PossiblyNonNegInst::NonNeg) |
@@ -357,6 +363,10 @@ void Instruction::dropPoisonGeneratingFlags() {
     cast<PossiblyExactOperator>(this)->setIsExact(false);
     break;
 
+  case Instruction::Or:
+    cast<PossiblyDisjointInst>(this)->setIsDisjoint(false);
+    break;
+
   case Instruction::GetElementPtr:
     cast<GetElementPtrInst>(this)->setIsInBounds(false);
     break;
@@ -419,6 +429,11 @@ bool Instruction::isExact() const {
   return cast<PossiblyExactOperator>(this)->isExact();
 }
 
+bool Instruction::isDisjoint() const {
+  assert(isa<PossiblyDisjointInst>(this) && "Must be or");
+  return (SubclassOptionalData & PossiblyDisjointInst::IsDisjoint) != 0;
+}
+
 void Instruction::setFast(bool B) {
   assert(isa<FPMathOperator>(this) && "setting fast-math flag on invalid op");
   cast<FPMathOperator>(this)->setFast(B);
@@ -532,6 +547,10 @@ void Instruction::copyIRFlags(const Value *V, bool IncludeWrapFlags) {
     if (isa<PossiblyExactOperator>(this))
       setIsExact(PE->isExact());
 
+  if (auto *PD = dyn_cast<PossiblyDisjointInst>(V))
+    if (isa<PossiblyDisjointInst>(this))
+      setIsDisjoint(PD->isDisjoint());
+
   // Copy the fast-math flags.
   if (auto *FP = dyn_cast<FPMathOperator>(V))
     if (isa<FPMathOperator>(this))
@@ -558,6 +577,10 @@ void Instruction::andIRFlags(const Value *V) {
     if (isa<PossiblyExactOperator>(this))
       setIsExact(isExact() && PE->isExact());
 
+  if (auto *PE = dyn_cast<PossiblyDisjointInst>(V))
+    if (isa<PossiblyDisjointInst>(this))
+      setIsDisjoint(isDisjoint() && PE->isDisjoint());
+
   if (auto *FP = dyn_cast<FPMathOperator>(V)) {
     if (isa<FPMathOperator>(this)) {
       FastMathFlags FM = getFastMathFlags();
diff --git a/llvm/test/Assembler/flags.ll b/llvm/test/Assembler/flags.ll
index 6ab5e1bfb9c4f46..04bddd02f50c814 100644
--- a/llvm/test/Assembler/flags.ll
+++ b/llvm/test/Assembler/flags.ll
@@ -256,3 +256,8 @@ define i64 @test_zext(i32 %a) {
   ret i64 %res
 }
 
+define i64 @test_or(i64 %a, i64 %b) {
+; CHECK: %res = or disjoint i64 %a, %b
+  %res = or disjoint i64 %a, %b
+  ret i64 %res
+}
diff --git a/llvm/test/Bitcode/flags.ll b/llvm/test/Bitcode/flags.ll
index a6e368b7e76327f..e3fc827d865d7e2 100644
--- a/llvm/test/Bitcode/flags.ll
+++ b/llvm/test/Bitcode/flags.ll
@@ -18,6 +18,8 @@ second:                                           ; preds = %first
   %z = add i32 %a, 0                              ; <i32> [#uses=0]
   %hh = zext nneg i32 %a to i64
   %ll = zext i32 %s to i64
+  %jj = or disjoint i32 %a, 0
+  %oo = or i32 %a, 0
   unreachable
 
 first:                                            ; preds = %entry
@@ -28,5 +30,7 @@ first:                                            ; preds = %entry
   %zz = add i32 %a, 0                             ; <i32> [#uses=0]
   %kk = zext nneg i32 %a to i64
   %rr = zext i32 %ss to i64
+  %mm = or disjoint i32 %a, 0
+  %nn = or i32 %a, 0
   br label %second
 }
diff --git a/llvm/test/Transforms/InstCombine/freeze.ll b/llvm/test/Transforms/InstCombine/freeze.ll
index dd9272b4b35f193..da59101d5710cb5 100644
--- a/llvm/test/Transforms/InstCombine/freeze.ll
+++ b/llvm/test/Transforms/InstCombine/freeze.ll
@@ -1127,6 +1127,17 @@ define i32 @freeze_zext_nneg(i8 %x) {
   ret i32 %fr
 }
 
+define i32 @propagate_drop_flags_or(i32 %arg) {
+; CHECK-LABEL: @propagate_drop_flags_or(
+; CHECK-NEXT:    [[ARG_FR:%.*]] = freeze i32 [[ARG:%.*]]
+; CHECK-NEXT:    [[V1:%.*]] = or i32 [[ARG_FR]], 2
+; CHECK-NEXT:    ret i32 [[V1]]
+;
+  %v1 = or disjoint i32 %arg, 2
+  %v1.fr = freeze i32 %v1
+  ret i32 %v1.fr
+}
+
 !0 = !{}
 !1 = !{i64 4}
 !2 = !{i32 0, i32 100}

``````````

</details>


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