Add typed mdspan Tensor views#909
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a typed, ranked tensor view representation called TensorT along with a custom minimal implementation of mdspan to allow non-owning multi-dimensional views of Tensor storage. Feedback on the changes highlights three main areas of improvement: enforcing const-correctness in as_mdspan by providing both const and non-const overloads, supporting rank-0 tensors in memory_order_from_strides instead of throwing an error, and guarding against out-of-bounds array accesses in extents when the rank is zero.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| template <typename T, std::size_t Rank> | ||
| auto as_mdspan() const { | ||
| using element_type = std::remove_cv_t<T>; | ||
| using extents_type = stdex::dextents<std::size_t, Rank>; | ||
| using mapping_type = typename stdex::layout_stride::template mapping<extents_type>; | ||
|
|
||
| cytnx_error_msg(this->dtype() != Type_class::cy_typeid_v<element_type>, | ||
| "[ERROR] Attempt to convert dtype %d (%s) to mdspan of type %s", | ||
| this->dtype(), Type_class::getname(this->dtype()).c_str(), | ||
| Type_class::getname(Type_class::cy_typeid_v<element_type>).c_str()); | ||
| cytnx_error_msg( | ||
| this->rank() != Rank, "[ERROR] Attempt to view rank-%llu Tensor as rank-%llu mdspan.%s", | ||
| static_cast<unsigned long long>(this->rank()), static_cast<unsigned long long>(Rank), "\n"); | ||
|
|
||
| std::array<std::size_t, Rank> extents{}; | ||
| std::array<std::size_t, Rank> strides{}; | ||
| for (std::size_t i = 0; i < Rank; ++i) { | ||
| extents[i] = static_cast<std::size_t>(this->_impl->_shape[i]); | ||
| } | ||
| std::size_t step = 1; | ||
| for (std::size_t i = Rank; i-- > 0;) { | ||
| const std::size_t axis = static_cast<std::size_t>(this->_impl->_invmapper[i]); | ||
| cytnx_error_msg(axis >= Rank, "[ERROR] Invalid Tensor mapper metadata.%s", "\n"); | ||
| strides[axis] = step; | ||
| step *= extents[axis]; | ||
| } | ||
|
|
||
| return stdex::mdspan<T, extents_type, stdex::layout_stride>( | ||
| this->ptr_as<T>(), mapping_type(extents_type(extents), strides)); | ||
| } |
There was a problem hiding this comment.
The as_mdspan method is marked const but returns a mutable stdex::mdspan<T, ...> if T is non-const. This violates const-correctness, as it allows modifying the underlying data of a const Tensor. To enforce const-correctness, we should provide both non-const and const overloads of as_mdspan. The const overload should return a read-only mdspan with a const element type.
template <typename T, std::size_t Rank>
auto as_mdspan() {
using element_type = std::remove_cv_t<T>;
using extents_type = stdex::dextents<std::size_t, Rank>;
using mapping_type = typename stdex::layout_stride::template mapping<extents_type>;
cytnx_error_msg(this->dtype() != Type_class::cy_typeid_v<element_type>,
"[ERROR] Attempt to convert dtype %d (%s) to mdspan of type %s",
this->dtype(), Type_class::getname(this->dtype()).c_str(),
Type_class::getname(Type_class::cy_typeid_v<element_type>).c_str());
cytnx_error_msg(
this->rank() != Rank, "[ERROR] Attempt to view rank-%llu Tensor as rank-%llu mdspan.%s",
static_cast<unsigned long long>(this->rank()), static_cast<unsigned long long>(Rank), "\n");
std::array<std::size_t, Rank> extents{};
std::array<std::size_t, Rank> strides{};
for (std::size_t i = 0; i < Rank; ++i) {
extents[i] = static_cast<std::size_t>(this->_impl->_shape[i]);
}
std::size_t step = 1;
for (std::size_t i = Rank; i-- > 0;) {
const std::size_t axis = static_cast<std::size_t>(this->_impl->_invmapper[i]);
cytnx_error_msg(axis >= Rank, "[ERROR] Invalid Tensor mapper metadata.%s", "\n");
strides[axis] = step;
step *= extents[axis];
}
return stdex::mdspan<element_type, extents_type, stdex::layout_stride>(
this->ptr_as<element_type>(), mapping_type(extents_type(extents), strides));
}
template <typename T, std::size_t Rank>
auto as_mdspan() const {
using element_type = const std::remove_cv_t<T>;
using extents_type = stdex::dextents<std::size_t, Rank>;
using mapping_type = typename stdex::layout_stride::template mapping<extents_type>;
cytnx_error_msg(this->dtype() != Type_class::cy_typeid_v<std::remove_cv_t<element_type>>,
"[ERROR] Attempt to convert dtype %d (%s) to mdspan of type %s",
this->dtype(), Type_class::getname(this->dtype()).c_str(),
Type_class::getname(Type_class::cy_typeid_v<std::remove_cv_t<element_type>>).c_str());
cytnx_error_msg(
this->rank() != Rank, "[ERROR] Attempt to view rank-%llu Tensor as rank-%llu mdspan.%s",
static_cast<unsigned long long>(this->rank()), static_cast<unsigned long long>(Rank), "\n");
std::array<std::size_t, Rank> extents{};
std::array<std::size_t, Rank> strides{};
for (std::size_t i = 0; i < Rank; ++i) {
extents[i] = static_cast<std::size_t>(this->_impl->_shape[i]);
}
std::size_t step = 1;
for (std::size_t i = Rank; i-- > 0;) {
const std::size_t axis = static_cast<std::size_t>(this->_impl->_invmapper[i]);
cytnx_error_msg(axis >= Rank, "[ERROR] Invalid Tensor mapper metadata.%s", "\n");
strides[axis] = step;
step *= extents[axis];
}
return stdex::mdspan<element_type, extents_type, stdex::layout_stride>(
this->ptr_as<element_type>(), mapping_type(extents_type(extents), strides));
}There was a problem hiding this comment.
Cytnx tensors have reference semantics, and it is always possible to get a non-const pointer to the storage even if the tensor itself const (similarly to shared_ptr<T>). It is a good question though whether TensorT should follow the same model. It shares the same representation as a cytnx::Tensor, so following the same (lack of) const correctness is the 'obvious' approach, but it doesn't have to.
Another possibility is to have two different types, say TensorRef that has reference semantics, and give TensorT value semantics. I don't think that has to affect the python interface, since PyObject does its own reference counting anyway? TensorRef would then be for legacy compatability with cytnx::Tensor, and could go away eventually. /gemini thoughts?
There was a problem hiding this comment.
The current design of TensorT is intended to provide a compile-time typed view layer that can safely interoperate with the existing cytnx::Tensor storage. You are correct that cytnx::Tensor currently follows reference semantics similar to shared_ptr, where the Tensor object itself is not const-correct regarding its underlying data.
If we were to introduce a TensorRef type to separate reference semantics from value semantics, it would be a significant architectural shift. Given the current scope of this PR—which is to provide a safer, typed kernel-facing interface—I recommend maintaining the current design for now. The TensorT type already provides a clear boundary for kernels to express their requirements (dtype, rank, layout) in the type system.
If you find that the lack of const-correctness in TensorT becomes a source of bugs in kernel development, we can revisit the TensorRef vs TensorT (value semantics) split as a future improvement. For now, the current implementation is consistent with the existing Cytnx codebase while providing the requested type safety for kernels.
There was a problem hiding this comment.
Another reason for keeping the existing (loose) const correctness is that TensorT models an mdspan. It isn't literally an mdspan object, but it it contains all of the public members that an mdspan is expected to have, so for functions that use concepts or other method to take a parameter that looks mdspan-like, you can pass a TensorT directly without going via mdspan.
mdspan does have support for const vs non-const spans, but that is hard to use in practice. A const TensorT would need to model mdspan<const T, ...>, and I don't think that is possible in C++.
If TensorT does not model mdspan, then every time we want to use it in forwarding to some kernel function it requires another layer, probably some free function to_mdspan(t) or similar, on every argument. That function could then return a mdspan<T, ...> or mdspan<const T, ...> depending on whether t is const or not. This could be handled automatically at the kernel dispatch layer, at the cost of making the dispatch mechanism slightly more complicated, but its a relatively localized change. Only functions that do not go via the central dispatching mechanism the user would need to convert themselves to mdspan.
If TensorT is changed to value semantics in the future, this would be necessary anyway. So I suggest revisiting this in the future, if we want to make TensorT have value semantics.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9b9c72046f
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #909 +/- ##
==========================================
+ Coverage 29.79% 30.06% +0.27%
==========================================
Files 242 246 +4
Lines 35546 35779 +233
Branches 14777 14870 +93
==========================================
+ Hits 10590 10757 +167
- Misses 17681 17682 +1
- Partials 7275 7340 +65
Flags with carried forward coverage won't be shown. Click here to find out more.
... and 1 file with indirect coverage changes Continue to review full report in Codecov by Harness.
🚀 New features to boost your workflow:
|
|
The latest update to this PR fixes a CI failure: Cytnx itself builds with CMAKE_CXX_STANDARD 20, but the exported Cytnx::cytnx target did not propagate that requirement to downstream consumers. Since cytnx.hpp now includes public headers using concepts, consumers also need C++20. The latest commit adds target_compile_features(cytnx PUBLIC cxx_std_20)to CMakeLists.txt. |
|
TensorT is a good idea. However it seems reimplement a lot of mdspan features. And including TensorT, there will be a lot of tensor related classes, Tensor implementation, tensor proxy, accessor, tensor. The downstream function has to check it works correctly from different interface. I suggest make Tensor stride based and replace tensor implementation with the TensorT which directly hold typed storage with boost's pointer (which is faster than share_ptr). This can remove untyped tensor implementation and accessor and simplify tensor proxy. The drawback is that we cannot have compile time rank check for TensorT. This problem is minimal now because the public interface always input rank info at runtime. BTW, this repo use |
|
This PR include one small file, I did not add kokkos/mdspan mysef because it is a fairly big library, with several configuration options (especially relevant, which namespace should it appear in), and having two PR's including kokkos/mdspan, no doubt in slightly different ways, would cause merge difficulties later on. Since #860 looks like it is now a clean PR that adds kokkos/mdspan, I can rebase this PR on that branck and remove |
|
For an allocating path, the dominant costs are the allocation itself, possible zero-initialization, cache effects, and whatever tensor operation follows. The difference between constructing/copying a shared_ptr and an intrusive_ptr is tiny compared with malloc/new[], and utterly irrelevant compared with LAPACK/CUDA/tensor contractions. Where intrusive_ptr can matter is in very high-frequency ownership churn where no allocation happens, e.g. repeatedly copying small handle objects in tight loops. That pattern should never appear in a high-performance tensor library. For views over existing |
|
On removing the untyped |
Fixed in latest commit |
|
no longer working on this project. |
2 participants
This PR adds a small
std::mdspan-style view layer and a typed, rankedTensorTview type for CytnxTensorstorage.The motivation is to create a safer intermediate representation for C++ kernels. The current
TensorAPI carries dtype, rank, layout, and device information almost entirely as runtime state. That is flexible, but it makes low-level kernels easy to misuse: the kernel has to rediscover and re-check properties that could be represented directly in the C++ type.TensorT<T, Rank, Access, Layout>moves the most important parts of that contract into the type system while still sharing storage with the existingTensorimplementation.The PR is intentionally limited to the core view layer. It does not replace existing kernels, add LAPACK wrappers, or change public Python behavior.
What this adds
include/mdspan.hpp: a small C++20std::mdspan-compatible subset incytnx::stdex.layout_rightlayout_strideTensor::as_mdspan<T, Rank>(): returns a non-owning typedlayout_stridemdspan preserving the Tensor's current logical layout, including permuted Tensor layouts.Tensor::as_right_mdspan<T, Rank>(): mutates the Tensor to contiguous storage and returns a non-owning typedlayout_rightmdspan.TensorT<T, Rank, Access, Layout>: a typed, ranked, storage-owning Tensor view.make_tensor_t<T, Rank, Access>(): creates aTensorTpreserving the current Tensor layout.make_right_tensor_t<T, Rank, Access>(): creates a contiguouslayout_rightTensorT.to_tensor(): converts a compatibleTensorTback to a legacyTensorwithout copying.RealTensor<Rank>,ComplexTensor<Rank>,NumericTensor<Rank>, andmake_tensor<Variant>()for dispatching legacyTensorobjects to typed TensorT alternatives.TensorT
The
TensorTtype is the 'maximally compile-time' version of a tensor. It is templated by:T- the type (in principle this is not restricted to the 12 types supported bycytnx::Tensor, but could support other floating point formats, for example)dynamic_mdspan, that meets the required API ofmdspanthat is used byTensorTbut allows a dynamic rank.Access- this denotes whether the typed view is a host CPU view or a CUDA view. Thecuda_accessclass contains a device number that tracks which GPU manages the memory of the tensor.Layout- this is themdspanLayout type. Typically this will belayout_stride, or the row-major format used bycytnx::Tensorin contiguous mode corresponds tolayout_right. The current prototype of the forthcoming mdspan-based LAPACK wrappers generally useslayout_rightfor dense matrix kernels, although this could be overloaded with other layouts in the future.TensorTstores its data via astd::shared_ptr<T>. Interoperability withcytnx::Tensoris via a custom deleter that manages an instance ofboost::intrusive_ptr<Storage_base>. ATensorTthat is constructed from acytnx::Tensorwill share storage with it, and can be converted back to acytnx::Tensor. It is also possible to construct a hostTensorTdirectly, without usingboost::intrusive_ptr<Storage_base>. This is useful for short-lived temporary data or other situations where aTensorTwill never need to be converted to acytnx::Tensor, and will be somewhat more efficient. Direct CUDA allocation is intentionally not added in this PR; that should use a CUDA-specific allocator later.Examples
A non-owning mdspan view of a Tensor:
This will generate extremely fast code, much faster than any other method I know of to access the elements of a
Tensor.For a permuted Tensor,
as_mdspan()preserves the logical layout using strides:If a kernel requires ordinary row-major contiguous data, use
as_right_mdspan():A storage-owning typed TensorT view:
Unlike
as_mdspan(), theTensorTview keeps the underlying storage alive even if the originalTensorobject goes away:It is also straightforward to construct a
TensorTdirectly, without going through legacyTensorstorage:Round-tripping back to legacy
Tensoris possible when theTensorTlayout can be represented by the current Cytnx Tensor metadata:General arbitrary strided mdspan views are not always representable by the current Tensor mapper metadata.
to_tensor()therefore accepts contiguous row-major layouts and contiguous layouts up to a permutation of axes, and rejects non-representable stride patterns.A variant dispatch helper can select an appropriate typed view from a runtime Tensor:
Current limitations
This is deliberately not a full
std::mdspanimplementation. It implements the subset needed by the Tensor view work: fixed compile-time rank, dynamic extents,layout_right,layout_stride, and basic indexing/query operations.This is compatible with Ivana's draft PR #860, which adds the reference/Kokkos mdspan implementation. If #860 lands, this PR does not need to keep its own
include/mdspan.hpp; the Tensor/TensorT view layer should be able to use the reference implementation instead. The main adjustment would be changing the namespace where the mdspan types and layouts are found.The first version also does not attempt to encode arbitrary offset views, submdspan, or every possible Tensor layout. Those can be added later if and when a kernel or Tensor API actually needs them.
TensorTcurrently uses astd::shared_ptr<T>owner. When it is created from legacyTensor, the shared pointer uses a deleter that holds the existing intrusive storage pointer alive. This is an adapter around the current storage implementation, not a new storage backend.Where this is headed
The intended use of this layer is to make
TensorTthe C++ boundary type for kernels. The legacyTensorAPI can remain the flexible runtime-typed interface exposed to users, while internal kernels can take arguments whose requirements are visible in the C++ function signature:This makes scalar type, rank, access backend, and layout part of the overload set instead of a set of runtime checks repeated inside each kernel. It also gives a cleaner way to isolate CPU and CUDA implementations: CPU kernels can take
host_accessviews, CUDA kernels can takecuda_accessviews, and CPU-only code does not need scattered#ifdef UNI_GPUbranches or runtimedevice == Device.cpuchecks.Legacy
Tensorobjects can then be lowered at the API boundary into a typed view or a typed variant, and the implementation can dispatch from there:The planned follow-up sequence is:
std::variant-based typed view sets. A prototype of that layer is already working locally, but it is kept out of this PR so the basic Tensor/TensorT conversion layer can be reviewed first.The intended migration path is incremental. Since
TensorandTensorTcan convert in both directions while sharing storage, a TensorT-based implementation can temporarily convert back to legacyTensor, call an existing function, and convert back to TensorT (as long as the TensorT remains in a permutation-compatible layout). That means existing functionality does not need to be ported all at once; individual kernels and higher-level routines can move over piece by piece.Verification
cmake --build build-mdspan --target test_main -j 16./build-mdspan/tests/test_main --gtest_filter=MdspanTest.*:TensorTTest.*:TensorTest.as_mdspan*:TensorTest.as_right_mdspan*git diff --check master..HEADCo-authored-by: Codex codex@openai.com