[AutoDiff] Start linear_function canonicalization skeleton

Start `linear_function` canonicalization skeleton copying from `differentiable_function`
canonicalization. For now, transpose function operands are filled in with `undef`.

Author: efremale

include/swift/SILOptimizer/Differentiation/ADContext.h

@@ -219,8 +219,7 @@ class ADContext {
   /// pointer value as a previously processed and deleted instruction.
   LinearFunctionInst *
   createLinearFunction(SILBuilder &builder, SILLocation loc,
-                       IndexSubset *parameterIndices,
-                       IndexSubset *resultIndices, SILValue original,
+                       IndexSubset *parameterIndices, SILValue original,
                        Optional<SILValue> transposeFunction = None);
 
   // Given a `differentiable_function` instruction, finds the corresponding

lib/SILOptimizer/Differentiation/ADContext.cpp

@@ -123,6 +123,15 @@ DifferentiableFunctionInst *ADContext::createDifferentiableFunction(
   return dfi;
 }
 
+LinearFunctionInst *ADContext::createLinearFunction(
+    SILBuilder &builder, SILLocation loc, IndexSubset *parameterIndices,
+    SILValue original, Optional<SILValue> transposeFunction) {
+  auto *lfi = builder.createLinearFunction(loc, parameterIndices, original,
+                                           transposeFunction);
+  processedLinearFunctionInsts.erase(lfi);
+  return lfi;
+}
+
 DifferentiableFunctionExpr *
 ADContext::findDifferentialOperator(DifferentiableFunctionInst *inst) {
   return inst->getLoc().getAsASTNode<DifferentiableFunctionExpr>();

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -755,19 +755,6 @@ emitDerivativeFunctionReference(
   return None;
 }
 
-/// Emits a reference to the transpose function of `originalFunction`,
-/// differentiated with respect to exactly `desiredIndices`. Returns the
-/// transpose function `SILValue`.
-///
-/// Returns `None` on failure, signifying that a diagnostic has been emitted
-/// using `invoker`.
-static Optional<SILValue> emitTransposeFunctionReference(
-    DifferentiationTransformer &transformer, SILBuilder &builder,
-    SILAutoDiffIndices desiredIndices, SILValue originalFunction,
-    DifferentiationInvoker invoker) {
-  // TODO: Fill in.
-}
-
 //===----------------------------------------------------------------------===//
 // `SILDifferentiabilityWitness` processing
 //===----------------------------------------------------------------------===//
@@ -1226,13 +1213,25 @@ SILValue DifferentiationTransformer::promoteToDifferentiableFunction(
 }
 
 SILValue DifferentiationTransformer::promoteToLinearFunction(
-    LinearFunctionInst *inst, SILBuilder &builder, SILLocation loc,
+    LinearFunctionInst *lfi, SILBuilder &builder, SILLocation loc,
     DifferentiationInvoker invoker) {
   // TODO: Fill in. Copy code from above.
   // For now, create a new `linear_function` instruction with an undef
   // transpose.
   // Eventually, use `emitTransposeFunctionReference` to fill in legitimately.
-  return inst;
+  auto origFnOperand = lfi->getOriginalFunction();
+  auto origFnCopy = builder.emitCopyValueOperation(loc, origFnOperand);
+  auto *parameterIndices = lfi->getParameterIndices();
+  auto originalType = origFnOperand->getType().castTo<SILFunctionType>();
+  auto transposeFnType = originalType->getAutoDiffTransposeFunctionType(
+      parameterIndices, context.getTypeConverter(),
+      LookUpConformanceInModule(builder.getModule().getSwiftModule()));
+  auto transposeType = SILType::getPrimitiveObjectType(transposeFnType);
+  auto transposeFn = SILUndef::get(transposeType, builder.getFunction());
+  auto *newLinearFn = context.createLinearFunction(
+      builder, loc, parameterIndices, origFnCopy, SILValue(transposeFn));
+  context.getLinearFunctionInstWorklist().push_back(lfi);
+  return newLinearFn;
 }
 
 /// Fold `differentiable_function_extract` users of the given
@@ -1390,7 +1389,8 @@ void Differentiation::run() {
 
   // If nothing has triggered differentiation, there's nothing to do.
   if (context.getInvokers().empty() &&
-      context.getDifferentiableFunctionInstWorklist().empty())
+      context.getDifferentiableFunctionInstWorklist().empty() &&
+      context.getLinearFunctionInstWorklist().empty())
     return;
 
   // Differentiation relies on the stdlib (the Swift module).

test/AutoDiff/SILOptimizer/linear_function_canonicalization.sil

@@ -0,0 +1,28 @@
+// RUN: %target-sil-opt -differentiation %s | %FileCheck %s
+
+sil_stage raw
+
+import Swift
+import Builtin
+
+import _Differentiation
+
+sil hidden @foo : $@convention(thin) (Float, Float, Float) -> Float {
+bb0(%0 : $Float, %1 : $Float, %2 : $Float):
+  return %2 : $Float
+}
+
+sil @make_linear_func : $@convention(thin) () -> () {
+bb0:
+  %orig = function_ref @foo : $@convention(thin) (Float, Float, Float) -> Float
+  %linear_fn_0 = linear_function [parameters 0] %orig : $@convention(thin) (Float, Float, Float) -> Float
+  %linear_fn_1 = linear_function [parameters 0 2] %orig : $@convention(thin) (Float, Float, Float) -> Float
+  return undef : $()
+}
+
+// CHECK-LABEL: sil @make_linear_func
+// CHECK: bb0:
+// CHECK:   [[ORIG_FN:%.*]] = function_ref @foo : $@convention(thin) (Float, Float, Float) -> Float
+// CHECK:   linear_function [parameters 0] %0 : $@convention(thin) (Float, Float, Float) -> Float with_transpose undef : $@convention(thin) (Float, Float, Float) -> Float
+// CHECK:   linear_function [parameters 0 2] %0 : $@convention(thin) (Float, Float, Float) -> Float with_transpose undef : $@convention(thin) (Float, Float) -> (Float, Float)
+// CHECK: }