diff --git a/src/pybind11_utils.h b/src/pybind11_utils.h index b1bfe344f..d27da4918 100644 --- a/src/pybind11_utils.h +++ b/src/pybind11_utils.h @@ -20,6 +20,57 @@ namespace infini::ops { namespace detail { +// Thread-local cache for tensor conversions to avoid repeated Python attribute access. +// Design follows the same pattern as Operator Cache (cache.key in operator.h). +// Key = data_ptr + shape + strides + dtype (all metadata that affects TensorView). +class TensorCache { + public: + struct Key { + void* data_ptr; + std::vector shape; + std::vector strides; + DataType dtype; + + bool operator==(const Key& other) const { + return data_ptr == other.data_ptr && shape == other.shape && + strides == other.strides && dtype == other.dtype; + } + }; + + struct KeyHash { + std::size_t operator()(const Key& k) const { + std::size_t seed = std::hash{}(k.data_ptr); + for (auto d : k.shape) { + seed ^= std::hash{}(d) + 0x9e3779b9 + (seed << 6) + (seed >> 2); + } + for (auto s : k.strides) { + seed ^= std::hash{}(s) + 0x9e3779b9 + (seed << 6) + (seed >> 2); + } + seed ^= std::hash{}(static_cast(k.dtype)) + 0x9e3779b9 + (seed << 6) + (seed >> 2); + return seed; + } + }; + + static TensorCache& Instance() { + thread_local TensorCache cache; + return cache; + } + + const Tensor* Find(const Key& key) const { + auto it = cache_.find(key); + return (it != cache_.end()) ? it->second.get() : nullptr; + } + + void Insert(const Key& key, Tensor tensor) { + cache_[key] = std::make_unique(std::move(tensor)); + } + + void Clear() { cache_.clear(); } + + private: + std::unordered_map, KeyHash> cache_; +}; + inline py::handle InternedName(const char* value) { // Keep one reference for the process lifetime; Python objects must not be // decref'd by a static destructor after interpreter finalization. @@ -274,7 +325,33 @@ inline Device DeviceFromPybind11Handle(py::handle obj) { inline Tensor TensorFromPybind11Handle(py::handle obj) { [[maybe_unused]] HostRangeScope host_range_tensor_conversion{ HostRangeLayer::kTensorConversion}; - return detail::TensorFromPybind11HandleImpl(obj); + + const auto& names{detail::GetInternedNames()}; + auto data_ptr{reinterpret_cast( + detail::CallMethodNoArgs(obj, names.data_ptr).cast())}; + + auto shape{detail::VectorFromPybind11Handle( + py::getattr(obj, names.shape))}; + + auto dtype{detail::DataTypeFromPybind11HandleImpl(py::getattr(obj, names.dtype))}; + + auto strides{detail::VectorFromPybind11Handle( + detail::CallMethodNoArgs(obj, names.stride))}; + + detail::TensorCache::Key key{data_ptr, {shape.begin(), shape.end()}, + {strides.begin(), strides.end()}, dtype}; + + const Tensor* cached = detail::TensorCache::Instance().Find(key); + if (cached != nullptr) { + return *cached; + } + + // Cache miss: full conversion. + auto device = detail::DeviceFromPybind11HandleImpl(obj); + Tensor result{data_ptr, std::move(shape), dtype, device, std::move(strides)}; + detail::TensorCache::Instance().Insert(key, std::move(result)); + + return *detail::TensorCache::Instance().Find(key); } inline std::optional OptionalTensorFromPybind11Handle(