| 132 | // tensor(index < seq_length) will insert instead of overwrite when onnxInsert=true |
| 133 | GeometryTensorArrayWrite(bool insertMode) : mInsertMode(insertMode) { } |
| 134 | virtual bool onCompute(const Op* op, const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs, |
| 135 | Context& context, CommandBuffer& res) const override { |
| 136 | auto tensorArrayInput = inputs[3]; |
| 137 | auto inDes = TensorUtils::getDescribe(tensorArrayInput); |
| 138 | if (inDes->tensorArrayAttr == nullptr) { |
| 139 | MNN_ASSERT(false); |
| 140 | return false; |
| 141 | } |
| 142 | auto output = outputs[0]; |
| 143 | auto outDes = TensorUtils::getDescribe(output); |
| 144 | outDes->memoryType = Tensor::InsideDescribe::MEMORY_VIRTUAL; |
| 145 | |
| 146 | int oldSize = inDes->tensorArrayAttr->arraySize; |
| 147 | int writeIndex = inputs[1]->host<uint32_t>()[0]; |
| 148 | // mInsertMode=true mean onnx mode, which position tensor is int32 instead of uint32 |
| 149 | if (mInsertMode) { |
| 150 | writeIndex = inputs[1]->host<int32_t>()[0]; |
| 151 | writeIndex += (writeIndex < 0 ? inDes->tensorArrayAttr->arraySize: 0); // [-n, n] |
| 152 | } |
| 153 | auto elemSize = getElemSize(output, writeIndex); |
| 154 | outDes->regions.clear(); |
| 155 | // support insertMode=true/false, easier to understand |
| 156 | int regionSize = (writeIndex > 0) + 1 + (writeIndex < outDes->tensorArrayAttr->arraySize - 1); |
| 157 | outDes->regions.reserve(regionSize); |
| 158 | /* |
| 159 | src: [leftData][writeIndex][rightData] |
| 160 | dst: [leftData][writeTensor][rightData] |
| 161 | */ |
| 162 | // 1. write Tensor to dst TensorArray [must] |
| 163 | if (elemSize.second == 0) { |
| 164 | return true; |
| 165 | } |
| 166 | { |
| 167 | Tensor::InsideDescribe::Region writeTensorRegion; |
| 168 | writeTensorRegion.origin = inputs[2]; |
| 169 | writeTensorRegion.src.offset = 0; |
| 170 | writeTensorRegion.src.stride[0] = 1; |
| 171 | writeTensorRegion.src.stride[1] = 1; |
| 172 | writeTensorRegion.src.stride[2] = 1; |
| 173 | writeTensorRegion.dst.offset = elemSize.first; |
| 174 | writeTensorRegion.dst.stride[0] = 1; |
| 175 | writeTensorRegion.dst.stride[1] = 1; |
| 176 | writeTensorRegion.dst.stride[2] = 1; |
| 177 | writeTensorRegion.size[0] = elemSize.second; |
| 178 | writeTensorRegion.size[1] = 1; |
| 179 | writeTensorRegion.size[2] = 1; |
| 180 | MNN_ASSERT(elemSize.second > 0); |
| 181 | outDes->regions.emplace_back(std::move(writeTensorRegion)); |
| 182 | } |
| 183 | if (regionSize == 1) { |
| 184 | return true; |
| 185 | } |
| 186 | // first write data, set pre zero |
| 187 | bool firstWrite = isFirstWrite(inDes); |
| 188 | if (firstWrite) { |
| 189 | auto type = tensorArrayInput->getType(); |
| 190 | auto zeroConst = context.allocConst(op, {}, type); |
| 191 | if (type == halide_type_of<float>()) { |
nothing calls this directly
no test coverage detected