def BoxingToSingleDevice(builder):
parallel_conf = placement_cfg.ParallelConf()
parallel_conf.set_device_tag(
blob_object.parallel_desc_symbol.device_tag)
parallel_conf.add_device_name("{}:{}".format(0, 0))
tmp_parallel_desc_symbol = builder.GetParallelDescSymbol(
parallel_conf)
tmp_op_arg_parallel_attr = oneflow_api.OpArgParallelAttribute(
tmp_parallel_desc_symbol,
str(blob_object.op_arg_parallel_attr.sbp_parallel),
str(blob_object.op_arg_parallel_attr.opt_mirrored_parallel
),
)
with oneflow.scope.placement(
self.parallel_conf.device_tag(),
list(self.parallel_conf.device_name()),
):
tmp_blob_object = boxing_util.BoxingTo(
builder, blob_object, tmp_op_arg_parallel_attr)
nonlocal consistent_blob_name
consistent_blob_name = "{}-consistent".format(
self.logical_blob_name)
if not blob_register.HasObject4BlobName(consistent_blob_name):
blob_register.SetObject4BlobName(consistent_blob_name,
tmp_blob_object)
def GetInBlobObject(builder, ibn, bn_in_op2blob_object):
origin_blob_object = bn_in_op2blob_object[ibn]
in_op_parallel_desc_sym = symbol_storage.GetSymbol4Id(
parallel_sig[ibn])
in_op_arg_parallel_attr = op_arg_util.GetOpArgParallelAttribute(
in_op_parallel_desc_sym, op_attribute, ibn)
return boxing_util.BoxingTo(builder, origin_blob_object,
in_op_arg_parallel_attr)
def GetInBlobObject(builder, ibn, bn_in_op2blob_object):
origin_blob_object = bn_in_op2blob_object[ibn]
in_op_parallel_desc_sym = oneflow_api.GetPlacementSymbol(
parallel_sig[ibn])
in_op_arg_parallel_attr = oneflow_api.GetOpArgParallelAttribute(
in_op_parallel_desc_sym, str(op_attribute), ibn)
return boxing_util.BoxingTo(builder, origin_blob_object,
in_op_arg_parallel_attr)
def BuildAssignInstruction(builder):
new_parallel_desc_symbol = boxing_util.TryReplaceDeviceTag(
builder, var_blob_object.parallel_desc_symbol, "cpu")
consumer_op_arg_parallel_attr = op_arg_util.OpArgParallelAttribute(
new_parallel_desc_symbol,
var_blob_object.op_arg_parallel_attr.sbp_parallel,
var_blob_object.op_arg_parallel_attr.opt_mirrored_parallel,
)
tmp_blob_object = boxing_util.BoxingTo(builder, value_blob_object,
consumer_op_arg_parallel_attr)
boxing_util.Assign(builder, var_blob_object, tmp_blob_object)
def Fetch(x_blob_object, op_arg_parallel_attr):
return boxing_util.BoxingTo(builder, x_blob_object,
op_arg_parallel_attr)
def FetchDelegateBlobObject(x_blob_object, op_arg_parallel_attr):
return boxing_util.BoxingTo(self, x_blob_object, op_arg_parallel_attr)
