differentiate struct_extract instructions using VJP #21567
Conversation
|
@swift-ci please test tensorflow |
| SmallVector<SILValue, 8> args; | ||
| auto seed = getAdjointValue(sei); | ||
| auto *seedBuf = builder.createAllocStack(loc, seed.getType()); | ||
| materializeAdjointIndirectHelper(seed, seedBuf); |
There was a problem hiding this comment.
Instead of always allocating stack and emitting a load when the adjoint is direct, could you call materializeAdjointDirect directly? It'd simplify both your logic and possibly generated code.
There was a problem hiding this comment.
Oh cool, in fact, I should always be able to materializeAdjointDirect because struct_extract only works on object types. This simplifies everything.
| SmallVector<SILValue, 8> args; | ||
| auto seed = getAdjointValue(sei); | ||
| auto *seedBuf = builder.createAllocStack(loc, seed.getType()); | ||
| materializeAdjointIndirectHelper(seed, seedBuf); |
There was a problem hiding this comment.
The helper materializeAdjointIndirectHelper is not supposed to be used outside of materializeAdjointIndirect. It expects the buffer argument to be the result of begin_access. Please switch to materializeAdjointIndirect if that does what you need.
| void visitStructExtractInst(StructExtractInst *sei) { | ||
| auto *structDecl = sei->getStructDecl(); |
There was a problem hiding this comment.
When we have derived conformances to Differentiable, it'd be much more efficient to use the original struct_extract differentiation logic when we know that the cotangent type equals the original type and that the property vjp is compiler-synthesized (this can be checked via Decl::isImplicit(), I think).
There was a problem hiding this comment.
In fact, this gives me an idea. The original struct_extract differentiation logic does not need getter VJPs. So if I can get the original struct_extract differentiation logic working on structs with derived conformance to Differentiable (even those with cotangent type not equal to self), then I don't need to write any code that synthesizes getter VJPs.
I think that the original struct_extract logic can be made to work even on derived conformances to Differentiable with cotangent type not equal to self, because it should be possible to determine which fields correspond.
I will try out this approach.
There was a problem hiding this comment.
Sounds reasonable for internal structs. Public structs need property VJPs in any case because of resilience requirements.
| loc, getterVJPRef, /*substitutionMap*/ {}, | ||
| /*args*/ {getMappedValue(sei->getOperand())}, /*isNonThrowing*/ false); | ||
|
|
||
| // Get the VJP results (original results and pullback) |
There was a problem hiding this comment.
Nit
| // Get the VJP results (original results and pullback) | |
| // Get the VJP results (original results and pullback). |
| assert(getterDecl); | ||
| auto *getterFn = getContext().getModule().lookUpFunction( | ||
| SILDeclRef(getterDecl, SILDeclRef::Kind::Func)); | ||
| if (!getterFn) { |
There was a problem hiding this comment.
When will a getter not exist?
There was a problem hiding this comment.
One example is the case that I added to test/AutoDiff/autodiff_diagnostics.swift.
I don't know why the getter doesn't exist in that case, and if I put the same code in a file an compile it with a plain call to swiftc, the getter does exist and the code ends up triggering the case where the getter exists but doesn't have a VJP.
I haven't investigated any farther than that, so I don't really know what's going on but the logic seems to handle all the cases I can think of correctly.
| mapValue(sei, originalDirectResult); | ||
|
|
||
| // Checkpoint the original results. | ||
| getPrimalInfo().addStaticPrimalValueDecl(sei); |
There was a problem hiding this comment.
reminder to self: must add a retain here
getBuilder().createRetainValue(loc, originalDirectResult,
getBuilder().getDefaultAtomicity());
| public static prefix func - (x: ${Self}) -> ${Self} { | ||
| return ${Self}(Builtin.fneg_FPIEEE${bits}(x._value)) | ||
| } | ||
| } | ||
|
|
||
| // SWIFT_ENABLE_TENSORFLOW | ||
| extension ${Self} { | ||
| @inlinable |
There was a problem hiding this comment.
I don’t think inlinable is necessary.
There was a problem hiding this comment.
I get a SILVerifier error about referencing a non-ABI-public function from a fragile function if I remove inlinable. (simple_math.swift triggers it)
I think what's happening is that _transparent doesn't make the function ABI public, and mandatory inlining isn't inlining the function because we're partially applying it.
There was a problem hiding this comment.
Gotcha. If you just need it to be ABI-public without needing it serializable, just use @usableFromInline.
marcrasi
force-pushed
the
ad-struct-extract
branch
from
0928c65 to
0c73b3e
Compare
|
@swift-ci please test tensorflow |
Differentiates
struct_extractinstructions using the VJP of the corresponding getter.Float._valuehas no VJP defined so AD can't differentiatepublic static prefix func - (x: ${Self}). This could be fixed by defining a VJP forFloat._valueor by adding a custom adjoint for-. The latter sounded less invasive so I did that.