Skip to content

[AutoDiff upstream] AST bits for @differentiable fn ty #28156

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Nov 13, 2019
Merged

[AutoDiff upstream] AST bits for @differentiable fn ty #28156

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 11 additions & 2 deletions include/swift/AST/AutoDiff.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,23 +2,26 @@
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
// Copyright (c) 2019 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//
//
// SWIFT_ENABLE_TENSORFLOW
// This file defines AST support for automatic differentiation.
//
//===----------------------------------------------------------------------===//

#ifndef SWIFT_AST_AUTODIFF_H
#define SWIFT_AST_AUTODIFF_H

#include <cstdint>

#include "swift/AST/Identifier.h"
#include "swift/AST/IndexSubset.h"
#include "swift/Basic/SourceLoc.h"
#include "swift/Basic/Range.h"

namespace swift {
Expand Down Expand Up @@ -86,6 +89,12 @@ class ParsedAutoDiffParameter {
}
};

enum class DifferentiabilityKind : uint8_t {
NonDifferentiable = 0,
Normal = 1,
Linear = 2
};

} // end namespace swift

#endif // SWIFT_AST_AUTODIFF_H
67 changes: 53 additions & 14 deletions include/swift/AST/Types.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
#ifndef SWIFT_TYPES_H
#define SWIFT_TYPES_H

#include "swift/AST/AutoDiff.h"
#include "swift/AST/DeclContext.h"
#include "swift/AST/GenericParamKey.h"
#include "swift/AST/Identifier.h"
Expand Down Expand Up @@ -300,8 +301,8 @@ class alignas(1 << TypeAlignInBits) TypeBase {
}

protected:
enum { NumAFTExtInfoBits = 6 };
enum { NumSILExtInfoBits = 6 };
enum { NumAFTExtInfoBits = 8 };
enum { NumSILExtInfoBits = 8 };
union { uint64_t OpaqueBits;

SWIFT_INLINE_BITFIELD_BASE(TypeBase, bitmax(NumTypeKindBits,8) +
Expand Down Expand Up @@ -2875,14 +2876,16 @@ class AnyFunctionType : public TypeBase {
// If bits are added or removed, then TypeBase::AnyFunctionTypeBits
// and NumMaskBits must be updated, and they must match.
//
// |representation|noEscape|throws|
// | 0 .. 3 | 4 | 5 |
// |representation|noEscape|throws|differentiability|
// | 0 .. 3 | 4 | 5 | 6 .. 7 |
//
enum : unsigned {
RepresentationMask = 0xF << 0,
NoEscapeMask = 1 << 4,
ThrowsMask = 1 << 5,
NumMaskBits = 6
RepresentationMask = 0xF << 0,
NoEscapeMask = 1 << 4,
ThrowsMask = 1 << 5,
DifferentiabilityMaskOffset = 6,
DifferentiabilityMask = 0x3 << DifferentiabilityMaskOffset,
NumMaskBits = 8
};

unsigned Bits; // Naturally sized for speed.
Expand All @@ -2905,13 +2908,24 @@ class AnyFunctionType : public TypeBase {
// Constructor with no defaults.
ExtInfo(Representation Rep,
bool IsNoEscape,
bool Throws)
bool Throws,
DifferentiabilityKind DiffKind)
: ExtInfo(Rep, Throws) {
Bits |= (IsNoEscape ? NoEscapeMask : 0);
Bits |= ((unsigned)DiffKind << DifferentiabilityMaskOffset) &
DifferentiabilityMask;
}

bool isNoEscape() const { return Bits & NoEscapeMask; }
bool throws() const { return Bits & ThrowsMask; }
bool isDifferentiable() const {
return getDifferentiabilityKind() >
DifferentiabilityKind::NonDifferentiable;
}
DifferentiabilityKind getDifferentiabilityKind() const {
return DifferentiabilityKind((Bits & DifferentiabilityMask) >>
DifferentiabilityMaskOffset);
}
Representation getRepresentation() const {
unsigned rawRep = Bits & RepresentationMask;
assert(rawRep <= unsigned(Representation::Last)
Expand Down Expand Up @@ -3069,6 +3083,11 @@ class AnyFunctionType : public TypeBase {
return getExtInfo().throws();
}

bool isDifferentiable() const { return getExtInfo().isDifferentiable(); }
DifferentiabilityKind getDifferentiabilityKind() const {
return getExtInfo().getDifferentiabilityKind();
}

/// Returns a new function type exactly like this one but with the ExtInfo
/// replaced.
AnyFunctionType *withExtInfo(ExtInfo info) const;
Expand Down Expand Up @@ -3716,14 +3735,16 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode,
// If bits are added or removed, then TypeBase::SILFunctionTypeBits
// and NumMaskBits must be updated, and they must match.

// |representation|pseudogeneric| noescape |
// | 0 .. 3 | 4 | 5 |
// |representation|pseudogeneric| noescape |differentiability|
// | 0 .. 3 | 4 | 5 | 6 .. 7 |
//
enum : unsigned {
RepresentationMask = 0xF << 0,
PseudogenericMask = 1 << 4,
NoEscapeMask = 1 << 5,
NumMaskBits = 6
DifferentiabilityMaskOffset = 6,
DifferentiabilityMask = 0x3 << DifferentiabilityMaskOffset,
NumMaskBits = 8
};

unsigned Bits; // Naturally sized for speed.
Expand All @@ -3737,10 +3758,13 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode,
ExtInfo() : Bits(0) { }

// Constructor for polymorphic type.
ExtInfo(Representation rep, bool isPseudogeneric, bool isNoEscape) {
ExtInfo(Representation rep, bool isPseudogeneric, bool isNoEscape,
DifferentiabilityKind diffKind) {
Bits = ((unsigned) rep) |
(isPseudogeneric ? PseudogenericMask : 0) |
(isNoEscape ? NoEscapeMask : 0);
(isNoEscape ? NoEscapeMask : 0) |
(((unsigned)diffKind << DifferentiabilityMaskOffset) &
DifferentiabilityMask);
}

/// Is this function pseudo-generic? A pseudo-generic function
Expand All @@ -3750,6 +3774,16 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode,
// Is this function guaranteed to be no-escape by the type system?
bool isNoEscape() const { return Bits & NoEscapeMask; }

bool isDifferentiable() const {
return getDifferentiabilityKind() !=
DifferentiabilityKind::NonDifferentiable;
}

DifferentiabilityKind getDifferentiabilityKind() const {
return DifferentiabilityKind((Bits & DifferentiabilityMask) >>
DifferentiabilityMaskOffset);
}

/// What is the abstract representation of this function value?
Representation getRepresentation() const {
return Representation(Bits & RepresentationMask);
Expand Down Expand Up @@ -4154,6 +4188,11 @@ class SILFunctionType final : public TypeBase, public llvm::FoldingSetNode,
getRepresentation() == SILFunctionTypeRepresentation::Thick;
}

bool isDifferentiable() const { return getExtInfo().isDifferentiable(); }
DifferentiabilityKind getDifferentiabilityKind() const {
return getExtInfo().getDifferentiabilityKind();
}

bool isNoReturnFunction(SILModule &M) const; // Defined in SILType.cpp

/// Create a SILFunctionType with the same parameters, results, and attributes as this one, but with
Expand Down
6 changes: 3 additions & 3 deletions lib/AST/ASTDemangler.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -479,9 +479,9 @@ Type ASTBuilder::createImplFunctionType(
break;
}

auto einfo = SILFunctionType::ExtInfo(representation,
flags.isPseudogeneric(),
!flags.isEscaping());
auto einfo = SILFunctionType::ExtInfo(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a mangling for these flags? If so, it seems like it should be plumbed in here for the demangler to use.

Copy link
Author

@marcrasi marcrasi Nov 11, 2019 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We've got some mangling in the tensorflow branch. I will upstream it in a separate PR. There is a ticket for that: https://bugs.swift.org/browse/TF-948

representation, flags.isPseudogeneric(), !flags.isEscaping(),
DifferentiabilityKind::NonDifferentiable);

llvm::SmallVector<SILParameterInfo, 8> funcParams;
llvm::SmallVector<SILYieldInfo, 8> funcYields;
Expand Down
17 changes: 17 additions & 0 deletions lib/AST/ASTPrinter.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3789,6 +3789,14 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
if (Options.SkipAttributes)
return;

if (!Options.excludeAttrKind(TAK_differentiable) &&
info.isDifferentiable()) {
if (info.getDifferentiabilityKind() == DifferentiabilityKind::Linear) {
Printer << "@differentiable(linear) ";
} else {
Printer << "@differentiable ";
}
}

if (Options.PrintFunctionRepresentationAttrs &&
!Options.excludeAttrKind(TAK_convention) &&
Expand Down Expand Up @@ -3833,6 +3841,15 @@ class TypePrinter : public TypeVisitor<TypePrinter> {
if (Options.SkipAttributes)
return;

if (!Options.excludeAttrKind(TAK_differentiable) &&
info.isDifferentiable()) {
if (info.getDifferentiabilityKind() == DifferentiabilityKind::Linear) {
Printer << "@differentiable(linear) ";
} else {
Printer << "@differentiable ";
}
}

if (Options.PrintFunctionRepresentationAttrs &&
!Options.excludeAttrKind(TAK_convention) &&
info.getRepresentation() != SILFunctionType::Representation::Thick) {
Expand Down
8 changes: 4 additions & 4 deletions lib/SILGen/SILGen.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -415,10 +415,10 @@ SILGenModule::getKeyPathProjectionCoroutine(bool isReadAccess,
: ParameterConvention::Indirect_In_Guaranteed },
};

auto extInfo =
SILFunctionType::ExtInfo(SILFunctionTypeRepresentation::Thin,
/*pseudogeneric*/false,
/*non-escaping*/false);
auto extInfo = SILFunctionType::ExtInfo(
SILFunctionTypeRepresentation::Thin,
/*pseudogeneric*/ false,
/*non-escaping*/ false, DifferentiabilityKind::NonDifferentiable);

auto functionTy = SILFunctionType::get(sig, extInfo,
SILCoroutineKind::YieldOnce,
Expand Down
12 changes: 8 additions & 4 deletions lib/SILGen/SILGenExpr.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2669,7 +2669,8 @@ static SILFunction *getOrCreateKeyPathGetter(SILGenModule &SGM,
auto signature = SILFunctionType::get(genericSig,
SILFunctionType::ExtInfo(SILFunctionType::Representation::Thin,
/*pseudogeneric*/ false,
/*noescape*/ false),
/*noescape*/ false,
DifferentiabilityKind::NonDifferentiable),
SILCoroutineKind::None,
ParameterConvention::Direct_Unowned,
params, {}, result, None,
Expand Down Expand Up @@ -2811,7 +2812,8 @@ static SILFunction *getOrCreateKeyPathSetter(SILGenModule &SGM,
auto signature = SILFunctionType::get(genericSig,
SILFunctionType::ExtInfo(SILFunctionType::Representation::Thin,
/*pseudogeneric*/ false,
/*noescape*/ false),
/*noescape*/ false,
DifferentiabilityKind::NonDifferentiable),
SILCoroutineKind::None,
ParameterConvention::Direct_Unowned,
params, {}, {}, None,
Expand Down Expand Up @@ -2987,7 +2989,8 @@ getOrCreateKeyPathEqualsAndHash(SILGenModule &SGM,
auto signature = SILFunctionType::get(genericSig,
SILFunctionType::ExtInfo(SILFunctionType::Representation::Thin,
/*pseudogeneric*/ false,
/*noescape*/ false),
/*noescape*/ false,
DifferentiabilityKind::NonDifferentiable),
SILCoroutineKind::None,
ParameterConvention::Direct_Unowned,
params, /*yields*/ {}, results, None,
Expand Down Expand Up @@ -3162,7 +3165,8 @@ getOrCreateKeyPathEqualsAndHash(SILGenModule &SGM,
auto signature = SILFunctionType::get(genericSig,
SILFunctionType::ExtInfo(SILFunctionType::Representation::Thin,
/*pseudogeneric*/ false,
/*noescape*/ false),
/*noescape*/ false,
DifferentiabilityKind::NonDifferentiable),
SILCoroutineKind::None,
ParameterConvention::Direct_Unowned,
params, /*yields*/ {}, results, None,
Expand Down
11 changes: 6 additions & 5 deletions lib/SILOptimizer/Transforms/Outliner.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -288,7 +288,8 @@ CanSILFunctionType BridgedProperty::getOutlinedFunctionType(SILModule &M) {
ResultConvention::Owned));
auto ExtInfo =
SILFunctionType::ExtInfo(SILFunctionType::Representation::Thin,
/*pseudogeneric*/ false, /*noescape*/ false);
/*pseudogeneric*/ false, /*noescape*/ false,
DifferentiabilityKind::NonDifferentiable);
auto FunctionType = SILFunctionType::get(
nullptr, ExtInfo, SILCoroutineKind::None,
ParameterConvention::Direct_Unowned, Parameters, /*yields*/ {},
Expand Down Expand Up @@ -1108,10 +1109,10 @@ CanSILFunctionType ObjCMethodCall::getOutlinedFunctionType(SILModule &M) {
OrigSigIdx++;
}

auto ExtInfo =
SILFunctionType::ExtInfo(SILFunctionType::Representation::Thin,
/*pseudogeneric*/ false,
/*noescape*/ false);
auto ExtInfo = SILFunctionType::ExtInfo(
SILFunctionType::Representation::Thin,
/*pseudogeneric*/ false,
/*noescape*/ false, DifferentiabilityKind::NonDifferentiable);

SmallVector<SILResultInfo, 4> Results;
// If we don't have a bridged return we changed from @autoreleased to @owned
Expand Down
12 changes: 6 additions & 6 deletions lib/SILOptimizer/UtilityPasses/BugReducerTester.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,13 +83,13 @@ class BugReducerTester : public SILFunctionTransform {
ResultInfoArray.push_back(
SILResultInfo(EmptyTupleCanType, ResultConvention::Unowned));
auto FuncType = SILFunctionType::get(
nullptr, SILFunctionType::ExtInfo(SILFunctionType::Representation::Thin,
false /*isPseudoGeneric*/,
false /*noescape*/),
nullptr,
SILFunctionType::ExtInfo(SILFunctionType::Representation::Thin,
false /*isPseudoGeneric*/, false /*noescape*/,
DifferentiabilityKind::NonDifferentiable),
SILCoroutineKind::None, ParameterConvention::Direct_Unowned,
ArrayRef<SILParameterInfo>(), ArrayRef<SILYieldInfo>(),
ResultInfoArray, None,
SubstitutionMap(), false,
ArrayRef<SILParameterInfo>(), ArrayRef<SILYieldInfo>(), ResultInfoArray,
None, SubstitutionMap(), false,
getFunction()->getModule().getASTContext());

SILOptFunctionBuilder FunctionBuilder(*this);
Expand Down
12 changes: 8 additions & 4 deletions lib/Sema/ConstraintSystem.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1642,7 +1642,8 @@ resolveOverloadForDeclWithSpecialTypeCheckingSemantics(ConstraintSystem &CS,
auto bodyClosure = FunctionType::get(arg, result,
FunctionType::ExtInfo(FunctionType::Representation::Swift,
/*noescape*/ true,
/*throws*/ true));
/*throws*/ true,
DifferentiabilityKind::NonDifferentiable));
FunctionType::Param args[] = {
FunctionType::Param(noescapeClosure),
FunctionType::Param(bodyClosure, CS.getASTContext().getIdentifier("do")),
Expand All @@ -1651,7 +1652,8 @@ resolveOverloadForDeclWithSpecialTypeCheckingSemantics(ConstraintSystem &CS,
refType = FunctionType::get(args, result,
FunctionType::ExtInfo(FunctionType::Representation::Swift,
/*noescape*/ false,
/*throws*/ true));
/*throws*/ true,
DifferentiabilityKind::NonDifferentiable));
openedFullType = refType;
return true;
}
Expand All @@ -1674,15 +1676,17 @@ resolveOverloadForDeclWithSpecialTypeCheckingSemantics(ConstraintSystem &CS,
auto bodyClosure = FunctionType::get(bodyArgs, result,
FunctionType::ExtInfo(FunctionType::Representation::Swift,
/*noescape*/ true,
/*throws*/ true));
/*throws*/ true,
DifferentiabilityKind::NonDifferentiable));
FunctionType::Param args[] = {
FunctionType::Param(existentialTy),
FunctionType::Param(bodyClosure, CS.getASTContext().getIdentifier("do")),
};
refType = FunctionType::get(args, result,
FunctionType::ExtInfo(FunctionType::Representation::Swift,
/*noescape*/ false,
/*throws*/ true));
/*throws*/ true,
DifferentiabilityKind::NonDifferentiable));
openedFullType = refType;
return true;
}
Expand Down
Loading