[InstCombine] Fixed select + masked load fold failure

[DylanFleming-arm]

Summary

Fixed type assertion failure caused by trying to fold a masked load with a
select where the select condition is a scalar value

[sdesmalen]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

3234

nit: unnecessary parentheses

[lebedev.ri]

LG

[srj]

I can confirm that this fixes the known issue in Halide. Thanks!

[srj]

(I presume this will be backported to the release/13.x branch as well as main?)

[DylanFleming-arm]

Yeah, I think it definitely makes sense to have it put into the LLVM13 release

[paulwalker-arm]

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

3234

Not a big issue but looking at the definition of select I think this can just be isa<VectorType>(CondVal->getType) and having that first in the if's condition block would prevent unnecessary calls to match and thus improve compile time by the odd nanosecond.

llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp

if (match(FalseVal, m_Zero()) &&
     if (isa<UndefValue>(MaskedLoad->getArgOperand(3)))
       MaskedLoad->setArgOperand(3, FalseVal /* Zero */);
     return replaceInstUsesWith(SI, MaskedLoad);
   }
 
   Value *Mask;
   if (match(TrueVal, m_Zero()) &&
       match(FalseVal, m_MaskedLoad(m_Value(), m_Value(), m_Value(Mask),
-                                   m_CombineOr(m_Undef(), m_Zero())))) {
+                                   m_CombineOr(m_Undef(), m_Zero()))) &&
+      (CondVal->getType() == Mask->getType())) {
     // We can remove the select by ensuring the load zeros all lanes the
     // select would have.  We determine this by proving there is no overlap
     // between the load and select masks.
     // (i.e (load_mask & select_mask) == 0 == no overlap)
     bool CanMergeSelectIntoLoad = false;
     if (Value *V = SimplifyAndInst(CondVal, Mask, SQ.getWithInstruction(&SI)))
       CanMergeSelectIntoLoad = match(V, m_Zero());
 

llvm/test/Transforms/InstCombine/select-masked_load.ll

 ; CHECK-NEXT: ret <4 x float> %load
   %mask = xor <4 x i1> %inv_mask, <i1 true, i1 true, i1 true, i1 true>
   %pg = and <4 x i1> %mask, %cond
   %load = call <4 x float> @llvm.masked.load.v4f32.p0v4f32(<4 x float>* %ptr, i32 4, <4 x i1> %pg, <4 x float> undef)
   %masked = select <4 x i1> %inv_mask, <4 x float> zeroinitializer, <4 x float> %load
   ret <4 x float> %masked
 }
 
+define <8 x float> @masked_load_and_scalar_select_cond(<8 x float>* %ptr, <8 x i1> %mask, i1 %cond) {
+; CHECK-LABEL: @masked_load_and_scalar_select_cond(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = call <8 x float> @llvm.masked.load.v8f32.p0v8f32(<8 x float>* [[PTR:%.*]], i32 32, <8 x i1> [[MASK:%.*]], <8 x float> undef)
+; CHECK-NEXT:    [[TMP1:%.*]] = select i1 [[COND:%.*]], <8 x float> zeroinitializer, <8 x float> [[TMP0]]
+; CHECK-NEXT:    ret <8 x float> [[TMP1]]
+entry:
+  %0 = call <8 x float> @llvm.masked.load.v8f32.p0v8f32(<8 x float>* %ptr, i32 32, <8 x i1> %mask, <8 x float> undef)
+  %1 = select i1 %cond, <8 x float> zeroinitializer, <8 x float> %0
+  ret <8 x float> %1
+}
+
+declare <8 x float> @llvm.masked.load.v8f32.p0v8f32(<8 x float>*, i32 immarg, <8 x i1>, <8 x float>)
 declare <4 x i32> @llvm.masked.load.v4i32.p0v4i32(<4 x i32>*, i32 immarg, <4 x i1>, <4 x i32>)
 declare <4 x float> @llvm.masked.load.v4f32.p0v4f32(<4 x float>*, i32 immarg, <4 x i1>, <4 x float>)