NVIDIA / NVIDIA/cutlass

[BUG] Non-ZFILL CP_ASYNC copy trait cause buffer overwriting with predicated copy

Open
#1,716 7 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

? - Needs Triage bug inactive-30d inactive-90d
Dominant language
C++
Stars
10.5k
Forks
2.1k
Avg merge
3d 11h
Merged PRs (30d)
7

Description

Describe the bug

#include "cute/tensor.hpp"

using namespace cute;

__global__ void kernel(int *gmem) {
  int tid = threadIdx.x;
  gmem[tid * 4 + 0] = tid * 4 + 0;
  gmem[tid * 4 + 1] = tid * 4 + 1;
  gmem[tid * 4 + 2] = tid * 4 + 2;
  gmem[tid * 4 + 3] = tid * 4 + 3;
  __syncthreads();

  __shared__ int smem[128];
  smem[tid * 4 + 0] = -1;
  smem[tid * 4 + 1] = -1;
  smem[tid * 4 + 2] = -1;
  smem[tid * 4 + 3] = -1;
  __syncthreads();

  auto g = make_tensor(make_gmem_ptr(gmem), Layout<Shape<_1, _128>>{});
  auto s = make_tensor(make_smem_ptr(smem), Layout<Shape<_1, _128>>{});

  __syncthreads();
  if (thread0()) {
    print_tensor(s);
  }

  using CopyAtom = Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL<uint128_t>, int>;
  // using CopyAtom = Copy_Atom<UniversalCopy<uint128_t>, int>;

  auto tiled_copy = make_tiled_copy(CopyAtom{}, Layout<_32>{}, Layout<_4>{});
  auto thr_copy = tiled_copy.get_thread_slice(tid);
  auto tG = thr_copy.partition_S(coalesce(g));
  auto tS = thr_copy.partition_D(coalesce(s));
  auto p = make_tensor<bool>(size<1>(tG));
  for (int i = 0; i < size(p); i++) {
    p(i) = tid < 2;  // only activate two threads
  }

  copy_if(tiled_copy, p, tG, tS);
  cp_async_fence();
  cp_async_wait<0>();
  __syncthreads();

  if (thread0()) {
    print_tensor(s);
  }
}

int main() {
  int *buffer;
  cudaMalloc(&buffer, 32 * sizeof(int));
  kernel<<<1, 32>>>(buffer);
  cudaDeviceSynchronize();

  return cudaGetLastError() == cudaSuccess;
}

nvcc cp_async.cu -Iinclude -std=c++20 --expt-relaxed-constexpr -gencode=arch=compute_80,code=sm_80 && ./a.out

produces

smem_ptr[32b](0x7f6985000000) o (_1,_128):(_0,_1):
   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1   -1
smem_ptr[32b](0x7f6985000000) o (_1,_128):(_0,_1):
    0    1    2    3    4    5    6    7    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0    0

Expected behavior

When _ZFILL is not requested, those predicated out values should not be touched.
This might cause

  1. buffer overwritten when dst buffer is not large enough, but have another __shared__ buffer declare later and is large enough.
  2. memory fault when oob write.

This is because

https://github.com/NVIDIA/cutlass/blob/f93a69134ec8259fd235f220209d6f8734a5cb06/include/cute/atom/copy_traits_sm80.hpp#L77-L82

re-dispatch to _ZFILL trait silently and this will generally cause very very very subtle bug when the user is expecting an async version of Copy_Atom<UniversalCopy<uint128_t>, T> as a simple substitute, but not the ignore-src behavior!

Since the _ZFILL variants exists, this implicit behavior should be removed.

The only workaround is to replace copy_if as follows

// assuming only one iter mode
if (p(i)) {
  copy(tiled_copy, tG(_, i), tS(_, i));
}

Environment details (please complete the following information):
f93a69134ec8259fd235f220209d6f8734a5cb06

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the provided CUDA reproducer and inspect include/cute/atom/copy_traits_sm80.hpp at the linked dispatch around lines 77-82. Trace how copy_if handles predicated non-ZFILL CP_ASYNC copies, then verify that inactive values leave shared memory unchanged and do not write out of bounds when the reproducer is run with the shown nvcc command.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp
Domain
hpc
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Clearly specified
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.