gemini-code-assist[bot] commented on code in PR #111: URL: https://github.com/apache/tvm-ffi/pull/111#discussion_r2427457462
########## tests/python/test_dlpack_exchange_api.py: ########## @@ -0,0 +1,207 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file to +# you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + + +from __future__ import annotations + +import pytest + +try: + import torch # type: ignore[no-redef] + + # Import tvm_ffi to load the DLPack exchange API extension + # This sets torch.Tensor.__c_dlpack_exchange_api__ + import tvm_ffi # noqa: F401 + from torch.utils import cpp_extension # type: ignore + from tvm_ffi import libinfo +except ImportError: + torch = None + +# Check if DLPack Exchange API is available +_has_dlpack_api = torch is not None and hasattr(torch.Tensor, "__c_dlpack_exchange_api__") + + [email protected](not _has_dlpack_api, reason="PyTorch DLPack Exchange API not available") +def test_dlpack_exchange_api() -> None: + assert torch is not None + + assert hasattr(torch.Tensor, "__c_dlpack_exchange_api__") + api_ptr = torch.Tensor.__c_dlpack_exchange_api__ + assert isinstance(api_ptr, int), "API pointer should be an integer" + assert api_ptr != 0, "API pointer should not be NULL" + + tensor = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4) + + source = """ + #include <torch/extension.h> + #include <dlpack/dlpack.h> + #include <memory> + + void test_dlpack_api(at::Tensor tensor, int64_t api_ptr_int, bool cuda_available) { + DLPackExchangeAPI* api = reinterpret_cast<DLPackExchangeAPI*>(api_ptr_int); + + // Test 1: API structure and version + { + TORCH_CHECK(api != nullptr, "API pointer is NULL"); + TORCH_CHECK(api->header.version.major == DLPACK_MAJOR_VERSION, + "Expected major version ", DLPACK_MAJOR_VERSION, ", got ", api->header.version.major); + TORCH_CHECK(api->header.version.minor == DLPACK_MINOR_VERSION, + "Expected minor version ", DLPACK_MINOR_VERSION, ", got ", api->header.version.minor); + TORCH_CHECK(api->managed_tensor_allocator != nullptr, + "managed_tensor_allocator is NULL"); + TORCH_CHECK(api->managed_tensor_from_py_object_no_sync != nullptr, + "managed_tensor_from_py_object_no_sync is NULL"); + TORCH_CHECK(api->managed_tensor_to_py_object_no_sync != nullptr, + "managed_tensor_to_py_object_no_sync is NULL"); + TORCH_CHECK(api->dltensor_from_py_object_no_sync != nullptr, + "dltensor_from_py_object_no_sync is NULL"); + TORCH_CHECK(api->current_work_stream != nullptr, + "current_work_stream is NULL"); + } + + // Test 2: managed_tensor_allocator + { + DLTensor prototype; + prototype.device.device_type = kDLCPU; + prototype.device.device_id = 0; + prototype.ndim = 3; + int64_t shape[3] = {3, 4, 5}; + prototype.shape = shape; + prototype.strides = nullptr; + DLDataType dtype; + dtype.code = kDLFloat; + dtype.bits = 32; + dtype.lanes = 1; + prototype.dtype = dtype; + prototype.data = nullptr; + prototype.byte_offset = 0; + + DLManagedTensorVersioned* out_tensor = nullptr; + int result = api->managed_tensor_allocator(&prototype, &out_tensor, nullptr, nullptr); + TORCH_CHECK(result == 0, "Allocator failed with code ", result); + TORCH_CHECK(out_tensor != nullptr, "Allocator returned NULL"); + TORCH_CHECK(out_tensor->dl_tensor.ndim == 3, "Expected ndim 3, got ", out_tensor->dl_tensor.ndim); + TORCH_CHECK(out_tensor->dl_tensor.shape[0] == 3, "Expected shape[0] = 3, got ", out_tensor->dl_tensor.shape[0]); + TORCH_CHECK(out_tensor->dl_tensor.shape[1] == 4, "Expected shape[1] = 4, got ", out_tensor->dl_tensor.shape[1]); + TORCH_CHECK(out_tensor->dl_tensor.shape[2] == 5, "Expected shape[2] = 5, got ", out_tensor->dl_tensor.shape[2]); + TORCH_CHECK(out_tensor->dl_tensor.dtype.code == kDLFloat, "Expected dtype code kDLFloat, got ", out_tensor->dl_tensor.dtype.code); + TORCH_CHECK(out_tensor->dl_tensor.dtype.bits == 32, "Expected dtype bits 32, got ", out_tensor->dl_tensor.dtype.bits); + TORCH_CHECK(out_tensor->dl_tensor.device.device_type == kDLCPU, "Expected device type kDLCPU, got ", out_tensor->dl_tensor.device.device_type); + if (out_tensor->deleter) { + out_tensor->deleter(out_tensor); + } + } + + // Test 3: managed_tensor_from_py_object_no_sync + { + std::unique_ptr<PyObject, decltype(&Py_DECREF)> py_obj(THPVariable_Wrap(tensor), &Py_DECREF); + TORCH_CHECK(py_obj.get() != nullptr, "Failed to wrap tensor to PyObject"); + + DLManagedTensorVersioned* out_tensor = nullptr; + int result = api->managed_tensor_from_py_object_no_sync(py_obj.get(), &out_tensor); + + TORCH_CHECK(result == 0, "from_py_object_no_sync failed with code ", result); + TORCH_CHECK(out_tensor != nullptr, "from_py_object_no_sync returned NULL"); + TORCH_CHECK(out_tensor->version.major == DLPACK_MAJOR_VERSION, + "Expected major version ", DLPACK_MAJOR_VERSION, ", got ", out_tensor->version.major); + TORCH_CHECK(out_tensor->version.minor == DLPACK_MINOR_VERSION, + "Expected minor version ", DLPACK_MINOR_VERSION, ", got ", out_tensor->version.minor); + TORCH_CHECK(out_tensor->dl_tensor.ndim == 3, "Expected ndim 3, got ", out_tensor->dl_tensor.ndim); + TORCH_CHECK(out_tensor->dl_tensor.shape[0] == 2, "Expected shape[0] = 2, got ", out_tensor->dl_tensor.shape[0]); + TORCH_CHECK(out_tensor->dl_tensor.shape[1] == 3, "Expected shape[1] = 3, got ", out_tensor->dl_tensor.shape[1]); + TORCH_CHECK(out_tensor->dl_tensor.shape[2] == 4, "Expected shape[2] = 4, got ", out_tensor->dl_tensor.shape[2]); + TORCH_CHECK(out_tensor->dl_tensor.dtype.code == kDLFloat, "Expected dtype code kDLFloat, got ", out_tensor->dl_tensor.dtype.code); + TORCH_CHECK(out_tensor->dl_tensor.dtype.bits == 32, "Expected dtype bits 32, got ", out_tensor->dl_tensor.dtype.bits); + TORCH_CHECK(out_tensor->dl_tensor.data != nullptr, "Data pointer is NULL"); + + if (out_tensor->deleter) { + out_tensor->deleter(out_tensor); + } + } + + // Test 4: managed_tensor_to_py_object_no_sync + { + auto py_obj_deleter = [](PyObject* p) { if (p) Py_DECREF(p); }; + std::unique_ptr<PyObject, decltype(py_obj_deleter)> py_obj(THPVariable_Wrap(tensor), py_obj_deleter); + TORCH_CHECK(py_obj.get() != nullptr, "Failed to wrap tensor to PyObject"); + + DLManagedTensorVersioned* managed_tensor = nullptr; + int result = api->managed_tensor_from_py_object_no_sync(py_obj.get(), &managed_tensor); + TORCH_CHECK(result == 0, "from_py_object_no_sync failed"); + TORCH_CHECK(managed_tensor != nullptr, "from_py_object_no_sync returned NULL"); + + std::unique_ptr<PyObject, decltype(py_obj_deleter)> py_obj_out(nullptr, py_obj_deleter); + PyObject* py_obj_out_raw = nullptr; + result = api->managed_tensor_to_py_object_no_sync(managed_tensor, reinterpret_cast<void**>(&py_obj_out_raw)); + py_obj_out.reset(py_obj_out_raw); Review Comment:  In Test 4, the custom lambda `py_obj_deleter` for managing the `PyObject*` lifecycle with `std::unique_ptr` can be simplified. The `Py_DECREF` macro from the Python C API already handles `NULL` pointers, so the `if (p)` check is redundant. For better consistency with Test 3 and Test 5, you can use `&Py_DECREF` directly as the deleter. This makes the code more concise and easier to maintain. ```suggestion std::unique_ptr<PyObject, decltype(&Py_DECREF)> py_obj(THPVariable_Wrap(tensor), &Py_DECREF); TORCH_CHECK(py_obj.get() != nullptr, "Failed to wrap tensor to PyObject"); DLManagedTensorVersioned* managed_tensor = nullptr; int result = api->managed_tensor_from_py_object_no_sync(py_obj.get(), &managed_tensor); TORCH_CHECK(result == 0, "from_py_object_no_sync failed"); TORCH_CHECK(managed_tensor != nullptr, "from_py_object_no_sync returned NULL"); std::unique_ptr<PyObject, decltype(&Py_DECREF)> py_obj_out(nullptr, &Py_DECREF); PyObject* py_obj_out_raw = nullptr; result = api->managed_tensor_to_py_object_no_sync(managed_tensor, reinterpret_cast<void**>(&py_obj_out_raw)); py_obj_out.reset(py_obj_out_raw); ``` -- 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: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
