Merge pull request #65119 from slavapestov/sil-optimizer-peepholes-for-packs

Simple SILOptimizer peepholes for packs

Author: slavapestov

lib/AST/SubstitutionMap.cpp

@@ -31,6 +31,7 @@
 #include "swift/AST/InFlightSubstitution.h"
 #include "swift/AST/LazyResolver.h"
 #include "swift/AST/Module.h"
+#include "swift/AST/PackConformance.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/TypeCheckRequests.h"
 #include "swift/AST/Types.h"
@@ -792,13 +793,29 @@ bool SubstitutionMap::isIdentity() const {
   if (empty())
     return true;
 
+  for (auto conf : getConformances()) {
+    if (conf.isAbstract())
+      continue;
+
+    if (conf.isPack()) {
+      auto patternConfs = conf.getPack()->getPatternConformances();
+      if (patternConfs.size() == 1 && patternConfs[0].isAbstract())
+        continue;
+    }
+
+    return false;
+  }
+
   GenericSignature sig = getGenericSignature();
   bool hasNonIdentityReplacement = false;
   auto replacements = getReplacementTypesBuffer();
 
   sig->forEachParam([&](GenericTypeParamType *paramTy, bool isCanonical) {
     if (isCanonical) {
-      if (!paramTy->isEqual(replacements[0]))
+      Type wrappedParamTy = paramTy;
+      if (paramTy->isParameterPack())
+        wrappedParamTy = PackType::getSingletonPackExpansion(paramTy);
+      if (!wrappedParamTy->isEqual(replacements[0]))
         hasNonIdentityReplacement = true;
     }
 

lib/IRGen/Fulfillment.cpp

@@ -188,6 +188,8 @@ static CanType getSingletonPackExpansionParameter(CanPackType packType,
       return expansion.getPatternType();
     }
   }
+
+  return CanType();
 }
 
 bool FulfillmentMap::searchTypeMetadataPack(IRGenModule &IGM,

lib/IRGen/GenReflection.cpp

@@ -490,10 +490,8 @@ IRGenModule::emitWitnessTableRefString(CanType type,
             type = genericEnv->mapTypeIntoContext(type)->getCanonicalType();
           }
           if (origType->hasTypeParameter()) {
-            auto origSig = genericEnv->getGenericSignature();
             conformance = conformance.subst(origType,
-              QueryInterfaceTypeSubstitutions(genericEnv),
-              LookUpConformanceInSignature(origSig.getPointer()));
+                genericEnv->getForwardingSubstitutionMap());
           }
           auto ret = emitWitnessTableRef(IGF, type, conformance);
           IGF.Builder.CreateRet(ret);

lib/SIL/IR/SILInstruction.cpp

@@ -881,6 +881,22 @@ namespace {
       return true;
     }
 
+    bool visitScalarPackIndexInst(const ScalarPackIndexInst *RHS) {
+      auto *X = cast<ScalarPackIndexInst>(LHS);
+      return (X->getIndexedPackType() == RHS->getIndexedPackType() &&
+              X->getComponentIndex() == RHS->getComponentIndex());
+    }
+
+    bool visitDynamicPackIndexInst(const DynamicPackIndexInst *RHS) {
+      auto *X = cast<DynamicPackIndexInst>(LHS);
+      return X->getIndexedPackType() == RHS->getIndexedPackType();
+    }
+
+    bool visitTuplePackElementAddrInst(const TuplePackElementAddrInst *RHS) {
+      auto *X = cast<TuplePackElementAddrInst>(LHS);
+      return X->getElementType() == RHS->getElementType();
+    }
+
   private:
     const SILInstruction *LHS;
   };

lib/SILOptimizer/LoopTransforms/LoopUnroll.cpp

@@ -163,11 +163,21 @@ static Optional<uint64_t> getMaxLoopTripCount(SILLoop *Loop,
 
   SILValue RecNext = Cmp->getArguments()[0];
   SILPhiArgument *RecArg;
+
+  // Match signed add with overflow, unsigned add with overflow and
+  // add without overflow.
   if (!match(RecNext, m_TupleExtractOperation(
                           m_ApplyInst(BuiltinValueKind::SAddOver,
                                       m_SILPhiArgument(RecArg), m_One()),
-                          0)))
+                          0)) &&
+      !match(RecNext, m_TupleExtractOperation(
+                          m_ApplyInst(BuiltinValueKind::UAddOver,
+                                      m_SILPhiArgument(RecArg), m_One()),
+                          0)) &&
+      !match(RecNext, m_ApplyInst(BuiltinValueKind::Add,
+                                      m_SILPhiArgument(RecArg), m_One()))) {
     return None;
+  }
 
   if (RecArg->getParent() != Header)
     return None;

lib/SILOptimizer/SILCombiner/SILCombiner.h

@@ -300,6 +300,11 @@ class SILCombiner :
   SILInstruction *
   visitDifferentiableFunctionExtractInst(DifferentiableFunctionExtractInst *DFEI);
 
+  SILInstruction *visitPackLengthInst(PackLengthInst *PLI);
+  SILInstruction *visitPackElementGetInst(PackElementGetInst *PEGI);
+  SILInstruction *visitTuplePackElementAddrInst(TuplePackElementAddrInst *TPEAI);
+  SILInstruction *visitCopyAddrInst(CopyAddrInst *CAI);
+
   SILInstruction *legacyVisitGlobalValueInst(GlobalValueInst *globalValue);
 
 #define PASS(ID, TAG, DESCRIPTION)

lib/SILOptimizer/SILCombiner/SILCombinerMiscVisitors.cpp

@@ -2391,3 +2391,136 @@ SILCombiner::visitDifferentiableFunctionExtractInst(DifferentiableFunctionExtrac
   replaceInstUsesWith(*DFEI, newValue);
   return eraseInstFromFunction(*DFEI);
 }
+
+// Simplify `pack_length` with constant-length pack.
+//
+// Before:
+// %len = pack_length $Pack{Int, String, Float}
+//
+// After:
+// %len = integer_literal Builtin.Word, 3
+SILInstruction *SILCombiner::visitPackLengthInst(PackLengthInst *PLI) {
+  auto PackTy = PLI->getPackType();
+  if (!PackTy->containsPackExpansionType()) {
+    return Builder.createIntegerLiteral(PLI->getLoc(), PLI->getType(),
+                                        PackTy->getNumElements());
+  }
+
+  return nullptr;
+}
+
+// Simplify `pack_element_get` where the index is a `dynamic_pack_index` with
+// a constant operand.
+//
+// Before:
+// %idx = integer_literal Builtin.Word, N
+// %pack_idx = dynamic_pack_index %Pack{Int, String, Float}, %idx
+// %pack_elt = pack_element_get %pack_value, %pack_idx, @element("...")
+//
+// After:
+// %pack_idx = scalar_pack_index %Pack{Int, String, Float}, N
+// %concrete_elt = pack_element_get %pack_value, %pack_idx, <<concrete type>>
+// %pack_elt = unchecked_addr_cast %concrete_elt, @element("...")
+SILInstruction *SILCombiner::visitPackElementGetInst(PackElementGetInst *PEGI) {
+  auto *DPII = dyn_cast<DynamicPackIndexInst>(PEGI->getIndex());
+  if (DPII == nullptr)
+    return nullptr;
+
+  auto PackTy = PEGI->getPackType();
+  if (PackTy->containsPackExpansionType())
+    return nullptr;
+
+  auto *Op = dyn_cast<IntegerLiteralInst>(DPII->getOperand());
+  if (Op == nullptr)
+    return nullptr;
+
+  if (Op->getValue().uge(PackTy->getNumElements()))
+    return nullptr;
+
+  unsigned Index = Op->getValue().getZExtValue();
+  auto *SPII = Builder.createScalarPackIndex(
+      DPII->getLoc(), Index, DPII->getIndexedPackType());
+
+  auto ElementTy = SILType::getPrimitiveAddressType(
+      PEGI->getPackType().getElementType(Index));
+  auto *NewPEGI = Builder.createPackElementGet(
+      PEGI->getLoc(), SPII, PEGI->getPack(),
+      ElementTy);
+
+  return Builder.createUncheckedAddrCast(
+      PEGI->getLoc(), NewPEGI, PEGI->getElementType());
+}
+
+// Simplify `tuple_pack_element_addr` where the index is a `dynamic_pack_index`
+//with a constant operand.
+//
+// Before:
+// %idx = integer_literal Builtin.Word, N
+// %pack_idx = dynamic_pack_index %Pack{Int, String, Float}, %idx
+// %tuple_elt = tuple_pack_element_addr %tuple_value, %pack_idx, @element("...")
+//
+// After:
+// %concrete_elt = tuple_element_addr %tuple_value, N
+// %tuple_elt = unchecked_addr_cast %concrete_elt, @element("...")
+SILInstruction *
+SILCombiner::visitTuplePackElementAddrInst(TuplePackElementAddrInst *TPEAI) {
+  auto *DPII = dyn_cast<DynamicPackIndexInst>(TPEAI->getIndex());
+  if (DPII == nullptr)
+    return nullptr;
+
+  auto PackTy = DPII->getIndexedPackType();
+  if (PackTy->containsPackExpansionType())
+    return nullptr;
+
+  auto *Op = dyn_cast<IntegerLiteralInst>(DPII->getOperand());
+  if (Op == nullptr)
+    return nullptr;
+
+  if (Op->getValue().uge(PackTy->getNumElements()))
+    return nullptr;
+
+  unsigned Index = Op->getValue().getZExtValue();
+
+  auto *TEAI = Builder.createTupleElementAddr(
+      TPEAI->getLoc(), TPEAI->getTuple(), Index);
+  return Builder.createUncheckedAddrCast(
+      TPEAI->getLoc(), TEAI, TPEAI->getElementType());
+}
+
+// This is a hack. When optimizing a simple pack expansion expression which
+// forms a tuple from a pack, like `(repeat each t)`, after the above
+// peepholes we end up with:
+//
+// %src = unchecked_addr_cast %real_src, @element("...")
+// %dst = unchecked_addr_cast %real_dst, @element("...")
+// copy_addr %src, %dst
+//
+// Simplify this to
+//
+// copy_addr %real_src, %real_dst
+//
+// Assuming that %real_src and %real_dst have the same type.
+//
+// In this simple case, this eliminates the opened element archetype entirely.
+// However, a more principled peephole would be to transform an
+// open_pack_element with a scalar index by replacing all usages of the
+// element archetype with a concrete type.
+SILInstruction *
+SILCombiner::visitCopyAddrInst(CopyAddrInst *CAI) {
+  auto *Src = dyn_cast<UncheckedAddrCastInst>(CAI->getSrc());
+  auto *Dst = dyn_cast<UncheckedAddrCastInst>(CAI->getDest());
+
+  if (Src == nullptr || Dst == nullptr)
+    return nullptr;
+
+  if (Src->getType() != Dst->getType() ||
+      !Src->getType().is<ElementArchetypeType>())
+    return nullptr;
+
+  if (Src->getOperand()->getType() != Dst->getOperand()->getType())
+    return nullptr;
+
+  return Builder.createCopyAddr(
+      CAI->getLoc(), Src->getOperand(), Dst->getOperand(),
+      CAI->isTakeOfSrc(), CAI->isInitializationOfDest());
+}

lib/SILOptimizer/Transforms/CSE.cpp

@@ -515,6 +515,24 @@ class HashVisitor : public SILInstructionVisitor<HashVisitor, llvm::hash_code> {
         X->getKind(), tryLookThroughOwnershipInsts(&X->getOperandRef()),
         llvm::hash_combine_range(ConformsTo.begin(), ConformsTo.end()));
   }
+
+  hash_code visitScalarPackIndexInst(ScalarPackIndexInst *X) {
+    return llvm::hash_combine(
+        X->getKind(), X->getIndexedPackType(), X->getComponentIndex());
+  }
+
+  hash_code visitDynamicPackIndexInst(DynamicPackIndexInst *X) {
+    return llvm::hash_combine(
+        X->getKind(), X->getIndexedPackType(), &X->getOperandRef());
+  }
+
+  hash_code visitTuplePackElementAddrInst(TuplePackElementAddrInst *X) {
+    OperandValueArrayRef Operands(X->getAllOperands());
+    return llvm::hash_combine(
+        X->getKind(),
+        llvm::hash_combine_range(Operands.begin(), Operands.end()),
+        X->getElementType());
+  }
 };
 } // end anonymous namespace
 
@@ -1225,6 +1243,9 @@ bool CSE::canHandle(SILInstruction *Inst) {
   case SILInstructionKind::MarkDependenceInst:
   case SILInstructionKind::InitExistentialMetatypeInst:
   case SILInstructionKind::WitnessMethodInst:
+  case SILInstructionKind::ScalarPackIndexInst:
+  case SILInstructionKind::DynamicPackIndexInst:
+  case SILInstructionKind::TuplePackElementAddrInst:
     // Intentionally we don't handle (prev_)dynamic_function_ref.
     // They change at runtime.
 #define LOADABLE_REF_STORAGE(Name, ...) \

lib/SILOptimizer/Utils/ConstantFolding.cpp

@@ -588,8 +588,31 @@ static SILValue constantFoldBinary(BuiltinInst *BI,
   // Are there valid uses for these in stdlib?
   case BuiltinValueKind::Add:
   case BuiltinValueKind::Mul:
-  case BuiltinValueKind::Sub:
-    return nullptr;
+  case BuiltinValueKind::Sub: {
+    OperandValueArrayRef Args = BI->getArguments();
+    auto *LHS = dyn_cast<IntegerLiteralInst>(Args[0]);
+    auto *RHS = dyn_cast<IntegerLiteralInst>(Args[1]);
+    if (!RHS || !LHS)
+      return nullptr;
+    APInt LHSI = LHS->getValue();
+    APInt RHSI = RHS->getValue();
+
+    switch (ID) {
+    default: llvm_unreachable("Not all cases are covered!");
+    case BuiltinValueKind::Add:
+      LHSI += RHSI;
+      break;
+    case BuiltinValueKind::Mul:
+      LHSI *= RHSI;
+      break;
+    case BuiltinValueKind::Sub:
+      LHSI -= RHSI;
+      break;
+    }
+
+    SILBuilderWithScope B(BI);
+    return B.createIntegerLiteral(BI->getLoc(), BI->getType(), LHSI);
+  }
 
   case BuiltinValueKind::And:
   case BuiltinValueKind::AShr:

test/IRGen/variadic_generic_opaque.swift

@@ -0,0 +1,14 @@
+// RUN: %target-swift-frontend -emit-ir %s -enable-experimental-feature VariadicGenerics -disable-availability-checking
+
+// Because of -enable-experimental-feature VariadicGenerics
+// REQUIRES: asserts
+
+// FIXME: Add more tests
+
+public protocol P {}
+
+public struct G<each T>: P {}
+
+public func returnsG<each T>(_ t: repeat each T) -> some P {
+  return G<repeat each T>()
+}

test/SILOptimizer/constant_propagation.sil

@@ -102,6 +102,25 @@ bb0:
 // CHECK-NEXT: return [[RES]] : $Builtin.Int64
 }
 
+// Compute an expression using a chain of arithmetic with non-overflowing instructions: 2 * (2 + 3) - 3
+sil @fold_arithmetic_without_overflow : $@convention(thin) () -> Builtin.Int64 {
+bb0:
+ %0 = integer_literal $Builtin.Int64, 2
+ %110 = integer_literal $Builtin.Int64, 3
+ %19 = builtin "add_Int64"(%0 : $Builtin.Int64, %110 : $Builtin.Int64) : $Builtin.Int64
+ %21 = builtin "mul_Int64"(%0 : $Builtin.Int64, %19 : $Builtin.Int64) : $Builtin.Int64
+ %23 = builtin "sub_Int64"(%21 : $Builtin.Int64, %110 : $Builtin.Int64) : $Builtin.Int64
+ return %23 : $Builtin.Int64
+
+// CHECK-LABEL: sil @fold_arithmetic_without_overflow
+// CHECK-NOT: integer_literal $Builtin.Int64, 2
+// CHECK-NOT: integer_literal $Builtin.Int64, 3
+// CHECK-NOT: integer_literal $Builtin.Int64, 0
+// CHECK-NOT: builtin
+// CHECK: [[RES:%.*]] = integer_literal $Builtin.Int64, 7
+// CHECK-NEXT: return [[RES]] : $Builtin.Int64
+}
+
 // Fold casts. (This test assumes that DCE does not run, otherwise the unreachable blocks will get removed.)
 sil @fold_trunc : $@convention(thin) () -> Builtin.Int64 {
 bb0: