Skip to content

[AutoDiff] Support ref_element_addr differentiation. #29749

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 5 commits into from
Feb 11, 2020
Merged

[AutoDiff] Support ref_element_addr differentiation. #29749

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
cda8e1c
[AutoDiff] Support `ref_element_addr` differentiation.
dan-zheng Feb 10, 2020
4f7f9bf
NFC: fix formatting.
dan-zheng Feb 10, 2020
9a55217
Add test.
dan-zheng Feb 10, 2020
610214a
Make diagnostic more user-friendly.
dan-zheng Feb 11, 2020
358c1ce
Fix incorrect `ref_element_addr` derivative values.
dan-zheng Feb 11, 2020
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
9 changes: 6 additions & 3 deletions include/swift/AST/DiagnosticsSIL.def
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -491,6 +491,12 @@ WARNING(autodiff_nonvaried_result_fixit,none,
"result does not depend on differentiation arguments and will always "
"have a zero derivative; do you want to use 'withoutDerivative(at:)'?",
())
// TODO(TF-1149): Remove this diagnostic.
NOTE(autodiff_loadable_value_addressonly_tangent_unsupported,none,
"cannot yet differentiate value whose type %0 has a compile-time known "
"size, but whose 'TangentVector' contains stored properties of unknown "
"size; consider modifying %1 to use fewer generic parameters in stored "
"properties", (Type, Type))
NOTE(autodiff_enums_unsupported,none,
"differentiating enum values is not yet supported", ())
NOTE(autodiff_global_let_closure_not_differentiable,none,
Expand Down Expand Up @@ -543,9 +549,6 @@ NOTE(autodiff_jvp_control_flow_not_supported,none,
"forward-mode differentiation does not yet support control flow", ())
NOTE(autodiff_control_flow_not_supported,none,
"cannot differentiate unsupported control flow", ())
// TODO(TF-645): Remove when differentiation supports `ref_element_addr`.
NOTE(autodiff_class_property_not_supported,none,
"differentiating class properties is not yet supported", ())
// TODO(TF-1080): Remove when differentiation supports `begin_apply`.
NOTE(autodiff_coroutines_not_supported,none,
"differentiation of coroutine calls is not yet supported", ())
Expand Down
10 changes: 5 additions & 5 deletions include/swift/SILOptimizer/Analysis/DifferentiableActivityAnalysis.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -157,8 +157,8 @@ class DifferentiableActivityInfo {
/// Propagates variedness from the given operand to its user's results.
void propagateVaried(Operand *operand, unsigned independentVariableIndex);
/// Marks the given value as varied and recursively propagates variedness
/// inwards (to operands) through projections. Skips `@noDerivative` struct
/// field projections.
/// inwards (to operands) through projections. Skips `@noDerivative` field
/// projections.
void
propagateVariedInwardsThroughProjections(SILValue value,
unsigned independentVariableIndex);
Expand All @@ -172,9 +172,9 @@ class DifferentiableActivityInfo {
unsigned dependentVariableIndex);
/// Propagates usefulnesss to the operands of the given instruction.
void propagateUseful(SILInstruction *inst, unsigned dependentVariableIndex);
/// Marks the given address as useful and recursively propagates usefulness
/// inwards (to operands) through projections. Skips `@noDerivative` struct
/// field projections.
/// Marks the given address or class-typed value as useful and recursively
/// propagates usefulness inwards (to operands) through projections. Skips
/// `@noDerivative` field projections.
void propagateUsefulThroughAddress(SILValue value,
unsigned dependentVariableIndex);
/// If the given value is an `array.uninitialized_intrinsic` application,
Expand Down
20 changes: 16 additions & 4 deletions include/swift/SILOptimizer/Utils/Differentiation/PullbackEmitter.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -266,10 +266,17 @@ class PullbackEmitter final : public SILInstructionVisitor<PullbackEmitter> {

SILBasicBlock::iterator getNextFunctionLocalAllocationInsertionPoint();

/// Creates and returns a local allocation with the given type.
///
/// Local allocations are created uninitialized in the pullback entry and
/// deallocated in the pullback exit. All local allocations not in
/// `destroyedLocalAllocations` are also destroyed in the pullback exit.
AllocStackInst *createFunctionLocalAllocation(SILType type, SILLocation loc);

SILValue &getAdjointBuffer(SILBasicBlock *origBB, SILValue originalBuffer);

// Accumulates `rhsBufferAccess` into the adjoint buffer corresponding to
// `originalBuffer`.
/// Accumulates `rhsBufferAccess` into the adjoint buffer corresponding to
/// `originalBuffer`.
void addToAdjointBuffer(SILBasicBlock *origBB, SILValue originalBuffer,
SILValue rhsBufferAccess, SILLocation loc);

Expand Down Expand Up @@ -354,6 +361,13 @@ class PullbackEmitter final : public SILInstructionVisitor<PullbackEmitter> {
/// field in tangent space corresponding to #field
void visitStructExtractInst(StructExtractInst *sei);

/// Handle `ref_element_addr` instruction.
/// Original: y = ref_element_addr x, <n>
/// Adjoint: adj[x] += struct (0, ..., #field': adj[y], ..., 0)
/// ^~~~~~~
/// field in tangent space corresponding to #field
void visitRefElementAddrInst(RefElementAddrInst *reai);
Copy link
Contributor

Choose a reason for hiding this comment

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

We should get rid of this header soon...

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I can do it soon :)


/// Handle `tuple` instruction.
/// Original: y = tuple (x0, x1, x2, ...)
/// Adjoint: (adj[x0], adj[x1], adj[x2], ...) += destructure_tuple adj[y]
Expand Down Expand Up @@ -421,8 +435,6 @@ class PullbackEmitter final : public SILInstructionVisitor<PullbackEmitter> {
UnconditionalCheckedCastAddrInst *uccai);

#define NOT_DIFFERENTIABLE(INST, DIAG) void visit##INST##Inst(INST##Inst *inst);

NOT_DIFFERENTIABLE(RefElementAddr, autodiff_class_property_not_supported)
#undef NOT_DIFFERENTIABLE

#define NO_ADJOINT(INST) \
Expand Down
18 changes: 12 additions & 6 deletions lib/SIL/SILFunctionType.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -275,9 +275,12 @@ CanSILFunctionType SILFunctionType::getAutoDiffDerivativeFunctionType(
switch (origResConv) {
case ResultConvention::Owned:
case ResultConvention::Autoreleased:
conv = tl.isTrivial()
? ParameterConvention::Direct_Unowned
: ParameterConvention::Direct_Guaranteed;
if (tl.isAddressOnly()) {
conv = ParameterConvention::Indirect_In_Guaranteed;
} else {
conv = tl.isTrivial() ? ParameterConvention::Direct_Unowned
: ParameterConvention::Direct_Guaranteed;
}
break;
case ResultConvention::Unowned:
case ResultConvention::UnownedInnerPointer:
Expand All @@ -301,9 +304,12 @@ CanSILFunctionType SILFunctionType::getAutoDiffDerivativeFunctionType(
case ParameterConvention::Direct_Owned:
case ParameterConvention::Direct_Guaranteed:
case ParameterConvention::Direct_Unowned:
conv = tl.isTrivial()
? ResultConvention::Unowned
: ResultConvention::Owned;
if (tl.isAddressOnly()) {
conv = ResultConvention::Indirect;
} else {
conv = tl.isTrivial() ? ResultConvention::Unowned
: ResultConvention::Owned;
}
break;
case ParameterConvention::Indirect_In:
case ParameterConvention::Indirect_Inout:
Expand Down
56 changes: 32 additions & 24 deletions lib/SILOptimizer/Analysis/DifferentiableActivityAnalysis.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -166,18 +166,20 @@ void DifferentiableActivityInfo::propagateVaried(
setVariedAndPropagateToUsers(teai, i);
}
}
// Handle `struct_extract` and `struct_element_addr` instructions.
// Handle element projection instructions:
// `struct_extract`, `struct_element_addr`, `ref_element_addr`.
// - If the field is marked `@noDerivative`, do not set the result as
// varied because it does not need a derivative.
// - Otherwise, propagate variedness from operand to result as usual.
#define PROPAGATE_VARIED_FOR_STRUCT_EXTRACTION(INST) \
else if (auto *sei = dyn_cast<INST##Inst>(inst)) { \
if (isVaried(sei->getOperand(), i) && \
!sei->getField()->getAttrs().hasAttribute<NoDerivativeAttr>()) \
setVariedAndPropagateToUsers(sei, i); \
#define PROPAGATE_VARIED_FOR_ELEMENT_PROJECTION(INST) \
else if (auto *projInst = dyn_cast<INST##Inst>(inst)) { \
if (isVaried(projInst->getOperand(), i) && \
!projInst->getField()->getAttrs().hasAttribute<NoDerivativeAttr>()) \
setVariedAndPropagateToUsers(projInst, i); \
}
PROPAGATE_VARIED_FOR_STRUCT_EXTRACTION(StructExtract)
PROPAGATE_VARIED_FOR_STRUCT_EXTRACTION(StructElementAddr)
PROPAGATE_VARIED_FOR_ELEMENT_PROJECTION(StructExtract)
PROPAGATE_VARIED_FOR_ELEMENT_PROJECTION(StructElementAddr)
PROPAGATE_VARIED_FOR_ELEMENT_PROJECTION(RefElementAddr)
#undef PROPAGATE_VARIED_FOR_STRUCT_EXTRACTION
// Handle `br`.
else if (auto *bi = dyn_cast<BranchInst>(inst)) {
Expand Down Expand Up @@ -222,13 +224,14 @@ static Optional<AccessorKind> getAccessorKind(SILFunction *fn) {
void DifferentiableActivityInfo::propagateVariedInwardsThroughProjections(
SILValue value, unsigned independentVariableIndex) {
auto i = independentVariableIndex;
// Skip `@noDerivative` struct projections.
// Skip `@noDerivative` projections.
#define SKIP_NODERIVATIVE(INST) \
if (auto *sei = dyn_cast<INST##Inst>(value)) \
if (sei->getField()->getAttrs().hasAttribute<NoDerivativeAttr>()) \
if (auto *projInst = dyn_cast<INST##Inst>(value)) \
if (projInst->getField()->getAttrs().hasAttribute<NoDerivativeAttr>()) \
return;
SKIP_NODERIVATIVE(StructExtract)
SKIP_NODERIVATIVE(StructElementAddr)
SKIP_NODERIVATIVE(RefElementAddr)
#undef SKIP_NODERIVATIVE
// Set value as varied and propagate to users.
setVariedAndPropagateToUsers(value, i);
Expand Down Expand Up @@ -274,7 +277,8 @@ void DifferentiableActivityInfo::setUsefulAndPropagateToOperands(
// Skip already-useful values to prevent infinite recursion.
if (isUseful(value, dependentVariableIndex))
return;
if (value->getType().isAddress()) {
if (value->getType().isAddress() ||
value->getType().getClassOrBoundGenericClass()) {
propagateUsefulThroughAddress(value, dependentVariableIndex);
return;
}
Expand Down Expand Up @@ -331,15 +335,16 @@ void DifferentiableActivityInfo::propagateUseful(
PROPAGATE_USEFUL_THROUGH_STORE(CopyAddr)
PROPAGATE_USEFUL_THROUGH_STORE(UnconditionalCheckedCastAddr)
#undef PROPAGATE_USEFUL_THROUGH_STORE
// Handle struct element extraction, skipping `@noDerivative` fields:
// `struct_extract`, `struct_element_addr`.
#define PROPAGATE_USEFUL_THROUGH_STRUCT_EXTRACTION(INST) \
else if (auto *sei = dyn_cast<INST##Inst>(inst)) { \
if (!sei->getField()->getAttrs().hasAttribute<NoDerivativeAttr>()) \
setUsefulAndPropagateToOperands(sei->getOperand(), i); \
// Handle element projections, skipping `@noDerivative` fields:
// `struct_extract`, `struct_element_addr`, `ref_element_addr`.
#define PROPAGATE_USEFUL_THROUGH_ELEMENT_PROJECTION(INST) \
else if (auto *projInst = dyn_cast<INST##Inst>(inst)) { \
if (!projInst->getField()->getAttrs().hasAttribute<NoDerivativeAttr>()) \
setUsefulAndPropagateToOperands(projInst->getOperand(), i); \
}
PROPAGATE_USEFUL_THROUGH_STRUCT_EXTRACTION(StructExtract)
PROPAGATE_USEFUL_THROUGH_STRUCT_EXTRACTION(StructElementAddr)
PROPAGATE_USEFUL_THROUGH_ELEMENT_PROJECTION(StructExtract)
PROPAGATE_USEFUL_THROUGH_ELEMENT_PROJECTION(StructElementAddr)
PROPAGATE_USEFUL_THROUGH_ELEMENT_PROJECTION(RefElementAddr)
#undef PROPAGATE_USEFUL_THROUGH_STRUCT_EXTRACTION
// Handle everything else.
else {
Expand All @@ -350,7 +355,8 @@ void DifferentiableActivityInfo::propagateUseful(

void DifferentiableActivityInfo::propagateUsefulThroughAddress(
SILValue value, unsigned dependentVariableIndex) {
assert(value->getType().isAddress());
assert(value->getType().isAddress() ||
value->getType().getClassOrBoundGenericClass());
// Skip already-useful values to prevent infinite recursion.
if (isUseful(value, dependentVariableIndex))
return;
Expand All @@ -364,13 +370,15 @@ void DifferentiableActivityInfo::propagateUsefulThroughAddress(
propagateUseful(use->getUser(), dependentVariableIndex);
for (auto res : use->getUser()->getResults()) {
#define SKIP_NODERIVATIVE(INST) \
if (auto *sei = dyn_cast<INST##Inst>(res)) \
if (sei->getField()->getAttrs().hasAttribute<NoDerivativeAttr>()) \
if (auto *projInst = dyn_cast<INST##Inst>(res)) \
if (projInst->getField()->getAttrs().hasAttribute<NoDerivativeAttr>()) \
continue;
SKIP_NODERIVATIVE(StructExtract)
SKIP_NODERIVATIVE(StructElementAddr)
SKIP_NODERIVATIVE(RefElementAddr)
#undef SKIP_NODERIVATIVE
if (Projection::isAddressProjection(res) || isa<BeginAccessInst>(res))
if (Projection::isAddressProjection(res) || isa<BeginAccessInst>(res) ||
isa<BeginBorrowInst>(res))
propagateUsefulThroughAddress(res, dependentVariableIndex);
}
}
Expand Down
Loading