[DynamicCompilation] Add _AnyTensorHandle and Swift-C round trip entry points #19529
Conversation
|
@swift-ci please test tensorflow |
|
@swift-ci please test tensorflow Linux GPU |
| @@ -44,6 +42,20 @@ public final class TensorHandle<Scalar : AccelerableByTensorFlow> { | |||
| self.cTensorHandle = cTensorHandle | |||
| } | |||
|
|
|||
| deinit { | |||
| debugLog("De-initializing TensorHandle.") | |||
There was a problem hiding this comment.
Actually I'm going to move initializers back into TensorHandle<T> since only that would be correct, and _AnyTensorHandle is just a type-erased class for easier IRGen. So it's always a TensorHandle.
| cHandle = TFE_TensorHandleCopySharingTensor(cHandle, status) | ||
| expectEqual(TF_GetCode(status), TF_OK) | ||
| TF_DeleteStatus(status) | ||
| let anyHandle = _CreateTensorHandleFromC(cHandle, TF_FLOAT) |
There was a problem hiding this comment.
instead of passing in TF_FLOAT, can we compute that via TFE_TensorHandleDataType() within TFE_TensorHandleCopySharingTensor()? This way within IRGen, we won't need to remember the dtype for each tensor handle. (Otherwise we may want to return that dtype in _GetCTensorHandleFromSwift(), which complicates things and is not necessary for IRGen).
| @@ -99,6 +99,11 @@ static bool isUserIgnoredByPartitioning(SILInstruction *inst) { | |||
| /// unexpected. | |||
| static CanType getSingleElementDeclFieldType(NominalTypeDecl *decl) { | |||
| auto *field = tf::getFieldIfContainsSingleField(decl); | |||
| if (!field) | |||
There was a problem hiding this comment.
please extend the API comment to describe/cover this new case.
| let status = TF_NewStatus() | ||
| // We must do a copy, i.e. a retain on the tensor handle, to make sure it won't | ||
| // get double-free'd when both `zero` and `anyHandle` below go out of scope. | ||
| cHandle = TFE_TensorHandleCopySharingTensor(cHandle, status) |
There was a problem hiding this comment.
hmm this suggests using TensorHandle<...>(owning: cHandle) within _GetCTensorHandleFromSwift() is not correct. A c handle should not be owned by 2 or more TensorHandle objects.
In IRgen, for a given tensor handle x, we may need to call _GetCTensorHandleFromSwift() on it multiple times, if x is consumed by multiple tfops. We don't want IRGen code to call TFE_TensorHandleCopySharingTensor to rebalance the ref count.
Can we call TFE_TensorHandleCopySharingTensor() within _GetCTensorHandleFromSwift()?
There was a problem hiding this comment.
Will the corresponding TF call emitted by IRGen consume the retain count?
There was a problem hiding this comment.
As we discussed in person, making _GetCTensorHandleFromSwift not make a copy is the correct behavior since the lifetime of each graph_op argument is guaranteed by construction.
This test case is irrelevant because it's making a new TensorHandle own the extracted CTensorHandle, which is for testing purpose and will not be emitted by IRGen.
|
@swift-ci please test tensorflow |
Introduce a scalar-type-erased
_AnyTensorHandleas a base class forTensorHandle<T>, along with two runtime entry points that assistgraph_oplowering in IRGen._swift_tfc_GetCTensorHandleFromSwift_swift_tfc_CreateTensorHandleFromCIt is unfortunate that this function has to switch on all supported dtypes, but that's the only way to create a concrete
TensorHandle<T>type. TheswitchindumpTensorContentsshould be unified with this.