Fix expression slices after dropping a dimension - #1254
Arthur031221 wants to merge 3 commits into
Conversation
 |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Pass slice starts through runtime JIT storage and version the generated class name to exclude kernels built with the old indexing. Cover reuse across start offsets and ordinary CUDA expression assignments.
|
/build |
| #ifdef MATX_EN_JIT | ||
| struct JIT_Storage { | ||
| typename detail::inner_storage_or_self_t<detail::base_type_t<T>> op_; | ||
| cuda::std::array<shape_type, T::Rank()> starts_; |
There was a problem hiding this comment.
For the JIT path starts and strides are intentionally encoded in the name so they can be compile-time parameters. What's the reason this is runtime now?
There was a problem hiding this comment.
I moved them to runtime storage to avoid cache collisions. That was a broader change than needed: on the upstream header, the class name records dimensions and sizes but omits starts and strides.
The generated slice now uses constexpr starts and strides, with their resolved values in the versioned class name. JIT launch storage now contains the nested operator, replacing the storage layout described in the PR body.
The JIT-enabled slice and slice-stride targets passed 62 and 20 tests locally. The cache regression varies starts on retained and dropped axes and strides for two-element outputs.
|
Coverage is 93.598% — Arthur031221:fix-expression-slice-dropped-dimensions into NVIDIA:main. No base build found for |
| for (int i = 0; i < input_rank; i++) { | ||
| params_str += std::format("b{}_", starts_[i]); | ||
| if constexpr (!cuda::std::is_same_v<StrideType, NoStride>) { | ||
| params_str += std::format("t{}_", strides_[i]); | ||
| } | ||
| } |
There was a problem hiding this comment.
When a JIT workload moves a slice window across an expression, each new start or stride changes the generated class name and cache key. Even if the output shape stays the same, each position needs a separate NVRTC compilation and persistent cubin rather than reusing a kernel with runtime offsets. This can slow repeated windowed slicing and grow the disk cache.
Knowledge Base Used: CUDA and JIT executors
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
This branch has not been deployed
3 participants
Users slicing an expression after dropping an earlier dimension get values from the wrong offset. For a fixed 3 by 5 input,
slice<1>(input + 0, {1, 2}, {matxDropDim, matxEnd})returned11, 12, 13instead of12, 13, 14. Adding a stride to this expression failed to compile.SliceOpread start offsets and calculated strided output sizes using an output dimension index where it needed an input dimension index. Its strided overload declared a reference member that it could not initialize. The fix stores strides by value and includes them in JIT argument storage, which the generated operator reads.The new test checks literal values for ordinary and strided expression slices, the
downsamplewrapper on an expression, a tensor-view control, and JIT execution of both slice expressions. From the worktree,build/test/test_00_operators_slice_test --gtest_brief=1passed 49 of 49 tests andbuild-jit/test/test_00_operators_slice_test --gtest_brief=1passed 61 of 61 tests on an RTX 5090 with CUDA 13.2.51. The slice-stride, slice-and-reduce, slice-and-reshape, and up/downsample targets passed 16 of 16, 48 of 48, 8 of 8, and 96 of 96 tests. The full test suite was not built.