[llvm] [Transform] Treat umul + extract pattern as cheap single instruction. (PR #124933)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Gábor Spaits (spaits)

<details>
<summary>Changes</summary>

This pr was created to fix #<!-- -->115683 .

In the comments @<!-- -->dtcxzyw has suggested, to treat the the pattern `extract oneuse(op.with.overflow),1` as one instruction.
If we want the `speculativelyExecuteBB` function to speculatively execute this block, then we must do that, since it only speculatively executes a block, if that has at most two instructions.

I was also thinking, maybe I should create a mechanism to add more cheap patterns like this one, and in the feature we just provide the pattern and don't have to create a add a new `if (pattern && len(thenBlock) == 1 + (len(pattern))) then break` to the `speculativelyExecuteBB` function.

What do you think? Is this solution correct? Or should I try to do this in a different way?
Any input is welcomed.

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


2 Files Affected:

- (modified) llvm/lib/Transforms/Utils/SimplifyCFG.cpp (+17-5) 
- (added) llvm/test/Transforms/SimplifyCFG/umul-extract-pattern.ll (+27) 


``````````diff
diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 12dd49da279b9c..5197db70285cec 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -3290,9 +3290,11 @@ bool SimplifyCFGOpt::speculativelyExecuteBB(BranchInst *BI,
   bool HoistLoadsStores = HoistLoadsStoresWithCondFaulting &&
                           Options.HoistLoadsStoresWithCondFaulting;
   SmallVector<Instruction *, 2> SpeculatedConditionalLoadsStores;
+  InstructionCost BlockCostSoFar = 0;
   Value *SpeculatedStoreValue = nullptr;
   StoreInst *SpeculatedStore = nullptr;
   EphemeralValueTracker EphTracker;
+  bool PatternFound = false;
   for (Instruction &I : reverse(drop_end(*ThenBB))) {
     // Skip debug info.
     if (isa<DbgInfoIntrinsic>(I)) {
@@ -3329,9 +3331,6 @@ bool SimplifyCFGOpt::speculativelyExecuteBB(BranchInst *BI,
     else
       ++SpeculatedInstructions;
 
-    if (SpeculatedInstructions > 1)
-      return false;
-
     // Don't hoist the instruction if it's unsafe or expensive.
     if (!IsSafeCheapLoadStore &&
         !isSafeToSpeculativelyExecute(&I, BI, Options.AC) &&
@@ -3339,10 +3338,23 @@ bool SimplifyCFGOpt::speculativelyExecuteBB(BranchInst *BI,
           (SpeculatedStoreValue =
                isSafeToSpeculateStore(&I, BB, ThenBB, EndBB))))
       return false;
-    if (!IsSafeCheapLoadStore && !SpeculatedStoreValue &&
-        computeSpeculationCost(&I, TTI) >
+
+    if (match(&I,
+              m_ExtractValue<1>(m_OneUse(
+                  m_Intrinsic<Intrinsic::umul_with_overflow>(m_Value())))) &&
+        ThenBB->size() <= 3) {
+      PatternFound = true;
+    }
+
+    BlockCostSoFar += computeSpeculationCost(&I, TTI);
+    if (! PatternFound && !IsSafeCheapLoadStore && !SpeculatedStoreValue &&
+        BlockCostSoFar >
             PHINodeFoldingThreshold * TargetTransformInfo::TCC_Basic)
       return false;
+    // If we don't find any pattern, that must be cheap, then only speculatively
+    // execute a single instruction (not counting the terminator).
+    if (!PatternFound && SpeculatedInstructions > 1)
+      return false;
 
     // Store the store speculation candidate.
     if (!SpeculatedStore && SpeculatedStoreValue)
diff --git a/llvm/test/Transforms/SimplifyCFG/umul-extract-pattern.ll b/llvm/test/Transforms/SimplifyCFG/umul-extract-pattern.ll
new file mode 100644
index 00000000000000..72610218c314ec
--- /dev/null
+++ b/llvm/test/Transforms/SimplifyCFG/umul-extract-pattern.ll
@@ -0,0 +1,27 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S < %s | FileCheck %s
+
+define dso_local signext range(i16 0, 2) i16 @func2(i64 noundef %x, i64 noundef %y) local_unnamed_addr #0 {
+; CHECK-LABEL: @func2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CMP_NOT:%.*]] = icmp eq i64 [[Y:%.*]], 0
+; CHECK-NEXT:    [[MUL:%.*]] = tail call { i64, i1 } @llvm.umul.with.overflow.i64(i64 [[Y]], i64 [[X:%.*]])
+; CHECK-NEXT:    [[MUL_OV:%.*]] = extractvalue { i64, i1 } [[MUL]], 1
+; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[CMP_NOT]], i1 false, i1 [[MUL_OV]]
+; CHECK-NEXT:    [[CONV:%.*]] = zext i1 [[SPEC_SELECT]] to i16
+; CHECK-NEXT:    ret i16 [[CONV]]
+;
+entry:
+  %cmp.not = icmp eq i64 %y, 0
+  br i1 %cmp.not, label %land.end, label %land.rhs
+
+land.rhs:                                         ; preds = %entry
+  %mul = tail call { i64, i1 } @llvm.umul.with.overflow.i64(i64 %y, i64 %x)
+  %mul.ov = extractvalue { i64, i1 } %mul, 1
+  br label %land.end
+
+land.end:                                         ; preds = %land.rhs, %entry
+  %0 = phi i1 [ false, %entry ], [ %mul.ov, %land.rhs ]
+  %conv = zext i1 %0 to i16
+  ret i16 %conv
+}

``````````

</details>


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