[llvm] 926ccfe - [SLP] ScalarizationOverheadBuilder - demand all elements for scalarization if the extraction index is unknown / out of bounds

[Simon Pilgrim via llvm-commits]

Author: Simon Pilgrim
Date: 2022-09-28T11:03:37+01:00
New Revision: 926ccfef032d206dcbcdf74ca1e3a9ebf4d1be45

URL: https://github.com/llvm/llvm-project/commit/926ccfef032d206dcbcdf74ca1e3a9ebf4d1be45
DIFF: https://github.com/llvm/llvm-project/commit/926ccfef032d206dcbcdf74ca1e3a9ebf4d1be45.diff

LOG: [SLP] ScalarizationOverheadBuilder - demand all elements for scalarization if the extraction index is unknown / out of bounds

Workaround for a chromium bug reported on D134605 - test case will be added later

Added: 
    

Modified: 
    llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
index 9584c421bcca..167858edc40c 100644
--- a/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp
@@ -5805,24 +5805,32 @@ class ScalarizationOverheadBuilder {
 public:
   /// Add an extraction from a specific source and element index.
   void addExtract(Value *Src, unsigned Idx) {
+    auto *Ty = cast<FixedVectorType>(Src->getType());
+    unsigned NumElts = Ty->getNumElements();
     if (m_ExtractsByClass.count(Src)) {
-      m_ExtractsByClass[Src].setBit(Idx);
+      if (Idx < NumElts)
+        m_ExtractsByClass[Src].setBit(Idx);
+      else
+        m_ExtractsByClass[Src].setAllBits();
       return;
     }
-    auto *Ty = cast<FixedVectorType>(Src->getType());
-    unsigned NumElts = Ty->getNumElements();
-    m_ExtractsByClass[Src] = APInt::getOneBitSet(NumElts, Idx);
+    m_ExtractsByClass[Src] = Idx < NumElts ? APInt::getOneBitSet(NumElts, Idx)
+                                           : APInt::getAllOnes(NumElts);
   }
 
   /// Add an extraction from a vector type and specific element index.
   /// We assume that all extractions from a given type are from the same source.
   void addExtract(FixedVectorType *VecTy, unsigned Idx) {
+    unsigned NumElts = VecTy->getNumElements();
     if (m_ExtractsByType.count(VecTy)) {
-      m_ExtractsByType[VecTy].setBit(Idx);
+      if (Idx < NumElts)
+        m_ExtractsByType[VecTy].setBit(Idx);
+      else
+        m_ExtractsByType[VecTy].setAllBits();
       return;
     }
-    unsigned NumElts = VecTy->getNumElements();
-    m_ExtractsByType[VecTy] = APInt::getOneBitSet(NumElts, Idx);
+    m_ExtractsByType[VecTy] = Idx < NumElts ? APInt::getOneBitSet(NumElts, Idx)
+                                            : APInt::getAllOnes(NumElts);
   }
 
   /// Add an extended extraction from a specific source and element index.