[Instructions] cache computed shufflevector properties

Prince781 · Prince781 · commit 6a332ad40498 · 2024-11-12T00:44:56.000-08:00
Cache computed properties of a shufflevector mask.
diff --git a/llvm/include/llvm/IR/Instructions.h b/llvm/include/llvm/IR/Instructions.h
@@ -1848,6 +1848,39 @@ class ShuffleVectorInst : public Instruction {
 
   SmallVector<int, 4> ShuffleMask;
   Constant *ShuffleMaskForBitcode;
+  struct ShuffleProperties {
+    bool isSingleSource : 1;
+    bool isSingleSource_set : 1;
+    bool isIdentityMask : 1;
+    bool isIdentityMask_set : 1;
+    bool isIdentityWithPadding : 1;
+    bool isIdentityWithPadding_set : 1;
+    bool isIdentityWithExtract : 1;
+    bool isIdentityWithExtract_set : 1;
+    bool isConcat : 1;
+    bool isConcat_set : 1;
+    bool isSelect : 1;
+    bool isSelect_set : 1;
+    bool isReverse : 1;
+    bool isReverse_set : 1;
+    bool isZeroEltSplat : 1;
+    bool isZeroEltSplat_set : 1;
+    bool isTranspose : 1;
+    bool isTranspose_set : 1;
+
+    void unset() {
+      isSingleSource_set = false;
+      isIdentityMask_set = false;
+      isIdentityWithPadding_set = false;
+      isIdentityWithExtract_set = false;
+      isConcat_set = false;
+      isSelect_set = false;
+      isReverse_set = false;
+      isZeroEltSplat_set = false;
+      isTranspose_set = false;
+    }
+  };
+  mutable ShuffleProperties CachedShuffleProperties = {};
 
 protected:
   // Note: Instruction needs to be a friend here to call cloneImpl.
@@ -1959,8 +1992,13 @@ class ShuffleVectorInst : public Instruction {
   /// Example: shufflevector <4 x n> A, <4 x n> B, <3,0,undef,3>
   /// TODO: Optionally allow length-changing shuffles.
   bool isSingleSource() const {
-    return !changesLength() &&
-           isSingleSourceMask(ShuffleMask, ShuffleMask.size());
+    if (CachedShuffleProperties.isSingleSource_set)
+      return CachedShuffleProperties.isSingleSource;
+
+    CachedShuffleProperties.isSingleSource_set = true;
+    return CachedShuffleProperties.isSingleSource =
+               !changesLength() &&
+               isSingleSourceMask(ShuffleMask, ShuffleMask.size());
   }
 
   /// Return true if this shuffle mask chooses elements from exactly one source
@@ -1987,12 +2025,18 @@ class ShuffleVectorInst : public Instruction {
   /// from its input vectors.
   /// Example: shufflevector <4 x n> A, <4 x n> B, <4,undef,6,undef>
   bool isIdentity() const {
+    if (CachedShuffleProperties.isIdentityMask_set)
+      return CachedShuffleProperties.isIdentityMask;
+
+    CachedShuffleProperties.isIdentityMask_set = true;
     // Not possible to express a shuffle mask for a scalable vector for this
     // case.
     if (isa<ScalableVectorType>(getType()))
-      return false;
+      return CachedShuffleProperties.isIdentityMask = false;
 
-    return !changesLength() && isIdentityMask(ShuffleMask, ShuffleMask.size());
+    return CachedShuffleProperties.isIdentityMask =
+               !changesLength() &&
+               isIdentityMask(ShuffleMask, ShuffleMask.size());
   }
 
   /// Return true if this shuffle lengthens exactly one source vector with
@@ -2033,7 +2077,13 @@ class ShuffleVectorInst : public Instruction {
   /// In that case, the shuffle is better classified as an identity shuffle.
   /// TODO: Optionally allow length-changing shuffles.
   bool isSelect() const {
-    return !changesLength() && isSelectMask(ShuffleMask, ShuffleMask.size());
+    if (CachedShuffleProperties.isSelect_set)
+      return CachedShuffleProperties.isSelect;
+
+    CachedShuffleProperties.isSelect_set = true;
+    return CachedShuffleProperties.isSelect =
+               !changesLength() &&
+               isSelectMask(ShuffleMask, ShuffleMask.size());
   }
 
   /// Return true if this shuffle mask swaps the order of elements from exactly
@@ -2054,7 +2104,13 @@ class ShuffleVectorInst : public Instruction {
   /// Example: shufflevector <4 x n> A, <4 x n> B, <3,undef,1,undef>
   /// TODO: Optionally allow length-changing shuffles.
   bool isReverse() const {
-    return !changesLength() && isReverseMask(ShuffleMask, ShuffleMask.size());
+    if (CachedShuffleProperties.isReverse_set)
+      return CachedShuffleProperties.isReverse;
+
+    CachedShuffleProperties.isReverse_set = true;
+    return CachedShuffleProperties.isReverse =
+               !changesLength() &&
+               isReverseMask(ShuffleMask, ShuffleMask.size());
   }
 
   /// Return true if this shuffle mask chooses all elements with the same value
@@ -2077,8 +2133,13 @@ class ShuffleVectorInst : public Instruction {
   /// TODO: Optionally allow length-changing shuffles.
   /// TODO: Optionally allow splats from other elements.
   bool isZeroEltSplat() const {
-    return !changesLength() &&
-           isZeroEltSplatMask(ShuffleMask, ShuffleMask.size());
+    if (CachedShuffleProperties.isZeroEltSplat_set)
+      return CachedShuffleProperties.isZeroEltSplat;
+
+    CachedShuffleProperties.isZeroEltSplat_set = true;
+    return CachedShuffleProperties.isZeroEltSplat =
+               !changesLength() &&
+               isZeroEltSplatMask(ShuffleMask, ShuffleMask.size());
   }
 
   /// Return true if this shuffle mask is a transpose mask.
@@ -2127,7 +2188,13 @@ class ShuffleVectorInst : public Instruction {
   /// exact specification.
   /// Example: shufflevector <4 x n> A, <4 x n> B, <0,4,2,6>
   bool isTranspose() const {
-    return !changesLength() && isTransposeMask(ShuffleMask, ShuffleMask.size());
+    if (CachedShuffleProperties.isTranspose_set)
+      return CachedShuffleProperties.isTranspose;
+
+    CachedShuffleProperties.isTranspose_set = true;
+    return CachedShuffleProperties.isTranspose =
+               !changesLength() &&
+               isTransposeMask(ShuffleMask, ShuffleMask.size());
   }
 
   /// Return true if this shuffle mask is a splice mask, concatenating the two
diff --git a/llvm/lib/IR/Instructions.cpp b/llvm/lib/IR/Instructions.cpp
@@ -1811,6 +1811,7 @@ void ShuffleVectorInst::getShuffleMask(const Constant *Mask,
 }
 
 void ShuffleVectorInst::setShuffleMask(ArrayRef<int> Mask) {
+  CachedShuffleProperties.unset();
   ShuffleMask.assign(Mask.begin(), Mask.end());
   ShuffleMaskForBitcode = convertShuffleMaskForBitcode(Mask, getType());
 }
@@ -2101,63 +2102,77 @@ bool ShuffleVectorInst::isInsertSubvectorMask(ArrayRef<int> Mask,
 }
 
 bool ShuffleVectorInst::isIdentityWithPadding() const {
+  if (CachedShuffleProperties.isIdentityWithPadding_set)
+    return CachedShuffleProperties.isIdentityWithPadding;
+
+  CachedShuffleProperties.isIdentityWithPadding_set = true;
   // FIXME: Not currently possible to express a shuffle mask for a scalable
   // vector for this case.
   if (isa<ScalableVectorType>(getType()))
-    return false;
+    return CachedShuffleProperties.isIdentityWithPadding = false;
 
   int NumOpElts = cast<FixedVectorType>(Op<0>()->getType())->getNumElements();
   int NumMaskElts = cast<FixedVectorType>(getType())->getNumElements();
   if (NumMaskElts <= NumOpElts)
-    return false;
+    return CachedShuffleProperties.isIdentityWithPadding = false;
 
   // The first part of the mask must choose elements from exactly 1 source op.
   ArrayRef<int> Mask = getShuffleMask();
   if (!isIdentityMaskImpl(Mask, NumOpElts))
-    return false;
+    return CachedShuffleProperties.isIdentityWithPadding = false;
 
   // All extending must be with undef elements.
   for (int i = NumOpElts; i < NumMaskElts; ++i)
     if (Mask[i] != -1)
-      return false;
+      return CachedShuffleProperties.isIdentityWithPadding = false;
 
-  return true;
+  return CachedShuffleProperties.isIdentityWithPadding = true;
 }
 
 bool ShuffleVectorInst::isIdentityWithExtract() const {
+  if (CachedShuffleProperties.isIdentityWithExtract_set)
+    return CachedShuffleProperties.isIdentityWithExtract;
+
+  CachedShuffleProperties.isIdentityWithExtract_set = true;
   // FIXME: Not currently possible to express a shuffle mask for a scalable
   // vector for this case.
   if (isa<ScalableVectorType>(getType()))
-    return false;
+    return CachedShuffleProperties.isIdentityWithExtract = false;
 
   int NumOpElts = cast<FixedVectorType>(Op<0>()->getType())->getNumElements();
   int NumMaskElts = cast<FixedVectorType>(getType())->getNumElements();
   if (NumMaskElts >= NumOpElts)
-    return false;
+    return CachedShuffleProperties.isIdentityWithExtract = false;
 
-  return isIdentityMaskImpl(getShuffleMask(), NumOpElts);
+  return CachedShuffleProperties.isIdentityWithExtract =
+             isIdentityMaskImpl(getShuffleMask(), NumOpElts);
 }
 
 bool ShuffleVectorInst::isConcat() const {
+  if (CachedShuffleProperties.isConcat_set)
+    return CachedShuffleProperties.isConcat;
+
+  CachedShuffleProperties.isConcat_set = true;
   // Vector concatenation is differentiated from identity with padding.
   if (isa<UndefValue>(Op<0>()) || isa<UndefValue>(Op<1>()))
-    return false;
+    return CachedShuffleProperties.isConcat = false;
 
   // FIXME: Not currently possible to express a shuffle mask for a scalable
   // vector for this case.
   if (isa<ScalableVectorType>(getType()))
-    return false;
+    return CachedShuffleProperties.isConcat = false;
 
   int NumOpElts = cast<FixedVectorType>(Op<0>()->getType())->getNumElements();
   int NumMaskElts = cast<FixedVectorType>(getType())->getNumElements();
   if (NumMaskElts != NumOpElts * 2)
-    return false;
+    return CachedShuffleProperties.isConcat = false;
 
   // Use the mask length rather than the operands' vector lengths here. We
   // already know that the shuffle returns a vector twice as long as the inputs,
   // and neither of the inputs are undef vectors. If the mask picks consecutive
   // elements from both inputs, then this is a concatenation of the inputs.
-  return isIdentityMaskImpl(getShuffleMask(), NumMaskElts);
+  return CachedShuffleProperties.isConcat =
+             isIdentityMaskImpl(getShuffleMask(), NumMaskElts);
 }
 
 static bool isReplicationMaskWithParams(ArrayRef<int> Mask,
