[SVE] Bail from VectorUtils heuristics for scalable vectors

[ctetreau]

Summary

Bail from maskIsAllZeroOrUndef and maskIsAllOneOrUndef prior to iterating over the number of
elements for scalable vectors.

Assert that the mask type is not scalable in possiblyDemandedEltsInMask .

Assert that the types are correct in all three functions.

[efriedma]

llvm/lib/Analysis/VectorUtils.cpp

920–925

I think I'd prefer to handle this in the callers: if the vector is scalable, just don't call it. Returning an Optional<> is more confusing, I think, since it can't fail for fixed vectors.

[efriedma]

llvm/include/llvm/Analysis/VectorUtils.h

552–553

If the input isn't i1, it isn't obvious what "true"/"false" means here.

[ctetreau]

llvm/include/llvm/Analysis/VectorUtils.h

552–553

I guess these functions make and don't check lots of assumptions. (element type is bool, type is not just int, etc...) I'm just trying to follow the boy scout code here and leave the code cleaner than I found it. It would be better if I just asserted that the mask type is a vector and the element type is a bool, then the function would actually do what it says on the tin.

llvm/lib/Analysis/VectorUtils.cpp

920–925

I can change it to assert for scalable vectors. If we ever add a notion of a scalable demanded elements (we do this thing with bitfields in lots of places), then it's possible that this function could return a value for scalable vectors. Though in that event the return type of this function would change anyways, so we might as well wait.

[ctetreau]

address code review issues

[ctetreau]

fix now unnecessary assignment to optional<APInt>

[ctetreau]

remove unneeded include

[efriedma]

LGTM

llvm/include/llvm/Analysis/VectorUtils.h

 //===- llvm/Analysis/VectorUtils.h - Vector utilities -----------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines some vectorizer utilities.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_VECTORUTILS_H
 #define LLVM_ANALYSIS_VECTORUTILS_H
 
 #include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Analysis/LoopAccessAnalysis.h"
 #include "llvm/Support/CheckedArithmetic.h"
 
 namespace llvm {
 class TargetLibraryInfo;
 
 /// Describes the type of Parameters

 ///
 /// This function generates code that concatenate the vectors in \p Vecs into a
 /// single large vector. The number of vectors should be greater than one, and
 /// their element types should be the same. The number of elements in the
 /// vectors should also be the same; however, if the last vector has fewer
 /// elements, it will be padded with undefs.
 Value *concatenateVectors(IRBuilderBase &Builder, ArrayRef<Value *> Vecs);
 
-/// Given a mask vector of the form <Y x i1>, Return true if all of the
-/// elements of this predicate mask are false or undef.  That is, return true
-/// if all lanes can be assumed inactive. 
+/// Given a mask vector of i1, Return true if all of the elements of this
+/// predicate mask are known to be false or undef.  That is, return true if all
+/// lanes can be assumed inactive.
 bool maskIsAllZeroOrUndef(Value *Mask);
 
-/// Given a mask vector of the form <Y x i1>, Return true if all of the
-/// elements of this predicate mask are true or undef.  That is, return true
-/// if all lanes can be assumed active. 
+/// Given a mask vector of i1, Return true if all of the elements of this
+/// predicate mask are known to be true or undef.  That is, return true if all
+/// lanes can be assumed active.
 bool maskIsAllOneOrUndef(Value *Mask);
 
 /// Given a mask vector of the form <Y x i1>, return an APInt (of bitwidth Y)
 /// for each lane which may be active.
 APInt possiblyDemandedEltsInMask(Value *Mask);
-  
 /// The group of interleaved loads/stores sharing the same stride and
 /// close to each other.
 ///
 /// Each member in this group has an index starting from 0, and the largest
 /// index should be less than interleaved factor, which is equal to the absolute
 /// value of the access's stride.
 ///
 /// E.g. An interleaved load group of factor 4:

llvm/lib/Analysis/VectorUtils.cpp

do {
     ResList = TmpList;
     NumVecs = ResList.size();
   } while (NumVecs > 1);
 
   return ResList[0];
 }
 
 bool llvm::maskIsAllZeroOrUndef(Value *Mask) {
+  assert(isa<VectorType>(Mask->getType()) &&
+         isa<IntegerType>(Mask->getType()->getScalarType()) &&
+         cast<IntegerType>(Mask->getType()->getScalarType())->getBitWidth() ==
+             1 &&
+         "Mask must be a vector of i1");
+
   auto *ConstMask = dyn_cast<Constant>(Mask);
   if (!ConstMask)
     return false;
   if (ConstMask->isNullValue() || isa<UndefValue>(ConstMask))
     return true;
+  if (isa<ScalableVectorType>(ConstMask->getType()))
+    return false;
   for (unsigned
            I = 0,
            E = cast<FixedVectorType>(ConstMask->getType())->getNumElements();
        I != E; ++I) {
     if (auto *MaskElt = ConstMask->getAggregateElement(I))
       if (MaskElt->isNullValue() || isa<UndefValue>(MaskElt))
         continue;
     return false;
   }
   return true;
 }
 
 
 bool llvm::maskIsAllOneOrUndef(Value *Mask) {
+  assert(isa<VectorType>(Mask->getType()) &&
+         isa<IntegerType>(Mask->getType()->getScalarType()) &&
+         cast<IntegerType>(Mask->getType()->getScalarType())->getBitWidth() ==
+             1 &&
+         "Mask must be a vector of i1");
+
   auto *ConstMask = dyn_cast<Constant>(Mask);
   if (!ConstMask)
     return false;
   if (ConstMask->isAllOnesValue() || isa<UndefValue>(ConstMask))
     return true;
+  if (isa<ScalableVectorType>(ConstMask->getType()))
+    return false;
   for (unsigned
            I = 0,
            E = cast<FixedVectorType>(ConstMask->getType())->getNumElements();
        I != E; ++I) {
     if (auto *MaskElt = ConstMask->getAggregateElement(I))
       if (MaskElt->isAllOnesValue() || isa<UndefValue>(MaskElt))
         continue;
     return false;
   }
   return true;
 }
 
 /// TODO: This is a lot like known bits, but for
 /// vectors.  Is there something we can common this with?
 APInt llvm::possiblyDemandedEltsInMask(Value *Mask) {
+  assert(isa<FixedVectorType>(Mask->getType()) &&
+         isa<IntegerType>(Mask->getType()->getScalarType()) &&
+         cast<IntegerType>(Mask->getType()->getScalarType())->getBitWidth() ==
+             1 &&
+         "Mask must be a fixed width vector of i1");
 
   const unsigned VWidth =
       cast<FixedVectorType>(Mask->getType())->getNumElements();
   APInt DemandedElts = APInt::getAllOnesValue(VWidth);
   if (auto *CV = dyn_cast<ConstantVector>(Mask))
     for (unsigned i = 0; i < VWidth; i++)
       if (CV->getAggregateElement(i)->isNullValue())
         DemandedElts.clearBit(i);

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

Instruction *InstCombinerImpl::simplifyMaskedStore(IntrinsicInst &II) {
   // If the mask is all ones, this is a plain vector store of the 1st argument.
   if (ConstMask->isAllOnesValue()) {
     Value *StorePtr = II.getArgOperand(1);
     Align Alignment = cast<ConstantInt>(II.getArgOperand(2))->getAlignValue();
     return new StoreInst(II.getArgOperand(0), StorePtr, false, Alignment);
   }
 
   // Use masked off lanes to simplify operands via SimplifyDemandedVectorElts
-  APInt DemandedElts = possiblyDemandedEltsInMask(ConstMask);
-  APInt UndefElts(DemandedElts.getBitWidth(), 0);
-  if (Value *V = SimplifyDemandedVectorElts(II.getOperand(0),
-                                            DemandedElts, UndefElts))
+  Optional<APInt> DemandedElts = possiblyDemandedEltsInMask(ConstMask);
+  if (DemandedElts) {
+    APInt UndefElts(DemandedElts->getBitWidth(), 0);
+    if (Value *V = SimplifyDemandedVectorElts(II.getOperand(0), *DemandedElts,
+                                              UndefElts))
       return replaceOperand(II, 0, V);
+  }
 
   return nullptr;
 }
 
 // TODO, Obvious Missing Transforms:
 // * Single constant active lane load -> load
 // * Dereferenceable address & few lanes -> scalarize speculative load/selects
 // * Adjacent vector addresses -> masked.load
Instruction *InstCombinerImpl::simplifyMaskedScatter(IntrinsicInst &II) {
   if (!ConstMask)
     return nullptr;
 
   // If the mask is all zeros, a scatter does nothing.
   if (ConstMask->isNullValue())
     return eraseInstFromFunction(II);
 
   // Use masked off lanes to simplify operands via SimplifyDemandedVectorElts
-  APInt DemandedElts = possiblyDemandedEltsInMask(ConstMask);
-  APInt UndefElts(DemandedElts.getBitWidth(), 0);
-  if (Value *V = SimplifyDemandedVectorElts(II.getOperand(0),
-                                            DemandedElts, UndefElts))
+  Optional<APInt> DemandedElts = possiblyDemandedEltsInMask(ConstMask);
+  if (DemandedElts) {
+    APInt UndefElts(DemandedElts->getBitWidth(), 0);
+    if (Value *V = SimplifyDemandedVectorElts(II.getOperand(0), *DemandedElts,
+                                              UndefElts))
       return replaceOperand(II, 0, V);
-  if (Value *V = SimplifyDemandedVectorElts(II.getOperand(1),
-                                            DemandedElts, UndefElts))
+    if (Value *V = SimplifyDemandedVectorElts(II.getOperand(1), *DemandedElts,
+                                              UndefElts))
       return replaceOperand(II, 1, V);
+  }
 
   return nullptr;
 }
 
 /// This function transforms launder.invariant.group and strip.invariant.group
 /// like:
 /// launder(launder(%x)) -> launder(%x)       (the result is not the argument)
 /// launder(strip(%x)) -> launder(%x)

llvm/test/Transforms/InstCombine/AArch64/VectorUtils_heuristics.ll

+; RUN: opt -S -instcombine < %s | FileCheck %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+target triple = "aarch64-unknown-linux-gnu"
+
+; This test checks that instcombine does not crash while invoking
+; maskIsAllOneOrUndef, maskIsAllZeroOrUndef, or possiblyDemandedEltsInMask.
+
+; CHECK-LABEL: novel_algorithm
+; CHECK: unreachable
+define void @novel_algorithm() {
+entry:
+  %a = call <vscale x 16 x i8> @llvm.masked.load.nxv16i8.p0nxv16i8(<vscale x 16 x i8>* undef, i32 1, <vscale x 16 x i1> shufflevector (<vscale x 16 x i1> insertelement (<vscale x 16 x i1> undef, i1 true, i32 0), <vscale x 16 x i1> undef, <vscale x 16 x i32> zeroinitializer), <vscale x 16 x i8> undef)
+  %b = add <vscale x 16 x i8> undef, %a
+  call void @llvm.masked.store.nxv16i8.p0nxv16i8(<vscale x 16 x i8> %b, <vscale x 16 x i8>* undef, i32 1, <vscale x 16 x i1> shufflevector (<vscale x 16 x i1> insertelement (<vscale x 16 x i1> undef, i1 true, i32 0), <vscale x 16 x i1> undef, <vscale x 16 x i32> zeroinitializer))
+  unreachable
+}
+
+declare <vscale x 16 x i8> @llvm.masked.load.nxv16i8.p0nxv16i8(<vscale x 16 x i8>*, i32 immarg, <vscale x 16 x i1>, <vscale x 16 x i8>)
+
+declare void @llvm.masked.store.nxv16i8.p0nxv16i8(<vscale x 16 x i8>, <vscale x 16 x i8>*, i32 immarg, <vscale x 16 x i1>)