tqchen commented on code in PR #16557: URL: https://github.com/apache/tvm/pull/16557#discussion_r1494590316
########## src/runtime/hexagon/hexagon_device_api.h: ########## @@ -199,6 +215,9 @@ class HexagonDeviceAPI final : public DeviceAPI { //! \brief Hexagon power manager std::unique_ptr<HexagonPowerManager> runtime_power_manager; + + //! \brief NDArray base -> Physical Shape map + std::map<void*, PhysicalShape> ndarray_physical_shape; Review Comment: usually we prefer unordered_map over map. my main comment is to embed the metadata like physicalshape into the Buffer data structure in void* data field, so we can remove this map altogether. ########## src/runtime/hexagon/hexagon_device_api.cc: ########## @@ -165,23 +179,44 @@ void HexagonDeviceAPI::CopyDataFromTo(DLTensor* from, DLTensor* to, TVMStreamHan << "HexagonDeviceAPI::CopyDataFromTo before initializing resources. " << "Please call HexagonDeviceAPI::AcquireResources"; - auto lookup_hexagon_buffer = [this](void* ptr) -> HexagonBuffer* { - return runtime_hexbuffs->FindHexagonBuffer(ptr); - }; + auto numBytes = GetDataSize(*from); + + size_t FlatShape = 1; + for (auto i = 0; i < from->ndim; ++i) FlatShape *= from->shape[i]; + + PhysicalShape source_shape = {1, 1, FlatShape}; + PhysicalShape dest_shape = {1, 1, FlatShape}; + auto it1 = ndarray_physical_shape.find(from->data); + if (it1 != ndarray_physical_shape.end()) source_shape = it1->second; + size_t src_rank = source_shape.ndim; + void* src_start = get_data_start(from); + void* dst_start = get_data_start(to); + BufferSet src((src_rank == 1) ? &(src_start) : static_cast<void**>(src_start), + source_shape.nblocks, numBytes / source_shape.nblocks); + auto it2 = ndarray_physical_shape.find(to->data); + if (it2 != ndarray_physical_shape.end()) dest_shape = it2->second; + size_t dest_rank = dest_shape.ndim; + BufferSet dest((dest_rank == 1) ? &(dst_start) : static_cast<void**>(dst_start), + dest_shape.nblocks, numBytes / dest_shape.nblocks); + hexagon_buffer_copy_across_regions(dest, src, numBytes, (it1 != ndarray_physical_shape.end()), + (it2 != ndarray_physical_shape.end())); + return; +} - HexagonBuffer* hex_from_buf = lookup_hexagon_buffer(from->data); - HexagonBuffer* hex_to_buf = lookup_hexagon_buffer(to->data); +void HexagonDeviceAPI::SetPhysicalShape(const DLTensor* tensor, const int64_t ndim, + const int64_t* shape) { + PhysicalShape physical_shape = {static_cast<size_t>(ndim), static_cast<size_t>(shape[0]), + static_cast<size_t>(shape[1])}; + SetPhysicalShape(tensor->data, physical_shape); +} - if (hex_from_buf && hex_to_buf) { - hex_to_buf->CopyFrom(*hex_from_buf, GetDataSize(*from)); - } else if (hex_to_buf) { - hex_to_buf->CopyFrom(from->data, GetDataSize(*from)); - } else if (hex_from_buf) { - hex_from_buf->CopyTo(to->data, GetDataSize(*to)); - } else { - CHECK(false) << "CopyDataFromTo requested between src and dst which are not managed by the " - "hexagon device api."; - } +void HexagonDeviceAPI::SetPhysicalShape(const void* data, const PhysicalShape& physical_shape) { + auto it = ndarray_physical_shape.find(const_cast<void*>(data)); + if (it != ndarray_physical_shape.end()) Review Comment: always circle with {} in multiple line ifthenelse ########## src/runtime/relax_vm/builtin.cc: ########## @@ -410,6 +410,10 @@ TVM_REGISTER_GLOBAL("vm.builtin.reshape").set_body_typed([](NDArray data, ShapeT return data.CreateView(new_shape, data->dtype); }); +TVM_REGISTER_GLOBAL("vm.builtin.copy_tensor").set_body_typed([](NDArray src, NDArray dst) { Review Comment: copy_tensor_from_to ########## src/runtime/relax_vm/hexagon/builtin.cc: ########## @@ -57,6 +57,44 @@ TVM_REGISTER_GLOBAL("vm.builtin.hexagon.dma_wait") ICHECK(inflight_dma >= 0); tvm::runtime::hexagon::HexagonDeviceAPI::Global()->UserDMA()->Wait(queue_id, inflight_dma); }); + +NDArray AllocNDArrayFromOffsets(TVMArgValue vm_ptr, Storage storage, uint64_t offset, + Storage data_storage, ShapeTuple data_storage_offsets, + ShapeTuple logical_shape, ShapeTuple shape_2d, DLDataType dtype) { + auto* storageObj = storage.operator->(); + + auto* container = storageObj->CreateNDArrayContainer(offset, logical_shape, dtype); Review Comment: storage_obj ########## src/runtime/hexagon/hexagon_buffer.h: ########## @@ -195,6 +195,10 @@ struct BufferSet { size_t region_size_bytes; }; +void hexagon_buffer_copy_across_regions(const BufferSet& dest, const BufferSet& src, Review Comment: CamelCase, please document the behavior here -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@tvm.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org