Repository: rust
Issue: Tracking Issue for NVPTX shared memory
Labels: T-libs-api,O-NVPTX,C-tracking-issue

<!--
Thank you for creating a tracking issue!

Tracking issues are for tracking a feature from implementation to stabilization.

Make sure to include the relevant RFC for the feature if it has one.

If the new feature is small, it may be fine to skip the RFC process. In that
case, you can use `issue = "none"` in your initial implementation PR. The
reviewer will ask you to open a tracking issue if they agree your feature can be
added without an RFC.
-->

Feature gate: `#![feature(stdarch_nvptx)]` (probably; see https://github.com/rust-lang/rust/issues/111199)

This is a tracking issue for access to shared memory from NVPTX device code.

There are two flavors of shared memory in NVPTX (see [this NVIDIA Technical Blog (2013)](https://developer.nvidia.com/blog/using-shared-memory-cuda-cc/) for a friendly introduction):
* **static** shared memory provides a fixed size amount of shared memory independent of the block size. There can be multiple static shared arrays (declared in CUDA C++ using the `__shared__` attribute).
* **dynamic** shared memory provides one base pointer, with a length (in bytes) that must be specified in the kernel launch parameters. If multiple arrays are needed, the single buffer must be manually partitioned (in CUDA C++).

In practice, the required amount of shared memory for GPU algorithms will depend on the block size. Since many domains expect to choose the block size at run-time, *static* shared memory is often considered to be of limited use.

<!--
Include a short description of the feature.
-->

### Public API

<!--
For most library features, it'd be useful to include a summarized version of the public API.
(E.g. just the public function signatures without their doc comments or implementation.)
-->
Note that other `core::arch::nvptx` intrinsics are covered by https://github.com/rust-lang/rust/issues/111199.

```rust
use core::arch::nvptx;

#[no_mangle]
pub unsafe extern "ptx-kernel" fn reverse_kernel(n: u32, a: *mut f32) {
    let t = nvptx::_thread_idx_x() as usize;
    let tr = n as usize - t - 1;
    let (s, sn) = nvptx::_dynamic_smem(); // <--- this issue
    let s = s as *mut f32;
    assert_eq!(sn as u32, n * 4); // requirement for the algorithm below to be correct
    *s.add(t) = *a.add(t);
    nvptx::_syncthreads();
    *a.add(t) = *s.add(tr);
}
```

#### A possible implementation for dynamic shared memory

Shared memory can be exposed using inline assembly ([reference](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#special-registers-dynamic-smem-size)).
```rs
core::arch::global_asm!(".extern .shared .align 16 .b8 _shared_data[];");

#[inline(always)]
pub fn _dynamic_smem() -> (*mut u8, u32) {
    // Dynamic shared memory.
    let size: u32;
    let saddr: u64;
    let ptr: *mut u8;
    unsafe {
        asm!("mov.u32 {}, %dynamic_smem_size;", out(reg32) size);
        asm!("mov.u64 {}, _shared_data;", out(reg64) saddr);
        asm!("cvta.shared.u64 {ptr}, {saddr};", ptr = out(reg64) ptr, saddr = in(reg64) saddr);
    }
    (ptr, size)
}
```

### Steps / History

<!--
For larger features, more steps might be involved.
If the feature is changed later, please add those PRs here as well.
-->

- [ ] Implementation: #...
- [ ] Final comment period (FCP)[^1]
- [ ] Stabilization PR

<!--
Once the feature has gone through a few release cycles and there are no
unresolved questions left, the feature might be ready for stabilization.

If this feature didn't go through the RFC process, a final comment period
(FCP) is always needed before stabilization. This works as follows:

A library API team member can kick off the stabilization process, at which point
the rfcbot will ask all the team members to verify they agree with
stabilization. Once enough members agree and there are no concerns, the final
comment period begins: this issue will be marked as such and will be listed
in the next This Week in Rust newsletter. If no blocking concerns are raised in
that period of 10 days, a stabilization PR can be opened by anyone.
-->

### Unresolved Questions

<!--
Include any open questions that need to be answered before the feature can be
stabilised. If multiple (unrelated) big questions come up, it can be a good idea
to open a separate issue for each, to make it easier to keep track of the
discussions.

It's useful to link any relevant discussions and conclusions (whether on GitHub,
Zulip, or the internals forum) here.
-->

- Do we prefer the `_dynamic_smem() -> (*mut u8, u32)` or should we have two separate intrinsics for accessing the base pointer and the size?
- Do we want to use `usize` for the length even though NVPTX does not (yet, and maybe ever) support that? I'm not sure about other vendors, if we're looking for the intrinsics to be as close as possible.
- Is it desirable to expose static shared memory or shall we focus on dynamic (which is less intrusive to implement, but may require more thinking to launch kernels with valid parameters)?
- It has also been suggested that this be implemented as an LLVM intrinsic by introducing `llvm.nvvm.read.ptx.sreg.dynamic_smem_size`. That could in principle allow accessing the shared memory base pointer without introducing a new symbol (`_shared_data above`). I lean toward using the inline assembly for now and possibly migrating it later, but no other inline assembly is used in `core::arch::nvptx` and I don't know if it should be kept that way.
- Do we want to tackle launch bounds in this issue? If so, it should not be under the `stdarch_nvptx` feature gate since it would need to verify launch parameters on the host (better user experience) or handle failure on the device (i.e., prevent UB if an insufficient amount of shared memory is provided). I prefer to handle only the unsafe primitives here since there are many open questions about such ergonomics.
- Should this issue consider partitioning dynamic shared memory into parts (e.g., multiple arrays or more general data structures, parametrized by block size)? This interacts with more types and would also make it inappropriate for the `stdarch_nvptx` feature gate.

[^1]: https://std-dev-guide.rust-lang.org/feature-lifecycle/stabilization.html