def _GetMut2OperandBlobObjects(self, op_attribute, parallel_desc_sym, bn_in_op2blob_object={}): mut2_operand_blob_objects = [] def GetOutBlobParallelDescSymbol(obn): parallel_signature = op_attribute.parallel_signature bn2symbol_id = parallel_signature.bn_in_op2parallel_desc_symbol_id if obn in bn2symbol_id: return symbol_storage.GetSymbol4Id(bn2symbol_id[obn]) else: return parallel_desc_sym for obn in op_attribute.output_bns: obn2modifier = op_attribute.arg_modifier_signature.obn2output_blob_modifier if obn2modifier[obn].header_infered_before_compute: continue obn_sym = self.GetSymbol4String(obn) op_arg_parallel_attr = op_arg_util.GetOpArgParallelAttribute( GetOutBlobParallelDescSymbol(obn), op_attribute, obn) op_arg_blob_attr = op_arg_util.GetOpArgBlobAttribute( op_attribute, obn) out_blob_object = self._NewBlobObject(op_arg_parallel_attr, op_arg_blob_attr) bn_in_op2blob_object[obn] = out_blob_object mut2_operand_blob_objects.append((obn_sym, out_blob_object)) return mut2_operand_blob_objects
def DistributeConcatOrAdd(op_attribute, parallel_conf, blob_register): op_parallel_desc_sym = symbol_storage.GetSymbol4Id( op_attribute.parallel_signature.op_parallel_desc_symbol_id) parallel_size = len(op_attribute.input_bns) op_arg_parallel_attr = op_arg_util.GetOpArgParallelAttribute( op_parallel_desc_sym, op_attribute, "out") op_arg_blob_attr = op_arg_util.GetOpArgBlobAttribute(op_attribute, "out") parallel_sig = op_attribute.parallel_signature.bn_in_op2parallel_desc_symbol_id def GetInBlobObject(builder, i, bn_in_op2blob_object): ibn = "in_%s" % i 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 BuildInstruction(builder): with blob_register.BnInOp2BlobObjectScope( op_attribute) as bn_in_op2blob_object: def GetPhysicalInBlob(i): return GetInBlobObject(builder, i, bn_in_op2blob_object) in_blob_objects = [ GetPhysicalInBlob(i) for i in range(parallel_size) ] bn_in_op2blob_object[ "out"] = builder.PackPhysicalBlobsToLogicalBlob( in_blob_objects, op_arg_parallel_attr, op_arg_blob_attr) vm_util.LogicalRun(BuildInstruction)
def MakeLazyRefBlobObject(self, interface_op_name): sess = session_ctx.GetDefaultSession() op_attribute = sess.OpAttribute4InterfaceOpName(interface_op_name) assert len(op_attribute.output_bns) == 1 obn = op_attribute.output_bns[0] blob_parallel_desc_sym_id = op_attribute.parallel_signature.bn_in_op2parallel_desc_symbol_id[ obn] blob_parallel_desc_sym = symbol_storage.GetSymbol4Id( blob_parallel_desc_sym_id) op_arg_parallel_attr = op_arg_util.GetOpArgParallelAttribute( blob_parallel_desc_sym, op_attribute, obn) op_arg_blob_attr = op_arg_util.GetOpArgBlobAttribute(op_attribute, obn) blob_object = self._NewBlobObject(op_arg_parallel_attr, op_arg_blob_attr) self._LazyReference(blob_object, interface_op_name) return blob_object
def _GetMut1OperandBlobObjects( self, op_attribute, parallel_desc_sym, bn_in_op2blob_object={} ): mut1_operand_blob_objects = [] for ibn in op_attribute.input_bns: ibn2modifier = op_attribute.arg_modifier_signature.ibn2input_blob_modifier if not ibn2modifier[ibn].is_mutable: continue ibn_sym = self.GetSymbol4String(ibn) ref_blob_object = bn_in_op2blob_object[ibn] mut1_operand_blob_objects.append((ibn_sym, ref_blob_object)) def GetOutBlobParallelDescSymbol(obn): parallel_signature = op_attribute.parallel_signature bn2symbol_id = parallel_signature.bn_in_op2parallel_desc_symbol_id if obn in bn2symbol_id: return symbol_storage.GetSymbol4Id(bn2symbol_id[obn]) else: return parallel_desc_sym def OutputBns(): obn2modifier = op_attribute.arg_modifier_signature.obn2output_blob_modifier for obn in op_attribute.output_bns: if obn2modifier[obn].header_infered_before_compute: yield obn for tmp_bn in op_attribute.tmp_bns: yield tmp_bn for obn in OutputBns(): obn_sym = self.GetSymbol4String(obn) op_arg_parallel_attr = op_arg_util.GetOpArgParallelAttribute( GetOutBlobParallelDescSymbol(obn), op_attribute, obn ) op_arg_blob_attr = op_arg_util.GetOpArgBlobAttribute(op_attribute, obn) out_blob_object = self._NewBlobObject( op_arg_parallel_attr, op_arg_blob_attr ) lbi = op_attribute.arg_signature.bn_in_op2lbi[obn] bn_in_op2blob_object[obn] = out_blob_object mut1_operand_blob_objects.append((obn_sym, out_blob_object)) return mut1_operand_blob_objects