Method onCompute

source/geometry/GeometryTensorArray.cpp:425–481

class GeometryTensorArrayConcat : public GeometryComputer {
public:
    virtual bool onCompute(const Op* op, const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs,
                           Context& context, CommandBuffer& res) const override {
        auto output    = outputs[0];
        auto outDes = TensorUtils::getDescribe(output);
        outDes->memoryType = Tensor::InsideDescribe::MEMORY_VIRTUAL;
        outDes->regions.clear();
        auto attr = TensorUtils::getDescribe(inputs[1])->tensorArrayAttr;
        auto tpParam = op->main_as_TensorArray();
        int concatAxis = tpParam->axis(), newAxis = tpParam->new_axis();
        int outputDimensions = output->dimensions();
        concatAxis = (concatAxis + outputDimensions) % outputDimensions;
        int outside = 1;
        int inside = 1;
        for (int i=0; i<concatAxis; ++i) {
            outside *= output->length(i);
        }
        for (int i=concatAxis+1; i<output->dimensions(); ++i) {
            inside *= output->length(i);
        }
        int concatFinal = output->length(concatAxis);
        for (int i = 0, concatSum = 0, concatLast = -1; i < attr->arraySize; ++i) {
            int shapeIndex = i;
            if (attr->isIdenticalShape) {
                shapeIndex = 0;
            }
            int concatLen = 1;
            if (newAxis == 0) {
                concatLen = attr->elemShape[shapeIndex][concatAxis];
            }
            if (1 == outside && outDes->regions.size() > 0) {
                // If outside is 1, fuse to one region
                outDes->regions[outDes->regions.size() - 1].size[2] += inside * concatLen;
                concatSum += concatLen;
                continue;
            }
            if (concatLast == concatLen) {
                // Fuse to last region
                outDes->regions[outDes->regions.size() - 1].size[0] += 1;
                concatSum += concatLen;
                continue;
            }
            Tensor::InsideDescribe::Region reg;
            reg.origin = inputs[1];
            reg.src.offset = inside * outside * concatSum;
            reg.src.stride[0] = outside * inside * concatLen;
            reg.src.stride[1] = inside * concatLen;
            reg.dst.offset = inside * concatSum;
            reg.dst.stride[0] = inside * concatLen;
            reg.dst.stride[1] = inside * concatFinal;
            reg.size[1] = outside;
            reg.size[2] = inside * concatLen;
            outDes->regions.push_back(reg);
            concatSum += concatLen;
            concatLast = concatLen;
        }
        return true;
    }
};