Ensure we are using mapped SIL type for switch_enum case and not the

original lowered one.

Fixes #73018

Author: asl

lib/IRGen/LoadableByAddress.cpp

@@ -2998,7 +2998,8 @@ namespace {
 }
 
 static void runPeepholesAndReg2Mem(SILPassManager *pm, SILModule *silMod,
-                                   IRGenModule *irgenModule);
+                                   IRGenModule *irgenModule,
+                                   LargeSILTypeMapper &mapperCache);
 
 /// The entry point to this function transformation.
 void LoadableByAddress::run() {  
@@ -3010,7 +3011,8 @@ void LoadableByAddress::run() {
     runOnFunction(&F);
 
   if (modFuncs.empty() && modApplies.empty()) {
-    runPeepholesAndReg2Mem(getPassManager(), getModule(), getIRGenModule());
+    runPeepholesAndReg2Mem(getPassManager(), getModule(), getIRGenModule(),
+                           MapperCache);
     return;
   }
 
@@ -3248,7 +3250,8 @@ void LoadableByAddress::run() {
   storeToBlockStorageInstrs.clear();
 
   getPassManager()->invalidateAllAnalysis();
-  runPeepholesAndReg2Mem(getPassManager(), getModule(), getIRGenModule());
+  runPeepholesAndReg2Mem(getPassManager(), getModule(), getIRGenModule(),
+                         MapperCache);
 }
 
 namespace {
@@ -3396,11 +3399,13 @@ class AddressAssignment {
   GenericEnvironment *genEnv;
   IRGenModule *irgenModule;
   SILFunction &currFn;
+  LargeSILTypeMapper &mapperCache;
 
 public:
   AddressAssignment(IRGenModule *irgenModule, GenericEnvironment *genEnv,
-                    SILFunction &currFn)
-      : genEnv(genEnv), irgenModule(irgenModule), currFn(currFn) {}
+                    SILFunction &currFn, LargeSILTypeMapper &mapperCache)
+      : genEnv(genEnv), irgenModule(irgenModule), currFn(currFn),
+        mapperCache(mapperCache) {}
 
   void assign(SILInstruction *inst);
 
@@ -3456,6 +3461,11 @@ class AddressAssignment {
     return false;
   }
 
+  SILType getNewSILType(CanSILFunctionType fnTy, SILType type) {
+    return mapperCache.getNewSILType(getSubstGenericEnvironment(fnTy), type,
+                                     *irgenModule);
+  }
+
   bool isLargeLoadableType(SILType ty) {
     if (ty.isAddress() || ty.isClassOrClassMetatype()) {
       return false;
@@ -4077,9 +4087,15 @@ class RewriteUser : SILInstructionVisitor<RewriteUser> {
              "caseBB has a payload argument");
 
       SILBuilder caseBuilder = assignment.getBuilder(caseBB->begin());
+      SILType eltType =
+        opdAddr->getType().getEnumElementType(caseDecl,
+                                              caseBuilder.getModule(),
+                                              caseBuilder.getTypeExpansionContext());
+      SILType mappedEltType =
+        assignment.getNewSILType(sw->getFunction()->getLoweredFunctionType(),
+                                 eltType);
       auto *caseAddr =
-          caseBuilder.createUncheckedTakeEnumDataAddr(loc, opdAddr, caseDecl);
-
+        caseBuilder.createUncheckedTakeEnumDataAddr(loc, opdAddr, caseDecl, mappedEltType);
       if (assignment.isLargeLoadableType(caseArg->getType())) {
         assignment.mapValueToAddress(caseArg, caseAddr);
         assignment.markBlockArgumentForDeletion(caseBB);
@@ -4274,7 +4290,8 @@ void AddressAssignment::assign(SILInstruction *inst) {
 }
 
 static void runPeepholesAndReg2Mem(SILPassManager *pm, SILModule *silMod,
-                                   IRGenModule *irgenModule) {
+                                   IRGenModule *irgenModule,
+                                   LargeSILTypeMapper &mapperCache) {
 
   if (!irgenModule->getOptions().EnableLargeLoadableTypesReg2Mem)
     return;
@@ -4301,7 +4318,7 @@ static void runPeepholesAndReg2Mem(SILPassManager *pm, SILModule *silMod,
     // Delete replaced instructions.
     opts.deleteInstructions();
 
-    AddressAssignment assignment(irgenModule, genEnv, currF);
+    AddressAssignment assignment(irgenModule, genEnv, currF, mapperCache);
 
     // Assign addresses to basic block arguments.
 

test/IRGen/loadable_by_address_issue73018.swift

@@ -0,0 +1,44 @@
+// REQUIRES: differentiable_programming
+// RUN: %target-swift-frontend -emit-ir -verify %s
+
+// This used to trigger an assertion due to LoadableByAddress not doing proper mapping of
+// switch_enum arguments during rewriting
+
+import _Differentiation
+struct O: Differentiable {var a: B.G; var b: Array<SIMD2<Float>>; var c: B.M; var d: B.M}
+func w(y: B.M) -> B.M {return y}
+@differentiable(reverse) func q(sd: Float, i: B.G, j: Array<SIMD2<Float>>, s: B.M, t: B.M) -> O {
+    let u1 = C(i.g4); let u2 = C(i.g2); let u3 = C(i.g3); let u4 = C(j)
+    var u5 = C(s.m4); var u6 = C(t.m4); var u7 = w(y: s); let u8 = w(y: t)
+    u6 = C(Array<Float>([0.0]))
+    u5 = u6
+    if sd > 0 {}
+    u7.m4 = u5.r
+    var u9 = i
+    u9.g1 = C(u9.g1).r
+    u9.g4 = u1.r
+    u9.g2 = u2.r
+    u9.g3 = u3.r
+    return O(a: u9, b: u4.r, c: u7, d: u8)
+}
+protocol N {}; protocol H: Differentiable {}
+struct B: D, N, Differentiable
+{
+    struct M: D, N, Differentiable {var m1 = Array<Float>(); var m2 = Array<SIMD2<Float>>(); var m3 = Array<SIMD2<Float>>(); var m4 = Array<Float>(); var m5 = Array<Float>()}
+    struct G: D, N, Differentiable {var g1 = Array<Float>(); var g2 = Array<Float>(); var g3 = Array<Float>(); var g4 = Array<Float>()}
+}
+protocol D: Differentiable & H & N where Self.TangentVector: N {}
+struct C<E>: Differentiable, AdditiveArithmetic where E: Differentiable, E: AdditiveArithmetic {
+    typealias TangentVector = C<E.TangentVector>
+    var v: [E]; var a: E
+    @differentiable(reverse) init(_ c: [E], t: E = .zero) {self.v = c; self.a = .zero}
+    @differentiable(reverse) var r: [E] {return v}
+    @inlinable @derivative(of: init(_:t:))
+    static func vjpInit(_ values: [E], t: E = .zero) -> (value: C, pullback: (Self.TangentVector) -> ([E].TangentVector, E.TangentVector)) {(value: Self(values, t: t), pullback: {tangentVector in return ([E].TangentVector(tangentVector.v), tangentVector.a)})}
+    @inlinable @derivative(of: r) func vjpArray() -> (value: [E], pullback: ([E].TangentVector) -> Self.TangentVector) {(value: self.v, pullback: {_ in return Self.TangentVector([E.TangentVector](), t: E.TangentVector.zero)})}
+    mutating func move(by offset: Self.TangentVector) {}
+    static func + (lhs: Self, rhs: Self) -> Self {return lhs}
+    static func - (lhs: Self, rhs: Self) -> Self {return lhs}
+    static var zero: Self { Self([], t: .zero) }
+}
+