Missing stride_ in vertical pointer arithmetic when advance_rank=0 (PredicatedTileAccessIterator)

[piaoyanglink]

I‘m researching cutlass！when advance_rank == 0 and tile_offset.strided() != 0, the pointer offset calculation is missing params_.stride_, which will lead to incorrect address computation.Code Location (Provide key snippets) is add_tile_offsetin predicated_tile_access_iterator.h
the code is below:

pointer_ += Shape::kStrided * tile_offset.strided() * sizeof_bits::value / 8;

maybe it's like this:

pointer_ += Shape::kStrided * tile_offset.strided() * params_.stride_ * sizeof_bits::value / 8;

This is my speculation; I’m not sure if it’s correct, and I hope to get a response.

[Arthur031221]

Your reading is right about the pointer, but that line alone doesn't make the move safe (details below). I checked it on main (0b55a2f) with a small kernel. It sets up a PredicatedTileAccessIterator over a pitch-linear float tensor with stride 1000, tile shape PitchLinearShape<64, 8> and PitchLinearStripminedThreadMap<Shape, 32, 1>. It then prints thread 0's element offset after add_tile_offset, next to an iterator constructed directly at the tile it should have moved to. The extent is 256 x 64, a multiple of the tile, so the first (residue) tile is a full tile and those reference positions are the right targets. This was built with CUDA 13.2 and run on an RTX 5090:

AdvanceRank=0, thread 0
  start {0,0}, add_tile_offset({1,1})          offset     72 = (contiguous 72, strided 0)
  start {0,0}, add({1,0}) then add({1,1})      offset    130 = (contiguous 130, strided 0)
  start {64,8} (reference, no add)             offset   8064 = (contiguous 64, strided 8)
  start {128,8} (reference for 2nd row)        offset   8128 = (contiguous 128, strided 8)
AdvanceRank=1, thread 0
  start {0,0}, add_tile_offset({1,1})          offset   8064 = (contiguous 64, strided 8)
  start {0,0}, add({0,1}) then add({1,1})      offset  16016 = (contiguous 16, strided 16)
  start {64,8} (reference, no add)             offset   8064 = (contiguous 64, strided 8)
  start {64,16} (reference for 2nd row)        offset  16064 = (contiguous 64, strided 16)
no error

With advance rank 0, the first add_tile_offset({1,1}) goes through the residue branch you quoted. It moves 8 elements along the contiguous dimension instead of 8 rows. With your line (the same expression times params_.stride_), that first row becomes 8064, which matches the reference. The advance rank 1 case is already right in that branch.

The steady-state branch a few lines below has the same problem in more places. With advance rank 0 it is missing both the stride and the element size (pointer_ is a char *), and lands at 130 instead of 8128. With advance rank 1 only the element size is missing, and it lands at 16016 instead of 16064. The element-size part is already #3487. That issue describes the residue branch as correct, but for the strided term with advance rank 0 it isn't. The history explains it: the fix for #931 in 3.1 (f079619) added the sizeof_bits factor to the two residue-branch lines and left the steady-state branch alone.

Fixing the pointer is not enough, though. add_tile_offset never adds the cross-rank part of the move to the offset the predicates are computed from. With advance rank 0 the steady-state guard only checks the strided coordinate, against the old rows. With your line applied, extent 256 x 12, and the same add_tile_offset({1,1}), thread 0 now points at rows 8 to 15, but all 16 of its accesses report valid(). The reference iterator at {64, 8} masks rows 12 to 15:

  access  6 row 11 moved valid 1 | reference row 11 valid 1
  access  7 row 11 moved valid 1 | reference row 11 valid 1
  access  8 row 12 moved valid 1 | reference row 12 valid 0
  access  9 row 12 moved valid 1 | reference row 12 valid 0

Every one of the 32 threads has 8 accesses past row 11 that are still marked valid. A diagonal move near the edge of the tensor can read or write past the extent, so a proper fix has to update the predicates too, not just the pointer.

As for why nobody has hit this: in #931, hwu36 said this cross term is never used in GEMM, because the mainloop only moves along the k dimension. It only matters if you call add_tile_offset yourself with a non-zero offset on the rank the iterator does not advance along.

Both kernels, if you want to run them (nvcc -std=c++17 -arch=sm_120 -I <cutlass>/include t.cu, the same for pred.cu, with sm_120 changed to your GPU's architecture). pred.cu only shows the problem with your line applied; on unmodified main the moved iterator stays on rows 0 to 7, which are in bounds.

t.cu (pointer offsets)

#include <cstdio>
#include "cutlass/cutlass.h"
#include "cutlass/layout/pitch_linear.h"
#include "cutlass/transform/pitch_linear_thread_map.h"
#include "cutlass/transform/threadblock/predicated_tile_access_iterator.h"

using Shape = cutlass::layout::PitchLinearShape<64, 8>;
using ThreadMap = cutlass::transform::PitchLinearStripminedThreadMap<Shape, 32, 1>;
using Access = cutlass::Array<float, 1>;
template <int R>
using Iter = cutlass::transform::threadblock::PredicatedTileAccessIterator<
    Shape, float, cutlass::layout::PitchLinear, R, ThreadMap, Access>;

constexpr int kStride = 1000;              // leading dimension, in elements
__device__ float buf[kStride * 64];

template <int R>
__global__ void k(long long *out) {
  typename Iter<R>::Params params{cutlass::layout::PitchLinear{kStride}};
  cutlass::layout::PitchLinearCoord extent(256, 64);
  // first call: residue path
  Iter<R> a(params, buf, extent, (int)threadIdx.x, {0, 0});
  a.add_tile_offset({1, 1});
  // second call: steady-state path
  Iter<R> b(params, buf, extent, (int)threadIdx.x, {0, 0});
  if (R == 0) b.add_tile_offset({1, 0}); else b.add_tile_offset({0, 1});
  b.add_tile_offset({1, 1});
  // reference: iterator started at the tile it should have moved to
  Iter<R> c(params, buf, extent, (int)threadIdx.x, {64, 8});
  Iter<R> d(params, buf, extent, (int)threadIdx.x, R == 0 ? cutlass::layout::PitchLinearCoord{128, 8} : cutlass::layout::PitchLinearCoord{64, 16});
  if (threadIdx.x == 0) {
    out[0] = (float *)a.get() - buf;
    out[1] = (float *)b.get() - buf;
    out[2] = (float *)c.get() - buf;
    out[3] = (float *)d.get() - buf;
  }
}

template <int R> void run() {
  long long *d, h[4];
  cudaMalloc(&d, sizeof h);
  k<R><<<1, 32>>>(d);
  cudaMemcpy(h, d, sizeof h, cudaMemcpyDeviceToHost);
  auto show = [](const char *w, long long o) {
    printf("  %-44s offset %6lld = (contiguous %lld, strided %lld)\n", w, o, o % kStride, o / kStride);
  };
  printf("AdvanceRank=%d, thread 0\n", R);
  show("start {0,0}, add_tile_offset({1,1})", h[0]);
  show(R == 0 ? "start {0,0}, add({1,0}) then add({1,1})" : "start {0,0}, add({0,1}) then add({1,1})", h[1]);
  show("start {64,8} (reference, no add)", h[2]);
  show(R == 0 ? "start {128,8} (reference for 2nd row)" : "start {64,16} (reference for 2nd row)", h[3]);
  cudaFree(d);
}

int main() { run<0>(); run<1>(); printf("%s\n", cudaGetErrorString(cudaDeviceSynchronize())); }

pred.cu (predicates after the move)

#include <cstdio>
#include "cutlass/cutlass.h"
#include "cutlass/layout/pitch_linear.h"
#include "cutlass/transform/pitch_linear_thread_map.h"
#include "cutlass/transform/threadblock/predicated_tile_access_iterator.h"

using Shape = cutlass::layout::PitchLinearShape<64, 8>;
using ThreadMap = cutlass::transform::PitchLinearStripminedThreadMap<Shape, 32, 1>;
using Access = cutlass::Array<float, 1>;
using Iter = cutlass::transform::threadblock::PredicatedTileAccessIterator<
    Shape, float, cutlass::layout::PitchLinear, 0, ThreadMap, Access>;

constexpr int kStride = 1000;
__device__ float buf[kStride * 64];

__global__ void k(int rows) {
  Iter::Params params{cutlass::layout::PitchLinear{kStride}};
  cutlass::layout::PitchLinearCoord extent(256, rows);
  Iter a(params, buf, extent, (int)threadIdx.x, {0, 0});
  a.add_tile_offset({1, 1});
  Iter r(params, buf, extent, (int)threadIdx.x, {64, 8});
  int bad = 0;
  for (int i = 0; i < ThreadMap::Iterations::kCount; ++i, ++a, ++r) {
    long long off = (float *)a.get() - buf;
    bool oob = off / kStride >= rows;
    if (threadIdx.x == 0)
      printf("  access %2d row %2lld moved valid %d | reference row %2lld valid %d\n", i, off / kStride, a.valid(), ((float *)r.get() - buf) / kStride, r.valid());
    bad += (a.valid() && oob);
  }
  if (bad) printf("  thread %d: %d accesses valid() but past row %d\n", threadIdx.x, bad, rows - 1);
}
int main() { printf("extent 256 x 12, AdvanceRank=0, add_tile_offset({1,1})\n"); k<<<1, 32>>>(12); printf("%s\n", cudaGetErrorString(cudaDeviceSynchronize())); }