[llvm] [InstCombine] fold addrcast+load through selects (PR #176352)

[Theodoros Theodoridis via llvm-commits]

================
@@ -1137,20 +1137,59 @@ Instruction *InstCombinerImpl::visitLoadInst(LoadInst &LI) {
     // but it would not be valid if we transformed it to load from null
     // unconditionally.
     //
+    // Check if both arms of a select are safe to load unconditionally.
+    auto CanSpeculateThroughSelect = [&](SelectInst *SI) {
+      Align Alignment = LI.getAlign();
+      return isSafeToLoadUnconditionally(SI->getTrueValue(), LI.getType(),
+                                         Alignment, DL, SI) &&
+             isSafeToLoadUnconditionally(SI->getFalseValue(), LI.getType(),
+                                         Alignment, DL, SI);
+    };
+
+    // Fold:
+    //   load(addrspacecast(select(Cond, &V1, &V2)))
+    // into:
+    //   load(select(Cond, addrspacecast(&V1), addrspacecast(&V2)))
+    // Then the transformation load(select) -> select(load) below will complete
+    // the fold:
+    //   select(Cond, load(addrspacecast(&V1)), load(addrspacecast(&V2)))
+    //
+    // We only fold if the subsequent load(select) transformation will succeed.
+    if (auto *ASC = dyn_cast<AddrSpaceCastInst>(Op)) {
+      unsigned TargetAddressSpace = ASC->getDestAddressSpace();
+      if (auto *SI = dyn_cast<SelectInst>(ASC->getOperand(0))) {
+        if (CanSpeculateThroughSelect(SI)) {
+          if (Instruction *NewSel = FoldOpIntoSelect(
+                  *ASC, SI, /*FoldWithMultiUse=*/false,
+                  /*SimplifyBothArms=*/false, /*FoldWithoutSimplify=*/true)) {
+            unsigned LHSTargetAddressSpace =
+                NewSel->getOperand(1)->getType()->getPointerAddressSpace();
+            unsigned RHSTargetAddressSpace =
+                NewSel->getOperand(2)->getType()->getPointerAddressSpace();
+            assert(LHSTargetAddressSpace == TargetAddressSpace &&
+                   RHSTargetAddressSpace == TargetAddressSpace &&
+                   "Expected same target address space for both arms");
+            NewSel->insertBefore(LI.getIterator());
+            // Replace the load's operand and continue processing with the new
+            // select. The load(select) transformation below will complete the
+            // fold.
+            return replaceOperand(LI, 0, NewSel);
+          }
+        }
+      }
+    }
----------------
thetheodor wrote:

Good catch, thanks. I've simplified the code.

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