def GetInterfaceBlobValue(op_name):
flow.sync_default_session()
sess = session_ctx.GetDefaultSession()
job_name = sess.JobName4InterfaceOpName(op_name)
def AsyncGetInterfaceBlobValue(Yield):
def build(builder):
blob_object = builder.MakeLazyRefBlobObject(op_name)
lbi = logical_blob_id_util.LogicalBlobId()
lbi.op_name = op_name
op_attribute = sess.OpAttribute4InterfaceOpName(op_name)
assert len(op_attribute.output_bns) == 1
lbi.blob_name = op_attribute.output_bns[0]
if blob_object.op_arg_parallel_attr.is_mirrored():
remote_blob = remote_blob_util.EagerMirroredBlob(
lbi, blob_object, job_name)
else:
remote_blob = remote_blob_util.EagerConsistentBlob(
lbi, blob_object, job_name)
if blob_object.op_arg_blob_attr.is_tensor_list:
value = remote_blob.numpy_list()
else:
value = remote_blob.numpy()
Yield(value)
vm_util.LogicalRun(build)
return async_util.Await(1, AsyncGetInterfaceBlobValue)[0]
def CreateBlob():
job_name = sess.JobName4InterfaceOpName(op_name)
def Build(builder, Yield):
blob_object = _GetInterfaceBlobObject(builder, op_name)
lbi = lbi_util.LogicalBlobId()
lbi.set_op_name(op_name)
op_attribute = sess.OpAttribute4InterfaceOpName(op_name)
assert len(op_attribute.output_bns) == 1
lbi.set_blob_name(op_attribute.output_bns[0])
if blob_object.op_arg_parallel_attr.is_mirrored():
remote_blob = oneflow._oneflow_internal.EagerMirroredBlob(
lbi, blob_object, blob_register, job_name)
else:
remote_blob = oneflow._oneflow_internal.EagerConsistentBlob(
lbi, blob_object, blob_register, job_name)
Yield(remote_blob)
def AsyncGetInterfaceBlob(Yield):
oneflow._oneflow_internal.deprecated.LogicalRun(
lambda builder: Build(builder, Yield))
blob = async_util.Await(1, AsyncGetInterfaceBlob)[0]
return blob
def GetInterfaceBlobValue(op_name):
sync_default_session_if_normal()
sess = session_ctx.GetDefaultSession()
job_name = sess.JobName4InterfaceOpName(op_name)
def AsyncGetInterfaceBlobValue(Yield):
def build(builder):
blob_object = GetEagerInterfaceBlob(op_name).blob_object
lbi = lbi_util.LogicalBlobId()
lbi.set_op_name(op_name)
op_attribute = sess.OpAttribute4InterfaceOpName(op_name)
assert len(op_attribute.output_bns) == 1
lbi.set_blob_name(op_attribute.output_bns[0])
if not isinstance(lbi, lbi_util.LogicalBlobId):
cfg_lbi = lbi_util.LogicalBlobId()
cfg_lbi.set_op_name(lbi.op_name)
cfg_lbi.set_blob_name(lbi.blob_name)
lbi = cfg_lbi
if blob_object.op_arg_parallel_attr.is_mirrored():
remote_blob = oneflow._oneflow_internal.EagerMirroredBlob(
lbi, blob_object, blob_register, job_name)
else:
remote_blob = oneflow._oneflow_internal.EagerConsistentBlob(
lbi, blob_object, blob_register, job_name)
value = remote_blob.numpy()
Yield(value)
oneflow._oneflow_internal.deprecated.LogicalRun(build)
return async_util.Await(1, AsyncGetInterfaceBlobValue)[0]
def FeedValueToInterfaceBlob(op_name, ndarray):
flow.sync_default_session()
def AsyncFeedValueToInterfaceBlob(Yield):
def build(builder):
blob_object = builder.MakeLazyRefBlobObject(op_name)
if blob_object.op_arg_blob_attr.is_tensor_list:
input_blob_def = input_blob_def_util.MirroredTensorListDef(
[x.shape for x in ndarray],
dtype=dtype_util.convert_numpy_dtype_to_oneflow_dtype(
ndarray.dtype),
)
elif blob_object.op_arg_parallel_attr.is_mirrored():
input_blob_def = input_blob_def_util.MirroredTensorDef(
ndarray.shape,
dtype=dtype_util.convert_numpy_dtype_to_oneflow_dtype(
ndarray.dtype),
)
else:
input_blob_def = input_blob_def_util.FixedTensorDef(
ndarray.shape,
dtype=dtype_util.convert_numpy_dtype_to_oneflow_dtype(
ndarray.dtype),
)
push_util.FeedValueToEagerBlob(blob_object, input_blob_def,
ndarray)
Yield()
vm_util.LogicalRun(build)
async_util.Await(1, AsyncFeedValueToInterfaceBlob)
def _LogicalSlice(
input_blob_object: oneflow_api.BlobObject,
start: Sequence[int],
stop: Sequence[int],
scope_symbol_id: int,
) -> np.ndarray:
"""
Construct a logical_slice op and run it by oneflow eager,
return the sliced result as a numpy ndarray
"""
op_name = id_util.UniqueStr(OP_PREFIX)
def AsyncSlice(Yield):
def build(builder):
op_conf = op_conf_pb.OperatorConf()
# device_tag doesn't matter for logical_slice op
device_tag = oneflow.current_scope(
).device_parallel_desc_symbol.device_tag
op_conf.device_tag = device_tag
op_conf.name = op_name
op_conf.user_conf.op_type_name = "logical_slice"
op_conf.user_conf.input["x"].s.append("{}/x_0".format(op_name))
op_conf.user_conf.output["y"].s.append("{}/y_0".format(op_name))
parallel_conf = input_blob_object.parallel_desc_symbol.parallel_conf
op_conf.user_conf.attr["parallel_conf"].at_string = str(
parallel_conf)
op_conf.user_conf.attr["start"].at_list_int64.val[:] = start
op_conf.user_conf.attr["stop"].at_list_int64.val[:] = stop
op_conf.user_conf.attr["step"].at_list_int64.val[:] = [
1
] * len(start)
bn_in_op2blob_object = oneflow_api.deprecated.BnInOp2BlobObject()
bn_in_op2blob_object["x_0"] = input_blob_object
op_attribute = op_infer_util.Infer(op_conf, bn_in_op2blob_object,
scope_symbol_id)
cfg_op_attribute = oneflow_api.deprecated.MakeOpAttributeByString(
str(op_attribute))
builder.StatelessCall(
cfg_op_attribute,
parallel_conf,
bn_in_op2blob_object,
boxing_util.BoxingTo,
)
Yield(bn_in_op2blob_object["y_0"])
oneflow_api.deprecated.LogicalRun(build)
lbi = lbi_util.LogicalBlobId()
lbi.set_op_name(op_name)
lbi.set_blob_name(op_name)
blob_object = async_util.Await(1, AsyncSlice)[0]
blob = oneflow_api.EagerConsistentBlob(
lbi,
blob_object=blob_object,
blob_register=blob_register,
job_name=FAKE_JOB_NAME,
)
return blob.numpy()
def FeedValueToInterfaceBlob(op_name, ndarray):
sync_default_session_if_normal()
def AsyncFeedValueToInterfaceBlob(Yield):
def build(builder):
blob_object = GetEagerInterfaceBlob(op_name).blob_object
FeedValueToInterfaceBlobObject(blob_object, ndarray)
Yield()
oneflow._oneflow_internal.deprecated.LogicalRun(build)
async_util.Await(1, AsyncFeedValueToInterfaceBlob)
def _FetchBlobHeader(blob_object):
def AsyncFetchBlobHeader(Yield):
fetcher = _MakeFetcherEagerPhysicalBlobHeaderFromOfBlob(Yield)
def BuildFetchBlobHeaderInstruction(builder):
builder.FetchBlobHeader(blob_object, fetcher)
builder.InsertRemoveForeignCallbackInstruction(
blob_object.object_id, fetcher)
vm_util.PhysicalRun(BuildFetchBlobHeaderInstruction)
return async_util.Await(1, AsyncFetchBlobHeader)[0]
def FetchTensorBlobAsNumpyList(parallel_size, blob_object):
def AsyncFetchBlobBody(Yield):
fetcher = _MakeFetcherEagerBlobBodyAsNumpyFromOfBlob(Yield)
def BuildFetchBlobBodyInstruction(builder):
builder.FetchBlobBody(blob_object, fetcher)
builder.InsertRemoveForeignCallbackInstruction(
blob_object.object_id, fetcher)
vm_util.PhysicalRun(BuildFetchBlobBodyInstruction)
return async_util.Await(parallel_size, AsyncFetchBlobBody)
def FeedValueToInterfaceBlob(op_name, ndarray):
flow.sync_default_session()
def AsyncFeedValueToInterfaceBlob(Yield):
def build(builder):
blob_object = GetEagerInterfaceBlob(op_name).blob_object
FeedValueToInterfaceBlobObject(blob_object, ndarray)
Yield()
vm_util.LogicalRun(build)
async_util.Await(1, AsyncFeedValueToInterfaceBlob)
def _LogicalSlice(input_blob: EagerBlobTrait, start: Sequence[int],
stop: Sequence[int]) -> np.ndarray:
"""
Construct a logical_slice op and run it by oneflow eager,
return the sliced result as a numpy ndarray
"""
op_name = id_util.UniqueStr(OP_PREFIX)
def AsyncSlice(Yield):
def build(builder):
op_conf = op_conf_pb.OperatorConf()
# device_tag doesn't matter for logical_slice op
device_tag = oneflow.current_scope(
).device_parallel_desc_symbol.device_tag
op_conf.device_tag = device_tag
op_conf.name = op_name
op_conf.user_conf.op_type_name = "logical_slice"
op_conf.user_conf.input["x"].s.append("{}/x_0".format(op_name))
op_conf.user_conf.output["y"].s.append("{}/y_0".format(op_name))
input_blob_object = input_blob.blob_object
parallel_conf = input_blob_object.parallel_desc_symbol.parallel_conf
op_conf.user_conf.attr["parallel_conf"].at_string = str(
parallel_conf)
op_conf.user_conf.attr["start"].at_list_int64.val[:] = start
op_conf.user_conf.attr["stop"].at_list_int64.val[:] = stop
op_conf.user_conf.attr["step"].at_list_int64.val[:] = [
1
] * len(start)
bn_in_op2blob_object = {"x_0": input_blob_object}
scope_symbol_id = _GetScopeSymbolIdFromEagerBlob(input_blob)
op_attribute = op_infer_util.Infer(op_conf, bn_in_op2blob_object,
scope_symbol_id)
builder.StatelessCall(
op_attribute,
parallel_conf=parallel_conf,
bn_in_op2blob_object=bn_in_op2blob_object,
)
Yield(bn_in_op2blob_object["y_0"])
vm_util.LogicalRun(build)
lbi = logical_blob_id_util.LogicalBlobId()
lbi.op_name = op_name
lbi.blob_name = op_name
blob_object = async_util.Await(1, AsyncSlice)[0]
blob = remote_blob_util.EagerConsistentBlob(lbi,
blob_object=blob_object,
job_name=FAKE_JOB_NAME)
return blob.numpy()
def _FetchBlobHeader(blob_object):
def AsyncFetchBlobHeader(Yield):
fetcher = _MakeFetcherEagerPhysicalBlobHeaderFromOfBlob(Yield)
def BuildFetchBlobHeaderInstruction(builder):
builder.FetchBlobHeader(
blob_object, python_callback.GetIdForRegisteredCallback(fetcher)
)
builder.InsertRemoveForeignCallbackInstruction(
blob_object.object_id,
python_callback.GetIdForRegisteredCallback(fetcher),
)
oneflow_api.deprecated.PhysicalRun(BuildFetchBlobHeaderInstruction)
return async_util.Await(1, AsyncFetchBlobHeader)[0]
def FetchTensorBlobAsNumpyList(parallel_size, blob_object):
def AsyncFetchBlobBody(Yield):
fetcher = _MakeFetcherEagerBlobBodyAsNumpyFromOfBlob(Yield)
def BuildFetchBlobBodyInstruction(builder):
builder.FetchBlobBody(
blob_object, python_callback.GetIdForRegisteredCallback(fetcher)
)
builder.InsertRemoveForeignCallbackInstruction(
blob_object.object_id,
python_callback.GetIdForRegisteredCallback(fetcher),
)
oneflow_api.deprecated.PhysicalRun(BuildFetchBlobBodyInstruction)
return async_util.Await(parallel_size, AsyncFetchBlobBody)
def _access_blob_by_callback(local_tensor, callback, modifier):
def AsyncAccess(Yield):
def MakeAccessor(Yield):
def AccessOfBlobPtr(ofblob_ptr):
ofblob = ofblob_util.OfBlob(ofblob_ptr)
Yield(callback(ofblob))
return AccessOfBlobPtr
accessor = MakeAccessor(Yield)
def BuildInstruction(builder):
builder.AccessBlobByCallback(local_tensor, accessor, modifier)
flow._oneflow_internal.deprecated.PhysicalRun(BuildInstruction)
return async_util.Await(1, AsyncAccess)[0]
def CreateBlob():
job_name = sess.JobName4InterfaceOpName(op_name)
def Build(builder, Yield):
blob_object = _GetInterfaceBlobObject(builder, op_name)
lbi = logical_blob_id_util.LogicalBlobId()
lbi.op_name = op_name
op_attribute = sess.OpAttribute4InterfaceOpName(op_name)
assert len(op_attribute.output_bns) == 1
lbi.blob_name = op_attribute.output_bns[0]
if blob_object.op_arg_parallel_attr.is_mirrored():
remote_blob = remote_blob_util.EagerMirroredBlob(
lbi, blob_object, job_name)
else:
remote_blob = remote_blob_util.EagerConsistentBlob(
lbi, blob_object, job_name)
Yield(remote_blob)
def AsyncGetInterfaceBlob(Yield):
vm_util.LogicalRun(lambda builder: Build(builder, Yield))
blob = async_util.Await(1, AsyncGetInterfaceBlob)[0]
return blob
