[BasicTTI] getInterleavedMemoryOpCost(): discount unused members of mask if mask for gap will be used

[lebedev.ri]

Summary

As it can be seen in InnerLoopVectorizer::vectorizeInterleaveGroup(),
in some cases (reported by UseMaskForGaps), the gaps in the interleaved load/store group
will be masked away by another constant mask, so there is no need to
account for the cost of replication of the mask for these.

[lebedev.ri]

ping

[RKSimon]

LGTM (by inspection) - I tend to prefer to use the getScalarizationOverhead wrapper when we're using all the elements of the type, but I think it's better for consistency here to always specify the demanded elts.

Ideally we'd have at least some test coverage for this, but I understand mask gaps codegen can be tricky.

llvm/include/llvm/CodeGen/BasicTTIImpl.h

1241

Worth pulling these out?

const APInt DemandedAllSubElts = APInt::getAllOnes(NumSubElts);
const APInt DemandedAllResultElts = APInt::getAllOnes(NumElts);

[lebedev.ri]

@RKSimon thank you for the review!
Applied nit suggestion.

Yeah, it would indeed be really great to have test coverage for this.

[lebedev.ri]

And we have a winner!
Test coverage added, will land immediately.

[RKSimon]

Thanks!

llvm/include/llvm/CodeGen/BasicTTIImpl.h

if (Cost.isValid() && VecTySize > VecTyLTSize) {
       // will be used.
       Cost = divideCeil(UsedInsts.count() * Cost.getValue().getValue(),
                         NumLegalInsts);
     }
 
     // Then plus the cost of interleave operation.
     assert(Indices.size() <= Factor &&
            "Interleaved memory op has too many members");
 
     APInt DemandedLoadStoreElts = APInt::getZero(NumElts);
     for (unsigned Index : Indices) {
       assert(Index < Factor && "Invalid index for interleaved memory op");
       for (unsigned Elm = 0; Elm < NumSubElts; Elm++)
         DemandedLoadStoreElts.setBit(Index + Elm * Factor);
     }
 
     if (Opcode == Instruction::Load) {
       // The interleave cost is similar to extract sub vectors' elements
       // from the wide vector, and insert them into sub vectors.
       //
       // E.g. An interleaved load of factor 2 (with one member of index 0):
       //      %vec = load <8 x i32>, <8 x i32>* %ptr
       //      %v0 = shuffle %vec, undef, <0, 2, 4, 6>         ; Index 0
       // The cost is estimated as extract elements at 0, 2, 4, 6 from the
       // <8 x i32> vector and insert them into a <4 x i32> vector.
-      InstructionCost InsSubCost =
-          getScalarizationOverhead(SubVT, /*Insert*/ true, /*Extract*/ false);
+      InstructionCost InsSubCost = thisT()->getScalarizationOverhead(
+          SubVT, APInt::getAllOnes(NumSubElts),
+          /*Insert*/ true, /*Extract*/ false);
       Cost += Indices.size() * InsSubCost;
       Cost +=
           thisT()->getScalarizationOverhead(VT, DemandedLoadStoreElts,
                                             /*Insert*/ false, /*Extract*/ true);
     } else {
       // The interleave cost is extract elements from sub vectors, and
       // insert them into the wide vector.
       //
       // E.g. An interleaved store of factor 3 with 2 members at indices 0,1:
       // (using VF=4):
       //    %v0_v1 = shuffle %v0, %v1, <0,4,undef,1,5,undef,2,6,undef,3,7,undef>
       //    %gaps.mask = <true, true, false, true, true, false,
       //                  true, true, false, true, true, false>
       //    call llvm.masked.store <12 x i32> %v0_v1, <12 x i32>* %ptr,
       //                           i32 Align, <12 x i1> %gaps.mask
       // The cost is estimated as extract all elements (of actual members,
       // excluding gaps) from both <4 x i32> vectors and insert into the <12 x
       // i32> vector.
-      InstructionCost ExtSubCost =
-          getScalarizationOverhead(SubVT, /*Insert*/ false, /*Extract*/ true);
+      InstructionCost ExtSubCost = thisT()->getScalarizationOverhead(
+          SubVT, APInt::getAllOnes(NumSubElts),
+          /*Insert*/ false, /*Extract*/ true);
       Cost += ExtSubCost * Indices.size();
       Cost += thisT()->getScalarizationOverhead(VT, DemandedLoadStoreElts,
                                                 /*Insert*/ true,
                                                 /*Extract*/ false);
     }
 
     if (!UseMaskForCond)
       return Cost;
 
     Type *I8Type = Type::getInt8Ty(VT->getContext());
     auto *MaskVT = FixedVectorType::get(I8Type, NumElts);
     SubVT = FixedVectorType::get(I8Type, NumSubElts);
 
     // The Mask shuffling cost is extract all the elements of the Mask
     // and insert each of them Factor times into the wide vector:
     //
     // E.g. an interleaved group with factor 3:
     //    %mask = icmp ult <8 x i32> %vec1, %vec2
     //    %interleaved.mask = shufflevector <8 x i1> %mask, <8 x i1> undef,
     //        <24 x i32> <0,0,0,1,1,1,2,2,2,3,3,3,4,4,4,5,5,5,6,6,6,7,7,7>
     // The cost is estimated as extract all mask elements from the <8xi1> mask
     // vector and insert them factor times into the <24xi1> shuffled mask
     // vector.
-    Cost += getScalarizationOverhead(SubVT, /*Insert*/ false, /*Extract*/ true);
     Cost +=
-        getScalarizationOverhead(MaskVT, /*Insert*/ true, /*Extract*/ false);
+        thisT()->getScalarizationOverhead(SubVT, APInt::getAllOnes(NumSubElts),
+                                          /*Insert*/ false, /*Extract*/ true);
+    Cost += thisT()->getScalarizationOverhead(
+        MaskVT,
+        UseMaskForGaps ? DemandedLoadStoreElts : APInt::getAllOnes(NumElts),
+        /*Insert*/ true, /*Extract*/ false);
 
     // The Gaps mask is invariant and created outside the loop, therefore the
     // cost of creating it is not accounted for here. However if we have both
     // a MaskForGaps and some other mask that guards the execution of the
     // memory access, we need to account for the cost of And-ing the two masks
     // inside the loop.
     if (UseMaskForGaps)
       Cost += thisT()->getArithmeticInstrCost(BinaryOperator::And, MaskVT,