[InstCombine] Fix miscompile when folding a select into a masked load

[Chennesxu]

visitSelectInst folds:

select(mask, masked.load(ptr, mask, PT), FV)

into:

masked.load(ptr, mask, FV)

The replacement load was previously created at the select, effectively moving the memory access past any intervening instructions. If one of them writes the loaded memory, the replacement load reads the updated value instead of the original one. This was also observed downstream in ispc/ispc#3891.

The fold was added in eb8589987267. The issue is labelled regression:22, so it affects LLVM 22.1 as well as current trunk.

Create the replacement load at the original load's position and require FV to be available there. Otherwise, leave the select unchanged. Requiring FV to be available at the original load means the fold no longer fires when FV is computed between the load and the select. No existing llvm/test/Transforms checks change as a result of this restriction.

Tests cover an intervening aliasing store and the case where FV is unavailable at the original load. They also guard against carrying over call-site attributes such as range and noundef when those attributes no longer apply.

Fixes #215453

[llvmorg-github-actions]

@llvm/pr-subscribers-llvm-transforms

Author: Chennes (Chennesxu)

Changes

visitSelectInst folds:

select(mask, masked.load(ptr, mask, PT), FV)

into:

masked.load(ptr, mask, FV)

The replacement load was previously created at the select, effectively moving the memory access past any intervening instructions. If one of them writes the loaded memory, the replacement load reads the updated value instead of the original one. This was also observed downstream in ispc/ispc#3891.

The fold was added in eb8589987267. The issue is labelled regression:22, so it affects LLVM 22.1 as well as current trunk.

Create the replacement load at the original load's position and require FV to be available there. Otherwise, leave the select unchanged. Requiring FV to be available at the original load means the fold no longer fires when FV is computed between the load and the select. No existing llvm/test/Transforms checks change as a result of this restriction.

Changing the insertion point picks up the original load's debug location, so explicitly restore the location of the select being replaced.

Tests cover an intervening aliasing store and the case where FV is unavailable at the original load. They also guard against carrying over call-site attributes such as range and noundef when those attributes no longer apply.

Fixes #215453

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp (+15-5)
• (modified) llvm/test/Transforms/InstCombine/select-masked_load.ll (+41)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index 558ad2ebccc37..75346989cb27c 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -5379,11 +5379,21 @@ Instruction *InstCombinerImpl::visitSelectInst(SelectInst &SI) {
   if (match(TrueVal, m_OneUse(m_MaskedLoad(m_Value(MaskedLoadPtr),
                                            m_Specific(CondVal), m_Value())))) {
     auto *LoadInst = cast<IntrinsicInst>(TrueVal);
-    Instruction *In = Builder.CreateMaskedLoad(
-        TrueVal->getType(), MaskedLoadPtr,
-        LoadInst->getParamAlign(0).valueOrOne(), CondVal, FalseVal);
-    In->setAAMetadata(LoadInst->getAAMetadata());
-    return replaceInstUsesWith(SI, In);
+    // Keep the load at its original position to avoid crossing writes. The new
+    // passthrough must therefore be available there.
+    // TODO: Sink the load when the passthrough is unavailable but no
+    // intervening instruction can modify memory.
+    if (DT.dominates(FalseVal, LoadInst)) {
+      Builder.SetInsertPoint(LoadInst);
+      // SetInsertPoint() took the debug location from the old load, but the new
+      // load replaces the select, so restore the select's location.
+      Builder.SetCurrentDebugLocation(SI.getDebugLoc());
+      Instruction *In = Builder.CreateMaskedLoad(
+          TrueVal->getType(), MaskedLoadPtr,
+          LoadInst->getParamAlign(0).valueOrOne(), CondVal, FalseVal);
+      In->setAAMetadata(LoadInst->getAAMetadata());
+      return replaceInstUsesWith(SI, In);
+    }
   }
 
   // Canonicalize sign function ashr pattern: select (icmp slt X, 1), ashr X,
diff --git a/llvm/test/Transforms/InstCombine/select-masked_load.ll b/llvm/test/Transforms/InstCombine/select-masked_load.ll
index cc6c48b29bf28..f0c2e8c8845cf 100644
--- a/llvm/test/Transforms/InstCombine/select-masked_load.ll
+++ b/llvm/test/Transforms/InstCombine/select-masked_load.ll
@@ -169,6 +169,46 @@ define <vscale x 4 x i32> @fold_sel_into_masked_load_drop_metadata(ptr %loc, <vs
   ret <vscale x 4 x i32> %sel
 }
 
+; Keep the folded load before an intervening aliasing store.
+define <4 x float> @fold_sel_into_masked_load_aliasing_store(ptr %ptr, <4 x i1> %mask, <4 x float> %passthrough) {
+; CHECK-LABEL: @fold_sel_into_masked_load_aliasing_store(
+; CHECK-NEXT:    [[SEL:%.*]] = call <4 x float> @llvm.masked.load.v4f32.p0(ptr align 4 [[PTR:%.*]], <4 x i1> [[MASK:%.*]], <4 x float> [[PASSTHROUGH:%.*]])
+; CHECK-NEXT:    store <4 x float> [[PASSTHROUGH]], ptr [[PTR]], align 16
+; CHECK-NEXT:    ret <4 x float> [[SEL]]
+;
+  %load = call <4 x float> @llvm.masked.load.v4f32.p0(ptr %ptr, i32 4, <4 x i1> %mask, <4 x float> zeroinitializer)
+  store <4 x float> %passthrough, ptr %ptr, align 16
+  %sel = select <4 x i1> %mask, <4 x float> %load, <4 x float> %passthrough
+  ret <4 x float> %sel
+}
+
+; Do not fold when the new passthrough is unavailable at the old load.
+define <4 x float> @neg_fold_sel_into_masked_load_passthrough_after_load(ptr %ptr, <4 x i1> %mask, <4 x float> %a) {
+; CHECK-LABEL: @neg_fold_sel_into_masked_load_passthrough_after_load(
+; CHECK-NEXT:    [[LOAD:%.*]] = call <4 x float> @llvm.masked.load.v4f32.p0(ptr align 4 [[PTR:%.*]], <4 x i1> [[MASK:%.*]], <4 x float> zeroinitializer)
+; CHECK-NEXT:    [[PASSTHROUGH:%.*]] = fadd <4 x float> [[A:%.*]], [[A]]
+; CHECK-NEXT:    [[SEL:%.*]] = select <4 x i1> [[MASK]], <4 x float> [[LOAD]], <4 x float> [[PASSTHROUGH]]
+; CHECK-NEXT:    ret <4 x float> [[SEL]]
+;
+  %load = call <4 x float> @llvm.masked.load.v4f32.p0(ptr %ptr, i32 4, <4 x i1> %mask, <4 x float> zeroinitializer)
+  %passthrough = fadd <4 x float> %a, %a
+  %sel = select <4 x i1> %mask, <4 x float> %load, <4 x float> %passthrough
+  ret <4 x float> %sel
+}
+
+; Do not copy result or passthrough attributes (range/noundef) to the new load.
+; Use the current intrinsic form because auto-upgrading the legacy form drops
+; these attributes before InstCombine.
+define <8 x i16> @fold_sel_into_masked_load_drop_attrs(ptr %ptr, <8 x i1> %mask, <8 x i16> %passthrough) {
+; CHECK-LABEL: @fold_sel_into_masked_load_drop_attrs(
+; CHECK-NEXT:    [[SEL:%.*]] = call <8 x i16> @llvm.masked.load.v8i16.p0(ptr align 2 [[PTR:%.*]], <8 x i1> [[MASK:%.*]], <8 x i16> [[PASSTHROUGH:%.*]])
+; CHECK-NEXT:    ret <8 x i16> [[SEL]]
+;
+  %load = call range(i16 0, 2) <8 x i16> @llvm.masked.load.v8i16.p0(ptr align 2 %ptr, <8 x i1> %mask, <8 x i16> noundef zeroinitializer)
+  %sel = select <8 x i1> %mask, <8 x i16> %load, <8 x i16> %passthrough
+  ret <8 x i16> %sel
+}
+
 !0 = !{!1, !1, i64 0}
 !1 = !{!"int", !2, i64 0}
 !2 = !{!"omnipotent char", !8, i64 0}
@@ -184,3 +224,4 @@ define <vscale x 4 x i32> @fold_sel_into_masked_load_drop_metadata(ptr %loc, <vs
 declare <8 x float> @llvm.masked.load.v8f32.p0(ptr, i32 immarg, <8 x i1>, <8 x float>)
 declare <4 x i32> @llvm.masked.load.v4i32.p0(ptr, i32 immarg, <4 x i1>, <4 x i32>)
 declare <4 x float> @llvm.masked.load.v4f32.p0(ptr, i32 immarg, <4 x i1>, <4 x float>)
+declare <8 x i16> @llvm.masked.load.v8i16.p0(ptr, <8 x i1>, <8 x i16>)

[nikic]

LGTM, thanks!

[nikic]

Remove the TODO. (AI risk mitigation.)

[Chennesxu]

Done, removed the TODO.

[nikic]

FWIW, I'm not sure this really makes sense. Both locations are not entirely correct here, but keeping the Load location will avoid jumpy debugging.

[Chennesxu]

Thanks for the review. I removed the explicit debug-location override.