MCPcopy Create free account
hub / github.com/Oneflow-Inc/oneflow / ConcatSplitDataContent

Function ConcatSplitDataContent

oneflow/core/kernel/boxing_kernel.cpp:170–194  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

168}
169
170void ConcatSplitDataContent(ep::Stream* stream,
171 const std::function<Blob*(const std::string&)>& BnInOp2Blob,
172 const PbRpf<std::string>& concat_bns, int32_t concat_axis,
173 const PbRpf<std::string>& split_bns, int32_t split_axis) {
174 DataContentDesc in_desc(BnInOp2Blob, &concat_bns, concat_axis);
175 DataContentDesc out_desc(BnInOp2Blob, &split_bns, split_axis);
176 CHECK_EQ(in_desc.TotalElemNum(), out_desc.TotalElemNum());
177 CHECK_EQ(in_desc.OneElemSize(), out_desc.OneElemSize());
178 static const size_t min_byte_one_part = 128;
179 int32_t part_num = in_desc.TotalElemNum() * in_desc.OneElemSize() / min_byte_one_part;
180 part_num = std::min(part_num, Singleton<ThreadPool>::Get()->thread_num());
181 if (part_num >= 2) {
182 BlockingCounter bc(part_num);
183 FOR_RANGE(int32_t, part_id, 0, part_num) {
184 Singleton<ThreadPool>::Get()->AddWork(
185 [stream, &in_desc, &out_desc, part_id, &part_num, &bc]() {
186 ConcatSplitPartDataContent(stream, in_desc, out_desc, part_id, part_num);
187 bc.Decrease();
188 });
189 }
190 bc.WaitForeverUntilCntEqualZero();
191 } else {
192 ConcatSplitPartDataContent(stream, in_desc, out_desc, 0, 1);
193 }
194}
195
196} // namespace
197

Callers 1

ForwardDataContentMethod · 0.85

Calls 6

GetFunction · 0.85
TotalElemNumMethod · 0.80
OneElemSizeMethod · 0.80
thread_numMethod · 0.80

Tested by

no test coverage detected