| 373 | class GeometryTensorArraySplit : public GeometryComputer { |
| 374 | public: |
| 375 | virtual bool onCompute(const Op* op, const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs, |
| 376 | Context& context, CommandBuffer& res) const override { |
| 377 | auto shape = inputs[1]->shape(); |
| 378 | int splitAxis = (op->main_as_TensorArray()->axis() + shape.size()) % shape.size(); |
| 379 | auto outside = std::accumulate(shape.begin(), shape.begin() + splitAxis, 1, |
| 380 | [](int a, int b) { return a * b; }); |
| 381 | auto inside = std::accumulate(shape.begin() + splitAxis + 1, shape.end(), 1, |
| 382 | [](int a, int b) { return a * b; }); |
| 383 | |
| 384 | auto value = inputs[1], lengths = inputs[2]; |
| 385 | bool scalarSplit = (lengths->elementSize() == 1); |
| 386 | int totalLen = value->shape()[splitAxis]; |
| 387 | int splitNum = (scalarSplit ? UP_DIV(totalLen, lengths->host<int>()[0]) : lengths->length(0)); |
| 388 | auto output = outputs[0]; |
| 389 | auto outDes = TensorUtils::getDescribe(output); |
| 390 | outDes->memoryType = Tensor::InsideDescribe::MEMORY_VIRTUAL; |
| 391 | outDes->regions.clear(); |
| 392 | for (int i = 0, splitSum = 0, splitLast = -1; i < splitNum; ++i) { |
| 393 | int splitLen; |
| 394 | if (scalarSplit) { |
| 395 | splitLen = ALIMIN(lengths->host<int>()[0], totalLen - splitSum); |
| 396 | } else { |
| 397 | splitLen = lengths->host<int>()[i]; |
| 398 | } |
| 399 | if (splitLast == splitLen) { |
| 400 | outDes->regions[outDes->regions.size() - 1].size[0] += 1; |
| 401 | splitSum += splitLen; |
| 402 | splitLast = splitLen; |
| 403 | continue; |
| 404 | } |
| 405 | Tensor::InsideDescribe::Region reg; |
| 406 | reg.origin = value; |
| 407 | reg.src.offset = inside * splitSum; |
| 408 | reg.src.stride[0] = inside * splitLen; |
| 409 | reg.src.stride[1] = inside * shape[splitAxis]; |
| 410 | reg.dst.offset = inside * outside * splitSum; |
| 411 | reg.dst.stride[0] = inside * outside * splitLen; |
| 412 | reg.dst.stride[1] = inside * splitLen; |
| 413 | reg.size[1] = outside; |
| 414 | reg.size[2] = inside * splitLen; |
| 415 | outDes->regions.push_back(reg); |
| 416 | splitSum += splitLen; |
| 417 | splitLast = splitLen; |
| 418 | } |
| 419 | return true; |
| 420 | } |
| 421 | }; |
| 422 | |
| 423 | class GeometryTensorArrayConcat : public GeometryComputer { |