diff --git a/python/paddle/base/dygraph/tensor_patch_methods.py b/python/paddle/base/dygraph/tensor_patch_methods.py index d3a60c0c04bc4..47f25e3e1191a 100644 --- a/python/paddle/base/dygraph/tensor_patch_methods.py +++ b/python/paddle/base/dygraph/tensor_patch_methods.py @@ -1476,6 +1476,21 @@ def __dlpack__(self, stream=None): return paddle.to_dlpack(self) + def __tvm_ffi_env_stream__(self) -> int: + """ + Returns the raw stream pointer of the current tensor's device context. + This is used for TVM FFI environment integration. + """ + if self.place.is_gpu_place(): + return paddle.base.libpaddle._get_current_raw_stream( + self.place.gpu_device_id() + ) + else: + # TODO: Add XPU and custom device support. + raise RuntimeError( + "Currently, the __tvm_ffi_env_stream__ method is only supported for GPU tensors." + ) + if not hasattr(core, "eager"): return @@ -1523,6 +1538,7 @@ def __dlpack__(self, stream=None): ("__cuda_array_interface__", __cuda_array_interface__), ("__dlpack__", __dlpack__), ("__dlpack_device__", __dlpack_device__), + ("__tvm_ffi_env_stream__", __tvm_ffi_env_stream__), ): setattr(core.eager.Tensor, method_name, method) diff --git a/test/dygraph_to_static/test_tensor_attr_consistency.py b/test/dygraph_to_static/test_tensor_attr_consistency.py index 98750369736ff..86a4437a7c69c 100644 --- a/test/dygraph_to_static/test_tensor_attr_consistency.py +++ b/test/dygraph_to_static/test_tensor_attr_consistency.py @@ -80,6 +80,7 @@ "__cuda_array_interface__", '__dlpack__', "__dlpack_device__", + "__tvm_ffi_env_stream__", ] ) STATIC_ONLY_TENSOR_ATTRS_ALLOW_LIST = OrderedSet( diff --git a/test/legacy_test/test_tvm_ffi.py b/test/legacy_test/test_tvm_ffi.py new file mode 100644 index 0000000000000..aa6a91b4aa24d --- /dev/null +++ b/test/legacy_test/test_tvm_ffi.py @@ -0,0 +1,38 @@ +# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved. +# +# Licensed 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. + +import unittest + +import paddle + + +class TestTVMFFI(unittest.TestCase): + def test_tvm_ffi_env_stream_for_gpu_tensor(self): + if not paddle.is_compiled_with_cuda(): + return + tensor = paddle.to_tensor([1.0, 2.0, 3.0]).cuda() + current_raw_stream_ptr = tensor.__tvm_ffi_env_stream__() + self.assertIsInstance(current_raw_stream_ptr, int) + self.assertNotEqual(current_raw_stream_ptr, 0) + + def test_tvm_ffi_env_stream_for_cpu_tensor(self): + tensor = paddle.to_tensor([1.0, 2.0, 3.0]).cpu() + with self.assertRaisesRegex( + RuntimeError, r"the __tvm_ffi_env_stream__ method" + ): + tensor.__tvm_ffi_env_stream__() + + +if __name__ == '__main__': + unittest.main()