[llvm] [CodeGenPrepare] Reject select address combining if a new select's operand does not dominate it (PR #226968)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Aochang Liu (HelloWorldU)

<details>
<summary>Changes</summary>

## Problem

When the address of a memory instruction is a select between two GEPs that differ only in the index, CodeGenPrepare combines them: it builds a new select of the two indices at the original select's position, and sinks one shared GEP to the memory instruction. For a loop-IV index, the combined form reuses the IV increment `(%iv.next)`, which is defined at the end of the loop body, after the original select. The new select then uses it before its definition：

    Instruction does not dominate all uses!
      %iv.next = add i64 %iv, 1
      %sel1 = select i1 false, i64 %iv.next, i64 0

Release builds (no verifier) propagate the invalid IR and crash in the register allocator on larger inputs.

## Fix

After the new selects are filled in, reject the combination if any operand does not dominate its select. Reuses the existing destroyNewNodes rollback pattern in the same function. The alternative of tightening the IV-increment reuse in the matcher was not taken: the matcher cannot know whether its result will feed a new select, nor where that select will be inserted, while the combiner knows both.

## Testing

- New test `sink-addrmode-select-iv.ll` fails on main and passes with this change; a positive control (increment defined before the select) still combines and sinks normally.
- `llvm/test/Transforms/CodeGenPrepare`: 79/79 pass; the issue's C++ reproducer compiles cleanly with `-O3 -fverify-intermediate-code`.

Fixes #<!-- -->226709


Assisted-by: Kimi K3

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


2 Files Affected:

- (modified) llvm/lib/CodeGen/CodeGenPrepare.cpp (+31-7) 
- (added) llvm/test/Transforms/CodeGenPrepare/X86/sink-addrmode-select-iv.ll (+43) 


``````````diff
diff --git a/llvm/lib/CodeGen/CodeGenPrepare.cpp b/llvm/lib/CodeGen/CodeGenPrepare.cpp
index 3e0b8a956ca81..6c6ae30a59212 100644
--- a/llvm/lib/CodeGen/CodeGenPrepare.cpp
+++ b/llvm/lib/CodeGen/CodeGenPrepare.cpp
@@ -4161,6 +4161,10 @@ class SimplificationTracker {
 
   unsigned countNewSelectNodes() const { return AllSelectNodes.size(); }
 
+  const SmallPtrSet<SelectInst *, 32> &newSelectNodes() const {
+    return AllSelectNodes;
+  }
+
   void destroyNewNodes(Type *CommonType) {
     // For safe erasing, replace the uses with dummy value first.
     auto *Dummy = PoisonValue::get(CommonType);
@@ -4203,9 +4207,15 @@ class AddressingModeCombiner {
   /// Common value among addresses
   Value *CommonValue = nullptr;
 
+  /// Deferred getter for the dominator tree, so that it is only computed
+  /// when it is actually needed.
+  const std::function<const DominatorTree &()> getDTFn;
+
 public:
-  AddressingModeCombiner(const DataLayout &DL, Value *OriginalValue)
-      : DL(DL), Original(OriginalValue) {}
+  AddressingModeCombiner(
+      const DataLayout &DL, Value *OriginalValue,
+      const std::function<const DominatorTree &()> &getDTFn)
+      : DL(DL), Original(OriginalValue), getDTFn(getDTFn) {}
 
   ~AddressingModeCombiner() { eraseCommonValueIfDead(); }
 
@@ -4389,6 +4399,20 @@ class AddressingModeCombiner {
       return nullptr;
     }
 
+    // New selects are inserted at the original selects' positions, so their
+    // operands must be available there. The addressing mode matcher only
+    // guarantees that a combined field dominates the memory instruction (e.g.
+    // when it reuses an IV increment), which may be defined after the original
+    // select in the same block. Reject the combination if any new select would
+    // use a value that does not dominate it.
+    const DominatorTree &DT = getDTFn();
+    for (SelectInst *SI : ST.newSelectNodes())
+      for (Value *Op : SI->operands())
+        if (!DT.dominates(Op, SI)) {
+          ST.destroyNewNodes(CommonType);
+          return nullptr;
+        }
+
     auto *Result = ST.Get(Map.find(Original)->second);
     if (Result) {
       NumMemoryInstsPhiCreated += ST.countNewPhiNodes() + PhiNotMatchedCount;
@@ -5916,7 +5940,11 @@ bool CodeGenPrepare::optimizeMemoryInst(Instruction *MemoryInst, Value *Addr,
   // the graph are compatible.
   bool PhiOrSelectSeen = false;
   SmallVector<Instruction *, 16> AddrModeInsts;
-  AddressingModeCombiner AddrModes(*DL, Addr);
+  // Defer the query (and possible computation of) the dom tree to point of
+  // actual use.  It's expected that most address matches don't actually need
+  // the domtree.
+  auto getDTFn = [this]() -> const DominatorTree & { return getDT(); };
+  AddressingModeCombiner AddrModes(*DL, Addr, getDTFn);
   TypePromotionTransaction TPT(RemovedInsts);
   TypePromotionTransaction::ConstRestorationPt LastKnownGood =
       TPT.getRestorationPoint();
@@ -5955,10 +5983,6 @@ bool CodeGenPrepare::optimizeMemoryInst(Instruction *MemoryInst, Value *Addr,
     AddrModeInsts.clear();
     std::pair<AssertingVH<GetElementPtrInst>, int64_t> LargeOffsetGEP(nullptr,
                                                                       0);
-    // Defer the query (and possible computation of) the dom tree to point of
-    // actual use.  It's expected that most address matches don't actually need
-    // the domtree.
-    auto getDTFn = [this]() -> const DominatorTree & { return getDT(); };
     ExtAddrMode NewAddrMode = AddressingModeMatcher::Match(
         V, AccessTy, AddrSpace, MemoryInst, AddrModeInsts, *TLI, *LI, getDTFn,
         *TRI, InsertedInsts, PromotedInsts, TPT, LargeOffsetGEP, OptSize, PSI,
diff --git a/llvm/test/Transforms/CodeGenPrepare/X86/sink-addrmode-select-iv.ll b/llvm/test/Transforms/CodeGenPrepare/X86/sink-addrmode-select-iv.ll
new file mode 100644
index 0000000000000..460af48dc4df0
--- /dev/null
+++ b/llvm/test/Transforms/CodeGenPrepare/X86/sink-addrmode-select-iv.ll
@@ -0,0 +1,43 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -S -passes="require<profile-summary>,function(codegenprepare)" -mtriple=x86_64 < %s | FileCheck %s
+
+; Address sinking through a select may reuse an IV increment as the combined
+; index, but the increment is only guaranteed to dominate the memory
+; instruction, not the position where the new select is inserted. The
+; combination must be rejected in that case; otherwise the new select uses
+; %iv.next before its definition.
+; https://github.com/llvm/llvm-project/issues/226709
+
+define i64 @f(ptr %p, i1 %cond) {
+; CHECK-LABEL: define i64 @f(
+; CHECK-SAME: ptr [[P:%.*]], i1 [[COND:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    br label %[[LOOP:.*]]
+; CHECK:       [[LOOP]]:
+; CHECK-NEXT:    [[IV:%.*]] = phi i64 [ 0, %[[ENTRY]] ], [ [[IV_NEXT:%.*]], %[[LOOP]] ]
+; CHECK-NEXT:    [[A:%.*]] = getelementptr [8 x i8], ptr [[P]], i64 [[IV]]
+; CHECK-NEXT:    [[B:%.*]] = getelementptr i8, ptr [[A]], i64 -8
+; CHECK-NEXT:    [[C:%.*]] = getelementptr [8 x i8], ptr [[P]], i64 -2
+; CHECK-NEXT:    [[SEL:%.*]] = select i1 false, ptr [[B]], ptr [[C]]
+; CHECK-NEXT:    [[IV_NEXT]] = add i64 [[IV]], 1
+; CHECK-NEXT:    br i1 [[COND]], label %[[EXIT:.*]], label %[[LOOP]]
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    [[V:%.*]] = load i64, ptr [[SEL]], align 8
+; CHECK-NEXT:    ret i64 [[V]]
+;
+entry:
+  br label %loop
+
+loop:
+  %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
+  %a = getelementptr [8 x i8], ptr %p, i64 %iv
+  %b = getelementptr i8, ptr %a, i64 -8
+  %c = getelementptr [8 x i8], ptr %p, i64 -2
+  %sel = select i1 false, ptr %b, ptr %c
+  %iv.next = add i64 %iv, 1
+  br i1 %cond, label %exit, label %loop
+
+exit:
+  %v = load i64, ptr %sel, align 8
+  ret i64 %v
+}

``````````

</details>


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