detect optimization pattern to stop reprocessing

Author: hassnaaHamdi

llvm/lib/CodeGen/CodeGenPrepare.cpp

@@ -338,10 +338,6 @@ class CodeGenPrepare {
   /// Keep track of instructions removed during promotion.
   SetOfInstrs RemovedInsts;
 
-  /// Keep track of seen mul_with_overflow intrinsics to avoid
-  // reprocessing them.
-  DenseMap<Instruction *, bool> SeenMulWithOverflowInstrs;
-
   /// Keep track of sext chains based on their initial value.
   DenseMap<Value *, Instruction *> SeenChainsForSExt;
 
@@ -778,7 +774,6 @@ bool CodeGenPrepare::_run(Function &F) {
     verifyBFIUpdates(F);
 #endif
 
-  SeenMulWithOverflowInstrs.clear();
   return EverMadeChange;
 }
 
@@ -6409,9 +6404,36 @@ bool CodeGenPrepare::optimizeMulWithOverflow(Instruction *I, bool IsSigned,
                                              ModifyDT &ModifiedDT) {
   if (!TLI->shouldOptimizeMulOverflowIntrinsic())
     return false;
-  // If we have already seen this instruction, don't process it again.
-  if (!SeenMulWithOverflowInstrs.insert(std::make_pair(I, true)).second)
-    return false;
+
+  // Check if we had already optimized this intrinsic by detecting the pattern of the changes we had made:
+  // Check if we are testing the high bits of the operands:
+  if (BasicBlock *BrBB = I->getParent()->getSinglePredecessor()) {
+    if (BranchInst *Br = dyn_cast<BranchInst>(BrBB->getTerminator()); Br && Br->isConditional()) {
+      if (IsSigned) {
+        // Check: cmp(or(xor(trunc(lshr(x))), xor(trunc(lshr(x)))))
+        if (match(Br->getCondition(),
+                  m_Cmp(m_Or(m_Xor(m_Trunc(m_LShr(m_Specific(I->getOperand(0)), m_Value())), m_Value()),
+                             m_Xor(m_Trunc(m_LShr(m_Specific(I->getOperand(1)), m_Value())), m_Value())),
+                        m_Value()))) {
+          LLVM_DEBUG(dbgs() << "CGP: pattern detected - bail out\n");
+          // Pattern detected, bail out.
+          return false;
+        }
+      }
+      else
+      {
+        // Check: or(cmp(trunc(lshr(x)), cmp(trunc(lshr(y))))
+        if (match(Br->getCondition(),
+                  m_Or(m_Cmp(m_Trunc(m_LShr(m_Specific(I->getOperand(0)), m_Value())), m_Value()),
+                       m_Cmp(m_Trunc(m_LShr(m_Specific(I->getOperand(1)), m_Value())), m_Value())))) {
+          LLVM_DEBUG(dbgs() << "CGP: pattern detected - bail out\n");
+          // Pattern detected, bail out.
+          return false;
+        }
+      }
+
+    }
+  }
 
   if (TLI->getTypeAction(
           I->getContext(),

llvm/test/CodeGen/AArch64/mul_with_overflow.ll

@@ -0,0 +1,132 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; REQUIRES: asserts
+; RUN: llc < %s -mtriple=aarch64-unknown-linux-gnu -debug-only=codegenprepare 2>%t | FileCheck %s --check-prefixes=CHECK
+; RUN: cat %t | FileCheck %s --check-prefix=DEBUG
+
+; DEBUG: CGP: pattern detected - bail out
+
+define i128 @test_umul_i128(i128 noundef %x, i128 noundef %y) {
+; CHECK-LABEL: test_umul_i128:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    orr x8, x1, x3
+; CHECK-NEXT:    cbz x8, .LBB0_3
+; CHECK-NEXT:  // %bb.1: // %overflow
+; CHECK-NEXT:    mul x9, x3, x0
+; CHECK-NEXT:    cmp x1, #0
+; CHECK-NEXT:    ccmp x3, #0, #4, ne
+; CHECK-NEXT:    umulh x10, x1, x2
+; CHECK-NEXT:    umulh x8, x3, x0
+; CHECK-NEXT:    madd x9, x1, x2, x9
+; CHECK-NEXT:    ccmp xzr, x10, #0, eq
+; CHECK-NEXT:    umulh x11, x0, x2
+; CHECK-NEXT:    ccmp xzr, x8, #0, eq
+; CHECK-NEXT:    cset w8, ne
+; CHECK-NEXT:    adds x1, x11, x9
+; CHECK-NEXT:    csinc w8, w8, wzr, lo
+; CHECK-NEXT:    cmp w8, #1
+; CHECK-NEXT:    b.ne .LBB0_4
+; CHECK-NEXT:  // %bb.2: // %if.then
+; CHECK-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    .cfi_offset w30, -16
+; CHECK-NEXT:    bl error
+; CHECK-NEXT:    // kill: def $w0 killed $w0 def $x0
+; CHECK-NEXT:    sxtw x0, w0
+; CHECK-NEXT:    asr x1, x0, #63
+; CHECK-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
+; CHECK-NEXT:    ret
+; CHECK-NEXT:  .LBB0_3: // %overflow.no
+; CHECK-NEXT:    umulh x1, x0, x2
+; CHECK-NEXT:  .LBB0_4:
+; CHECK-NEXT:    mul x0, x0, x2
+; CHECK-NEXT:    ret
+entry:
+  %0 = tail call { i128, i1 } @llvm.umul.with.overflow.i128(i128 %x, i128 %y)
+  %1 = extractvalue { i128, i1 } %0, 1
+  br i1 %1, label %if.then, label %if.end
+
+if.then:
+  %call = tail call i32 @error()
+  %conv1 = sext i32 %call to i128
+  br label %cleanup
+
+if.end:
+  %2 = extractvalue { i128, i1 } %0, 0
+  br label %cleanup
+
+cleanup:
+  %retval.0 = phi i128 [ %conv1, %if.then ], [ %2, %if.end ]
+  ret i128 %retval.0
+}
+
+; DEBUG: CGP: pattern detected - bail out
+
+define i128 @test_smul_i128(i128 noundef %x, i128 noundef %y) {
+; CHECK-LABEL: test_smul_i128:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    eor x8, x3, x2, asr #63
+; CHECK-NEXT:    eor x9, x1, x0, asr #63
+; CHECK-NEXT:    orr x8, x9, x8
+; CHECK-NEXT:    cbz x8, .LBB1_3
+; CHECK-NEXT:  // %bb.1: // %overflow
+; CHECK-NEXT:    asr x9, x1, #63
+; CHECK-NEXT:    umulh x10, x0, x2
+; CHECK-NEXT:    asr x13, x3, #63
+; CHECK-NEXT:    mul x11, x1, x2
+; CHECK-NEXT:    umulh x8, x1, x2
+; CHECK-NEXT:    mul x9, x9, x2
+; CHECK-NEXT:    adds x10, x11, x10
+; CHECK-NEXT:    mul x14, x0, x3
+; CHECK-NEXT:    umulh x12, x0, x3
+; CHECK-NEXT:    adc x8, x8, x9
+; CHECK-NEXT:    mov x9, x1
+; CHECK-NEXT:    mul x13, x0, x13
+; CHECK-NEXT:    asr x11, x8, #63
+; CHECK-NEXT:    mul x15, x1, x3
+; CHECK-NEXT:    adds x1, x14, x10
+; CHECK-NEXT:    smulh x9, x9, x3
+; CHECK-NEXT:    adc x10, x12, x13
+; CHECK-NEXT:    asr x12, x10, #63
+; CHECK-NEXT:    adds x8, x8, x10
+; CHECK-NEXT:    adc x10, x11, x12
+; CHECK-NEXT:    adds x8, x15, x8
+; CHECK-NEXT:    asr x11, x1, #63
+; CHECK-NEXT:    adc x9, x9, x10
+; CHECK-NEXT:    cmp x9, x11
+; CHECK-NEXT:    ccmp x8, x11, #0, eq
+; CHECK-NEXT:    b.eq .LBB1_4
+; CHECK-NEXT:  // %bb.2: // %if.then
+; CHECK-NEXT:    str x30, [sp, #-16]! // 8-byte Folded Spill
+; CHECK-NEXT:    .cfi_def_cfa_offset 16
+; CHECK-NEXT:    .cfi_offset w30, -16
+; CHECK-NEXT:    bl error
+; CHECK-NEXT:    // kill: def $w0 killed $w0 def $x0
+; CHECK-NEXT:    sxtw x0, w0
+; CHECK-NEXT:    asr x1, x0, #63
+; CHECK-NEXT:    ldr x30, [sp], #16 // 8-byte Folded Reload
+; CHECK-NEXT:    ret
+; CHECK-NEXT:  .LBB1_3: // %overflow.no
+; CHECK-NEXT:    smulh x1, x0, x2
+; CHECK-NEXT:  .LBB1_4:
+; CHECK-NEXT:    mul x0, x0, x2
+; CHECK-NEXT:    ret
+entry:
+  %0 = tail call { i128, i1 } @llvm.smul.with.overflow.i128(i128 %x, i128 %y)
+  %1 = extractvalue { i128, i1 } %0, 1
+  br i1 %1, label %if.then, label %if.end
+
+if.then:
+  %call = tail call i32 @error()
+  %conv1 = sext i32 %call to i128
+  br label %cleanup
+
+if.end:
+  %2 = extractvalue { i128, i1 } %0, 0
+  br label %cleanup
+
+cleanup:
+  %retval.0 = phi i128 [ %conv1, %if.then ], [ %2, %if.end ]
+  ret i128 %retval.0
+}
+
+declare i32 @error()