Skip to content
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.

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

[SR-9791] [AD] PrimalGen attempts to analyze an external function in multi-file compilation #52216

Closed
swift-ci opened this issue Jan 29, 2019 · 5 comments
Assignees
Labels
bug A deviation from expected or documented behavior. Also: expected but undesirable behavior. swift for tensorflow

Comments

@swift-ci
Copy link
Collaborator

Previous ID SR-9791
Radar None
Original Reporter jekbradbury (JIRA User)
Type Bug
Status Resolved
Resolution Done
Additional Detail from JIRA
Votes 0
Component/s Swift for TensorFlow
Labels Bug
Assignee @dan-zheng
Priority Medium

md5: 97eeb6a5e983d062674abda5d7020988

relates to:

  • TF-24 [AD] Differentiation task registration crasher

Issue Description:

When invoking the compiler as swiftc main.swift other.swift, with a function declared with @differentiable in other.swift and an attempted differentiation of that function in main.swift, PrimalGen crashes because it attempts to invoke itself on the external function.

// other.swift
@differentiable
func foo(_ x: Float) -> Float { return x }
// main.swift
let _ = gradient(at: Float(0), in: foo)
frame #&#8203;4: 0x0000000001067292 swift`swift::SILAnalysis::verifyFunction(F=<unavailable>) at Analysis.cpp:33:3                                                                                                   
   30                                                                                                                                                                                                         
   31   void SILAnalysis::verifyFunction(SILFunction *F) {                                                                                                                                                    
   32     // Only functions with bodies can be analyzed by the analysis.                                                                                                                                      
-> 33     assert(F->isDefinition() && "Can't analyze external functions");                                                                                                                                    
   34   }                                                                                                                                                                                                     
   35                                                                                                                                                                                                         
   36   SILAnalysis *swift::createDominanceAnalysis(SILModule *) {                                                                                                                                            
(lldb) up                                                                                                                                                                                                     
frame #&#8203;5: 0x00000000011e62b6 swift`swift::FunctionAnalysisBase<(anonymous namespace)::DifferentiableActivityInfo>::get(this=0x0000000007d7b280, f=0x0000000007cf6218) at Analysis.h:229:5                     
   226    /// Returns a function info structure for a specific function \p F.                                                                                                                                 
   227    FunctionInfoTy *get(SILFunction *f) {                                                                                                                                                               
   228      // Check that the analysis can handle this function.                                                                                                                                              
-> 229      verifyFunction(f);                                                                                                                                                                                
   230                                                                                                                                                                                                        
   231      auto &it = storage.FindAndConstruct(f);                                                                                                                                                           
   232      if (!it.second)                                                                                                                                                                                   
(lldb) up                                                                                                                                                                                                     
frame #&#8203;6: 0x00000000011d83ed swift`(anonymous namespace)::PrimalGen::run() [inlined] (anonymous namespace)::PrimalGen::performSynthesis(this=<unavailable>)::FunctionSynthesisItem) at Differentiation.cpp:24$
2:43                                                                                                                                                                                                          
   2459   auto &passManager = context.getPassManager();                                                                                                                                                       
   2460   auto *activityAnalysis =                                                                                                                                                                            
   2461       passManager.getAnalysis<DifferentiableActivityAnalysis>();                                                                                                                                      
-> 2462   auto &activityInfo = *activityAnalysis->get(item.original);                                                                                                                                         
   2463   // For debugging, dump the original function's activity analysis.                                                                                                                                   
   2464   LLVM_DEBUG(dumpActivityInfo(*item.original, item.task->getIndices(),                                                                                                                                
   2465                               activityInfo, getADDebugStream()));
@rxwei
Copy link
Member

rxwei commented Jan 29, 2019

I'm actually surprised this would even go into PrimalGen. This function already has a `@differentiable` attribute. So `emitAssociatedFunctionReference()` in the AD pass should be able to find the JVP and VJP. Investigating.

@rxwei
Copy link
Member

rxwei commented Jan 29, 2019

Ok, this is actually a bug in SILFunctionBuilder::addFunctionAttributes. It checks whether two the to-be-built function and the decl are in the different Swift modules, but this case is where two are in the same Swift module but different SIL modules.

@rxwei
Copy link
Member

rxwei commented Jan 29, 2019

I'll have a fix later today. It's not complicated, but it just needs multiple parts to coordinate well (TBDGen, AD pass, and SILGen).

@rxwei
Copy link
Member

rxwei commented Jan 30, 2019

Update: I became swamped with a memory bug. Dan is taking on this.

My WIP tree is https://github.com/rxwei/swift/tree/dont-analyse-external-functions. There are four remaining steps:

  1. make JVP/VJP serialized when the original is
  2. make primal/adjoint linkage be public when the original is serialized
  3. Extend AutoDiffAssociatedFunctionIdentifier to also accept primal/adjoint. Maybe rename this to AutoDiffFunctionIdentifier
  4. in TBDGen where it generates symbols for jvp and vjp, also generate symbols for primal/adjoint when the original function is @inlinable.

@dan-zheng
Copy link
Collaborator

Resolved in #22280

Fix description:

Previously, SILFunctionBuilder created SILDifferentiableAttr with the
expected mangled JVP/VJP names even if they weren't specified in the original
@differentiable attribute, when the original function is from a different
SIL module. This enabled the differentiation pass to find externally defined
JVP/VJPs.

However, this logic was hacky/unreliable. Now, SILFunctionBuilder simply
propagates JVP/VJP names if they are specified in the original @differentiable
attribute. The differentiation pass then declares external JVP/VJP functions
when the original function is an external declaration.

Note that the fix did not involve changing linkage or generating symbols for primal/adjoint. Those changes are staged at https://github.com/dan-zheng/swift/tree/autodiff-internal-function-identifer for revisiting.

@swift-ci swift-ci transferred this issue from apple/swift-issues Apr 25, 2022
This issue was closed.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bug A deviation from expected or documented behavior. Also: expected but undesirable behavior. swift for tensorflow
Projects
None yet
Development

No branches or pull requests

3 participants