refactor(trace): implement the GPU Tensor trace as CUDA kernels#850
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 benchmarks for linalg::Trace and refactors both the CPU and GPU implementations of Trace_internal to perform in-place diagonal summation using physical strides, eliminating the need to contract with an identity tensor. The review feedback identifies critical bugs where an empty tensor (Ndiag == 0) causes unsigned underflow and subsequent out-of-bounds memory access in the CPU functions. For the GPU implementation, an empty element count (Nelem == 0) can lead to an invalid CUDA configuration error, and the use of redundant cudaDeviceSynchronize() calls blocks host execution and degrades performance.
e355381 to
4c76053
Compare
f480e9b to
6a2cce9
Compare
IvanaGyro
left a comment
IvanaGyro
left a comment
There was a problem hiding this comment.
(Review written by Claude on behalf of @IvanaGyro.)
I verified the central correctness claim: PhysicalStrides reconstructs exactly the row-major-over-invmapper layout that the container layer already uses to address elements (Tensor_impl::get does vec_map(_shape, _invmapper) then row-major over that order), so summing the diagonal in place via strides[ax1] + strides[ax2] is correct, and it's strictly more correct than the old _trace_2d, which assumed contiguous storage (rawdata[i*Ldim + i]) and would have been wrong for a lazily-permuted input. The GPU 2D tree reduction is valid (block dim 512 is a power of two), and the nd offset precomputation mirrors the CPU path. Nice refactor — it removes the inverted layering and the full-tensor contiguous() copy.
A few non-blocking points, left inline:
PhysicalStridesis duplicated byte-for-byte between the CPU and GPU translation units.extent = (Ndiag - 1) * diag_stride + 1underflows ifNdiag == 0.- The GPU tree reduction silently assumes
_TNinB_TRACE_is a power of two.
Also a cross-PR note: this PR makes StrideView (from #849) live — the CPU trace now reduces a PairwiseSum over a StrideView temporary. The current call sites are safe because the StrideView temporary outlives the whole PairwiseSum(...) full-expression, but it means the dangling-pointer concern @gemini-code-assist raised on #849's StrideView::Iterator now has a real consumer; worth fixing the iterator (store the base iterator, not a pointer to the view) before more call sites appear in #834 follow-ups.
GPU build + test results (openblas-cuda / mkl-cuda) will follow in a separate comment.
|
(Comment written by Claude on behalf of @IvanaGyro.) Test-coverage assessment. The correctness-relevant paths this PR adds aren't all exercised by committed tests:
Suggestions rather than blockers — except (1) and (2): the permuted-input path is the core change of this PR and currently has no automated coverage. |
|
(Comment written by Claude on behalf of @IvanaGyro.) GPU build + test resultsBuilt and tested on two GPUs (GTX 1660 SUPER, RTX 2080 Ti), CUDA 12, conda toolchain. Builds — both succeeded (rc=0):
Tests (release, the authoritative result for this PR):
Canonical ASan run (
Conclusion: this PR is green on GPU for both BLAS backends; the only failures are unrelated ASan/debug pre-existing issues. Separately, while running the suite I noticed the GPU tests are dominated by a handful of extremely slow element-wise |
4c76053 to
38fdf4f
Compare
6a2cce9 to
95cb3b0
Compare
IvanaGyro
left a comment
IvanaGyro
left a comment
There was a problem hiding this comment.
(Review written by Claude on behalf of @IvanaGyro.)
This round resolves all four of my earlier code comments: PhysicalStrides is hoisted into a public Tensor::strides() (no more CPU/GPU duplication, and the impl matches the invmapper layout Tensor::at uses), the Ndiag == 0 underflow is guarded on both backends, the _TNinB_TRACE_ power-of-two invariant has a static_assert, and the reductions now use the | stride(...) adaptor. Nice. Remaining items are about test coverage — inline.
|
(Comment written by Claude on behalf of @IvanaGyro.) GPU build + test results (updated commit
|
38fdf4f to
3a5a680
Compare
95cb3b0 to
fd6fd1f
Compare
47869fa to
b295b04
Compare
b295b04 to
cce6d7d
Compare
IvanaGyro
changed the title
IvanaGyro
left a comment
IvanaGyro
left a comment
There was a problem hiding this comment.
(Review written by Claude on behalf of @IvanaGyro.)
Replacing the eye+Contract GPU trace with direct kernels is a solid win and mirrors the CPU TraceImpl offset math 1:1, so the two stay easy to keep in sync. The cuda::std::complex mapping for the complex dtypes is the right way to keep the kernels type-generic, and the power-of-two static_assert on the block size correctly guards the tree reduction. I built and ran this on both openblas-cuda and mkl-cuda: DenseUniTensorTest.gpu_Trace and BlockUniTensorTest.gpu_Trace pass on both. Comments inline: an async-launch error-check gap, and a GPU test-coverage gap that I think is the most important follow-up here.
621a730 to
915105b
Compare
…exits The zero-extent early return and the normal function exit both ended with the same three lines: Tensor::from_storage on the filled storage, a conditional reshape_ to output_shape, and the return. Factoring this into ComposeTraceOutput means a future change to how the output Tensor is composed only needs to be made in one place. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
src/backend/StorageImplementation.cpp includes cuFill_gpu.hpp unconditionally, regardless of whether the build has CUDA enabled -- the same pattern every other utils_internal_gpu header in this directory (cuAlloc_gpu.hpp, cuMovemem_gpu.hpp, etc.) already follows by wrapping their CUDA-specific declarations in #ifdef UNI_GPU. cuFill_gpu.hpp's unconditional `#include "cuda/std/complex"` and the ToCudaDType specializations built on cuda::std::complex broke every non-GPU build (no CUDA toolkit installed, so the header doesn't exist) with "fatal error: cuda/std/complex: No such file or directory". Wrapping both in #ifdef UNI_GPU, matching the sibling headers' convention, restores non-GPU builds while leaving GPU builds unchanged. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…read TraceImplGpu previously called cudaFree directly on the two small device-side layout arrays after the kernel launch. cudaFree is a synchronizing legacy-API call, so every GPU Trace() stalled the host thread until the whole device pipeline drained, even though the buffers themselves are only a handful of uint64s. Records a CUDA event right after the kernel launch and hands the wait-then-free off to a detached std::thread, so TraceImplGpu returns as soon as the kernel is queued instead of blocking on cudaFree. The wait must run on a detached std::thread rather than std::async: an unstored std::future returned by std::async blocks in its own destructor until the task completes, which would reintroduce the same synchronous stall at the end of the calling expression. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
|
(Comment written by Claude on behalf of @IvanaGyro.) GPU memory-management microbenchmarksStandalone microbenchmarks (not part of this PR's diff or test suite) reproducing this PR's 1. cudaFree strategy: direct vs. thread-deferred vs. cudaFreeAsyncSwept
What each case actually runs: Direct — what this PR replaced: cudaMalloc(&d_shape, ...); cudaMalloc(&d_stride, ...);
cudaMemcpy(d_shape, ...); cudaMemcpy(d_stride, ...);
TraceKernel<<<...>>>(...);
cudaDeviceSynchronize(); // explicit here for a clean timing boundary; a bare
// cudaFree() already does this internally
cudaFree(d_shape); cudaFree(d_stride);Thread-deferred — this PR's actual pattern: TraceKernel<<<...>>>(...);
cudaEvent_t free_ready;
cudaEventCreate(&free_ready);
cudaEventRecord(free_ready, 0);
std::thread([d_shape, d_stride, free_ready]() {
cudaEventSynchronize(free_ready);
cudaFree(d_shape); cudaFree(d_stride);
cudaEventDestroy(free_ready);
}).detach();
cudaMallocAsync(&d, ..., 0); cudaMemcpyAsync(d, ..., 0);
TraceKernel<<<...>>>(...);
cudaFreeAsync(d, 0);
cudaStreamSynchronize(0); // the trap: this also waits for the free, which
// exposes the stream as idle and trips the memory
// pool's default release-to-OS behavior
TraceKernel<<<...>>>(...);
cudaEventRecord(kernel_done, 0);
cudaEventSynchronize(kernel_done); // waits only for the kernel
cudaFreeAsync(d, 0); // enqueued, not waited onTakeaway: direct 2. Memory-block comparisonSame call pattern, comparing six ways to build (or avoid building) the small per-call
1. std::vector<uint64_t> host(rank);
for (i) host[i] = ...;
cudaMalloc(&d, sizeof(uint64_t) * rank);
cudaMemcpy(d, host.data(), sizeof(uint64_t) * rank, cudaMemcpyHostToDevice);2. thrust::universal_vector<uint64_t> v;
v.reserve(rank);
for (i) v.push_back(...);
// kernel reads thrust::raw_pointer_cast(v.data())3. cudaMalloc(&d, sizeof(uint64_t) * rank);
for (i) {
uint64_t val = ...;
cudaMemcpy(d + i, &val, sizeof(uint64_t), cudaMemcpyHostToDevice);
}4. cudaMallocManaged(&d, sizeof(uint64_t) * rank);
for (i) d[i] = ...; // plain host pointer write, no copy API needed5. // (a) reserve-then-append (Storage has no dedicated reserve(), but resize()
// only grows the buffer when newsize > capacity_, so shrinking back to 0
// leaves the grown capacity_ in place):
Storage s(0, Type.Uint64, Device.cuda);
s.resize(rank); s.resize(0);
for (i) s.append(cytnx_uint64(...)); // now hits append()'s cheap
// size_ != capacity_ branch, no realloc
// (b) sized constructor + at<>() write:
Storage s(rank, Type.Uint64, Device.cuda); // one allocation, size_ = rank up front
for (i) s.at<cytnx_uint64>(i) = ...; // valid host write: GPU Storage backs
// onto cudaMallocManaged, not cudaMalloc
// (c) presized + from_vector:
Storage s = Storage::from_vector(host, Device.cuda);6. no constexpr int kMaxTraceRank = 64;
struct TraceLayout {
uint64_t rank;
uint64_t data[kMaxTraceRank];
};
__global__ void TraceKernel(const __grid_constant__ TraceLayout layout, uint64_t* sink) {
uint64_t acc = 0;
for (uint64_t i = 0; i < layout.rank; ++i) acc += layout.data[i] * (i + 1);
if (threadIdx.x == 0 && blockIdx.x == 0) *sink = acc;
}
// per call:
TraceLayout layout{};
layout.rank = rank;
for (i) layout.data[i] = ...;
TraceKernel<<<grid, block>>>(layout, sink);
cudaDeviceSynchronize(); // no free to wait for, but the call must still not
// return until the kernel is doneTakeaway: skipping One correctness caveat worth flagging explicitly: this fast path requires a compile-time cap on Also worth flagging separately: the unreserved Not part of this PR's diff or CI — standalone benchmarks run locally against this branch, for informing the review discussion above. |
…uffers The prior commit deferred freeing device_surviving_shape and device_surviving_input_stride to a detached std::thread synchronized on a CUDA event, to avoid cudaFree's implicit device-wide sync stalling the host on every Trace() call. That trade only pays off if the deferred wait is shorter than the thread it runs on costs to spawn. Here it isn't: TraceKernel is already launched on the same in-order default stream cudaFree would wait on, so cudaFree's wait is bounded by the kernel's own completion -- there's no additional stall to defer. The std::thread construction and destruction on every call adds fixed overhead on top of that same wait, at every grid size and kernel duration. Reverts to a direct cudaFree and drops the now-unused <thread> include. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The prior commit reverted to a direct cudaFree for device_surviving_shape/device_surviving_input_stride, since with every op on the default stream, cudaFree's device-wide wait is already bounded by TraceKernel's own completion -- no additional stall to avoid. That equivalence only holds because everything currently runs on the default stream. Several internal GPU ops already create their own non-default streams (cuEig_internal.cu, cuQuantumQr_internal.cu, cuQuantumGeSvd_internal.cu, cutensornet.cu), so a genuinely multi-stream call pattern is plausible as this backend grows. A direct cudaFree synchronizes the whole device rather than just the stream that produced the freed buffers, so once concurrent unrelated stream work exists, it would stall that work too. Switches the allocation/copy/free of the two layout buffers to cudaMallocAsync/cudaMemcpyAsync/cudaFreeAsync on the default stream. cudaEventRecord+cudaEventSynchronize after the kernel launch still waits for TraceKernel specifically (not the whole device) before the buffers are released; cudaFreeAsync then enqueues the release into the stream-ordered memory pool without the host blocking on the deallocation itself. cudaFreeAsync requires memory allocated via cudaMallocAsync/cudaMallocFromPoolAsync, so the allocation and copy calls switch to their async counterparts too. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…u.hpp ToCudaDType is a generic scalar-to-device-type trait with no connection to Fill specifically; cuTrace_internal.cu including cuFill_gpu.hpp only for this trait meant the trace implementation depended on an unrelated operation's header. Moves the trait to a new cuTypeTraits_gpu.hpp, leaving cuFill_gpu.hpp with just the FillGpu declaration it's named for. cuFill_gpu.cu and cuTrace_internal.cu now include the trait header directly instead of picking it up as a side effect of including cuFill_gpu.hpp. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
TraceImplGpu mixed cytnx_uint64 (diagonal_length, diagonal_stride, the surviving-axis shape/stride arrays, surviving_rank, output_size) with the cytnx_int64 Tensor::strides() and output_shape already needed, so values crossed between the two with a static_cast at nearly every assignment. Switches every index/extent in TraceImplGpu, TraceKernel, and DecodeDiagonalStartOffset to cytnx_int64 throughout, mirroring the CPU TraceImpl (Trace_internal.cpp, PR #862): CheckedCastToInt64Gpu is now the only place a uint64 shape value crosses into the signed side, and std::ssize replaces a manual cast for surviving_rank. output_shape also now doubles as the surviving-axis extent array uploaded to the device, since after the cast it holds exactly the same values as the former host_surviving_shape, removing that separate vector. blockIdx/gridDim/threadIdx stay CUDA's native unsigned int, cast once at the point they enter the int64 arithmetic, since they aren't part of the host-side bookkeeping this mirrors and CUDA's execution model defines them as unsigned. Also names the previously-repeated literal `0` stream argument once as a local `stream` variable, used consistently across the cudaMallocAsync/cudaMemcpyAsync/cudaEventRecord/cudaFreeAsync calls, and passes it explicitly to the TraceKernel launch's stream slot (previously implicit) so every operation in this function is demonstrably on the same stream rather than relying on the default stream being the same thing by convention. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
|
I guess it copies similar style from elsewhere in Cytnx, but for small parameters do not use constexpr int kMaxTraceRank = 64; // a bit extreme, but OK.....
struct TraceLayout {
int64_t rank;
int64_t shape[kMaxTraceRank];
int64_t input_stride[kMaxTraceRank];
};
TraceLayout layout{};
layout.rank = surviving_rank;
for (int i = 0; i < surviving_rank; ++i) {
layout.shape[i] = output_shape[i];
layout.input_stride[i] = surviving_input_stride[i];
}
TraceKernel<<<grid, block, 0, stream>>>(out, in, layout, ...);
cudaStreamSynchronize(stream);The "Thread-deferred" code above is unsafe for the unified memory model. The function must not return until the GPU device has syncronized, otherwise host reads via unified memory will receive corrupted data. |
|
make sure template <class T>
__global__ void TraceKernel(
T* __restrict__ output_data,
const T* __restrict__ input_data,
const __grid_constant__ TraceLayout layout,
cytnx_int64 output_size,
cytnx_int64 diagonal_length,
cytnx_int64 diagonal_stride) {
...
} and if using a |
TraceImplGpu always shipped the surviving-axis shape/stride arrays to the device via cudaMallocAsync + cudaMemcpyAsync + cudaFreeAsync, even though those arrays are only a handful of int64s for the vast majority of tensor-network tensors (MPS/MPO rank 2-4, PEPS rank 4-6, MERA/TTN rank 3-4). That round trip's allocator and driver-call overhead dominates the actual reduction for calls this small. Adds TraceLayout, a fixed-size (kMaxTraceRank = 32) struct holding the surviving-axis rank/shape/stride, and TraceKernelFast, a kernel that takes it by value as a __grid_constant__ parameter -- placed in read-only per-grid constant memory rather than copied into each thread's local memory -- instead of two device pointers. For surviving_rank <= kMaxTraceRank (the common case), TraceImplGpu populates the struct directly in the same pass that builds output_shape and launches TraceKernelFast: no cudaMalloc, cudaMemcpy, or cudaFree at all. 32 mirrors NumPy's historical NPY_MAXDIMS -- far above any realistic tensor-network rank, and free in kernel parameter/constant-memory budget regardless. Above that cap, TraceImplGpu falls back to the previous cudaMallocAsync/cudaMemcpyAsync/TraceKernel/cudaFreeAsync path, since TraceLayout can't hold more than kMaxTraceRank entries. The fallback rebuilds the stride list as three contiguous std::back_inserter range copies around ax1/ax2 rather than a per-element branch, since this path is only reached when rank is large enough for that branch to cost something. Adds SurvivingRankAtFastPathCapMatchesReference and SurvivingRankAboveFastPathCapUsesFallback to tests/gpu/linalg_test/ Trace_test.cpp, pinning correctness on both sides of the kMaxTraceRank boundary against the CPU reference. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
|
Please don't add a fallback case that uses So that means the surviving_rank corresponds to a tensor of size at least The number of bits that an x86-64 machine can physically address is 52. So the maximum conceivable size of kMaxTraceRank = 50and remove the fallback path. It is irrelevant for the cost of the kernel launch. You could set it to 128 and it will still fit comfortably in pre-sm_70 devices that have much smaller |
Surviving ranks may be 1 for 3-rank tensor.
|
Not sure what you mean here - surviving ranks could be zero, for trace of rank-2 tensor (or higher rank where some of the dimensions are 1). Did you misread my previous comment? I notice that the quote has munged the LaTeX - "at least 2 surviving rank" was part of a Latex expression 2^{surviving rank + 2} as the minimum number of elements. |
Sorry, I misread. I read this to |
|
(Comment written by Claude on behalf of @IvanaGyro.) Validating the async-allocator + grid-constant fixes against the actual
|
| Benchmark | cudaMallocAsync/cudaFreeAsync, no grid-constant (1f745365) |
__grid_constant__ + async fallback (latest, d5b98c73) |
|---|---|---|
Strided_3D_Double/256/64 |
27.5us | 16.4us |
Strided_3D_Double/1024/16 |
28.9us | 18.0us |
Strided_3D_Double/2048/16 |
31.5us | 20.4us |
Strided_3D_Double/4096/8 |
33.8us | 23.1us |
Strided_3D_ComplexDouble/256/64 |
28.2us | 17.4us |
Strided_3D_ComplexDouble/1024/16 |
29.9us | 18.9us |
Strided_3D_ComplexDouble/2048/16 |
32.4us | 21.6us |
Strided_3D_ComplexDouble/4096/8 |
34.9us | 24.1us |
Strided_2D (rank-2, surviving_rank == 0) is flat across both commits, as expected — that path never allocated a layout buffer even before either fix, since the decode loop is empty when there are no surviving axes. Matvec_3D/Vecdot_2D/Reshape_2D (unrelated reference paths using other cytnx ops) are also unchanged across both commits, confirming no regression elsewhere.
Takeaway: skipping the allocation entirely via __grid_constant__ is a genuine additional win on top of the async-allocator switch — roughly a further 40% reduction (e.g. 27.5us → 16.4us), not a redundant micro-optimization once cudaMallocAsync/cudaMemcpyAsync/cudaFreeAsync are already in place. Those three calls still carry real per-call driver overhead even without a thread spawn or a full-device sync, and the grid-constant path removes them entirely for the common surviving_rank <= kMaxTraceRank case.
Not part of this PR's diff or CI — standalone benchmark runs against the two commits above, for informing the review discussion.
That is actually a key point. The overhead of passing |
…eRank to 50 TraceImplGpu kept a cudaMallocAsync/cudaMemcpyAsync/cudaFreeAsync fallback for surviving_rank > kMaxTraceRank, guarding a case that cannot occur on any physically constructible machine: every surviving axis has extent >= 2 (a size-1 axis contributes nothing to the diagonal decode and would be squeezed away rather than kept), so a tensor with surviving_rank axes plus the two traced axes has at least 2^(surviving_rank + 2) elements. x86-64 addresses at most 2^52 bytes, so even a 1-byte-per-element dtype caps surviving_rank at 50 -- no dtype/hardware combination can reach 51. Raising kMaxTraceRank from 32 to 50 costs nothing extra against CUDA's per-launch parameter/ constant-memory budget, so there was no reason to keep the fallback path duplicating the same reduction logic through a second kernel just to handle an unreachable case. Removes the fallback branch and the now-dead pointer-parameter TraceKernel, renaming TraceKernelFast to TraceKernel since it is once again the only kernel. Replaces the removed fallback with a defensive cytnx_error_msg guard: if surviving_rank ever did exceed kMaxTraceRank, TraceLayout's shape/stride arrays beyond that index were never populated by the build loop, so proceeding would read uninitialized stack memory on the device -- the guard turns that into a clear error instead of silent corruption. tests/gpu/linalg_test/Trace_test.cpp: replaces SurvivingRankAtFastPathCapMatchesReference (rank 32) and SurvivingRankAboveFastPathCapUsesFallback (rank 33, the now-removed path) with a single SurvivingRankAtTraceLayoutCapMatchesReference at the new rank-50 boundary, since there is no separate fallback path left to pin. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
![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: 8d28f0edc6
ℹ️ 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".
TraceImplGpu's kMaxTraceRank derivation assumed every surviving axis
has extent >= 2, but the code stored every surviving axis in
TraceLayout regardless of extent -- including size-1 axes, which the
public API does not disallow and does not squeeze away. A tensor with
many size-1 surviving axes (e.g. shape {2, 2, 1, 1, ...} with 51+
trailing 1s, traced over the first two axes) has only a handful of
elements, so it satisfies the address-space argument trivially, yet
the code rejected it once surviving_rank exceeded kMaxTraceRank -- a
regression the fallback path this branch removed would not have hit.
A surviving axis of extent 1 contributes nothing to
DecodeDiagonalStartOffset's per-output decode: its odometer step is
always (remaining_flat_index % 1 == 0, remaining_flat_index /= 1
leaves it unchanged, stride * 0 == 0), so omitting it from the decode
arrays changes nothing about the offsets computed for the other axes.
TraceImplGpu now only stores axes with extent > 1 in TraceLayout, and
counts only those toward kMaxTraceRank; output_shape (used for
Tensor::reshape_) still includes every surviving axis regardless of
extent, so a size-1 surviving axis still appears in the output shape
correctly. This is what makes kMaxTraceRank's derivation (every
counted axis has extent >= 2) true of what's actually stored, rather
than merely assumed of the input.
DecodeDiagonalStartOffset's parameters are renamed from
surviving_shape/surviving_input_stride/surviving_rank to
axis_shape/axis_input_stride/axis_count, since the count it receives
(TraceLayout::rank) is no longer the tensor's literal surviving rank
and the old names invited exactly this kind of confusion.
tests/gpu/linalg_test/Trace_test.cpp: replaces the previous cap
boundary test (which, after this fix, no longer exercises the cap at
all, since all its surviving axes are size-1) with
SingletonSurvivingAxesDoNotCountTowardTraceLayoutCap, matching this
regression directly, and MixedSingletonAndNonTrivialSurvivingAxesPreserveShape,
covering singleton and non-trivial surviving axes interleaved in the
same tensor.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7ab7ee0911
ℹ️ 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".
| const cytnx_int64 dim = CheckedCastToInt64Gpu(input_shape[axis], "input_shape[axis]"); | ||
| output_shape.push_back(dim); | ||
| if (dim > 1) { | ||
| cytnx_error_msg(non_trivial_rank >= kMaxTraceRank, |
There was a problem hiding this comment.
When any surviving axis has extent 0, output_size becomes 0 and the function should take the zero-output early return without launching TraceKernel, so the fixed-size TraceLayout is not needed. This check runs before output_size is computed, so a cheap zero-extent tensor such as reshape({2, 2, 0, 2, ...}) with more than 50 extent-2 surviving axes now errors on GPU while the CPU path returns an empty tensor with the requested shape.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Confirmed and fixed in bf13a1a. The kMaxTraceRank guard ran inside the axis-building loop, before output_size was known — so it could reject a tensor whose eventual output is empty, exactly as you describe.
Moved the guard to after output_size is computed and after the zero-extent early return, so it only fires when a real kernel launch (and thus TraceLayout) is actually about to happen. The loop no longer errors inline; it just stops writing into the fixed arrays once they're full while still tracking the true non_trivial_rank, so the guard has an accurate count whenever it does run.
Added ManyNonTrivialSurvivingAxesWithZeroExtentReturnsEmptyOutput to tests/gpu/linalg_test/Trace_test.cpp, close to your example: 51 non-trivial (extent-2) surviving axes plus one zero-extent surviving axis, built via the existing ZeroExtentGpuTensor helper so it stays cheap.
Generated by Claude Code
…ank guard TraceImplGpu's kMaxTraceRank guard ran inside the loop that builds output_shape and TraceLayout, before output_size was known. A tensor with more than kMaxTraceRank non-trivial (extent > 1) surviving axes plus one zero-extent surviving axis has output_size == 0 -- the zero-extent early return should apply, matching the CPU path's empty output -- but the guard rejected it first, since it never got to see that the eventual output was empty. Moves the kMaxTraceRank check to after output_size is computed and after the zero-extent early return, so it only fires when a real kernel launch (and thus TraceLayout) is actually about to happen. The loop itself no longer errors inline; it just stops writing into layout.shape/stride once the fixed arrays are full, tracking the true non_trivial_rank regardless, so the guard has an accurate count to check once it does run. Adds ManyNonTrivialSurvivingAxesWithZeroExtentReturnsEmptyOutput to tests/gpu/linalg_test/Trace_test.cpp: a tensor with 51 non-trivial surviving axes and one zero-extent surviving axis, built via ZeroExtentGpuTensor so it stays cheap despite the rank. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
4 participants
Fixes #834.
Summary
Implements the GPU Tensor trace as native CUDA kernels reading storage directly via
Tensor::strides(), replacing the previousContract-based path that built an identityUniTensorper call.Change
BlockTraceDiagonal<T>: sums one diagonal with the whole block. Each thread strides over the diagonal, each warp reduces its lanes with a__shfl_down_synctree, the per-warp sums land in shared memory, and the first warp reduces those. A single-warp launch returns directly after the warp shuffle (no shared memory or barrier).TraceKernel<T>: one block per output element. The block decodes its flat index into the surviving-axis multi-index usingDecodeDiagonalStartOffset, then reduces the corresponding diagonal viaBlockTraceDiagonal. The rank-2 trace is theoutput_size == 1/surviving_rank == 0case — the decode loop is empty, the diagonal starts at offset 0, and a single full block reduces the whole matrix diagonal.diagonal_length:threads_per_block = min(round_up_to_warp(diagonal_length), 256). A short diagonal launches a small block (no idle warps; many such blocks run concurrently across SMs); a long diagonal gets the full 256-thread block with each thread striding. The rank-2 trace puts every thread of its single block on the one diagonal.WarpShuffleDownAddwraps__shfl_down_sync, overloaded forcuda::std::complex(shuffled component-wise since__shfl_down_synconly moves scalar lanes). Small integer types promote tointthrough the shuffle and truncate back, preserving the modular-sum semantics the trace already has.TraceImplGpu<T>mirrors the CPUTraceImpl<T>: derives the reduced output shape and per-surviving-axis input strides in a single pass; fills a device-residentStorageand promotes it viaTensor::from_storage(no host round-trip).diagonal_length == 0/output_size == 0guards skip the kernel launch;checkCudaErrors(cudaGetLastError())after launch surfaces any launch/configuration failure at the trace call rather than the next unrelated CUDA call.surviving_rank-sized layout arrays (shape + per-axis input stride) are shipped to the device; the kernel decodes each element's offset itself rather than the host materializing anoutput_size-sized offset table.TviaT(0)/+=; complex storage isreinterpret_casttocuda::std::complex<{float,double}>, which provides deviceoperator+/ zero-construction (CUDA ≥ 12.6).Tests
tests/gpu/linalg_test/Trace_test.cppasLinalgGpuTraceTestmirrors the CPULinalgTraceTestmatrix:PermutedRank3MatchesContiguous/PermutedTracesAcrossRanksAndDtypes: permuted (non-contiguous) inputs over non-adjacent axes across ranks 3–5 and Double/ComplexDouble/Int32, compared against the trace of a contiguous CPU clone.Rank2Path: thesurviving_rank == 0single-block path.OutputRankIsInputMinusTwo,SwappedAxisOrderMatches: rank and axis-order invariants.MatchesCpuTraceOnContiguousInput: GPU result equals the CPU pairwise result on identical contiguous data across dtypes.BlockSizeBoundaries: rank-2 traces at diagonal lengths {1, 31, 32, 33, 64, 255, 256, 257, 1024} — straddles every block-sizing boundary (sub-warp, exact-warp, warp+1, the 256 cap and one past it, multi-stride).ShortDiagonalsManyOutputs:{n, middle, n}traced over (0, 2) for(n, middle) = (2, 1000), (32, 500), (33, 100)— many small blocks per launch, exercises block indexing across a large output together with sub-warp / exact-warp / warp+1 block sizes.LargeDiagonalAccuracyBound: upper bounds the diagonal-sum accuracy atdiagonal_length = 4096with a slack proportional ton · ε · |value|— generous enough that the current serial intra-thread accumulation passes, and any future intra-block tree reduction also passes (locks correctness, not the current error regime).Benchmark
benchmarks/linalg/Trace_bm.cpp,BM_gpu_Trace_*(Google Benchmark, releaseopenblas-cudabuild,-DRUN_BENCHMARKS=ON, GTX 1660 SUPER):Strided_3DStrided_3DStrided_3DStrided_3DMatvec_3DMatvec_3DMatvec_3DMatvec_3DStrided_2DStrided_2DStrided_2DStrided_2DVecdot_2DVecdot_2DVecdot_2DVecdot_2DReshape_2DReshape_2DReshape_2DReshape_2DStrided_3D/Strided_2D(this PR's new kernel path, traced over non-adjacent/permuted axes) stay roughly flat regardless of size — dominated by launch overhead rather than diagonal length.Matvec_3D,Vecdot_2D, andReshape_2Dare reference/comparison shapes that do proportionally more work and scale with size as expected.ComplexDoubletracksDoubleat roughly 1.0–1.5x throughout, consistent with thecuda::std::complexreinterpret-cast path.Test plan
openblas-cpufull suite 1072/1072 passed (the GPU TU is excluded; no CPU regressions from the GPU patch).nvcc 12.6compile-check ofcuTrace_internal.cu: 0 errors..curuntime can't be exercised locally; relying on CI to runtests/gpu/.Draft — opened for review.