[QST] Should `MMA_Atom::get_layoutC_TV()` mask or ignore `ThrK`?

What is your question?

I have a question about the implementation of MMA_Atom::get_layoutC_TV() in include/cute/atom/mma_atom.hpp.

For A and B, get_layoutA_TV() and get_layoutB_TV() explicitly mask the irrelevant thread-layout dimension before composing with thrfrg_A / thrfrg_B. For example, A masks the N dimension and B masks the M dimension.

However, get_layoutC_TV() currently composes the thread index mapping directly:

auto thridx_2_thrid = composition(
    make_layout(make_shape(size(thr_layout_vmnk_), Int<1>{}),
                make_stride(Int<1>{}, Int<0>{})),
    right_inverse(make_layout(thr_layout_vmnk_, complement(thr_layout_vmnk_))));

return thrfrg_C(ref_C).compose(thridx_2_thrid, _);

My understanding is that thridx_2_thrid maps:

thr_idx -> logical thread id in ThrLayoutVMNK

where that logical id still includes the ThrK coordinate. Since C logically depends only on (M,N) and not on K, I would expect get_layoutC_TV() to mask or ignore ThrK, similar to how partition_C() does.
partition_C() appears to explicitly drop ThrK:

auto thr_vmn = make_coord(get<0>(thr_vmnk_),
                          make_coord(get<1>(thr_vmnk_), get<2>(thr_vmnk_)));

return thr_tensor(thr_vmn, make_coord(_, repeat<rank<1,1>(thr_tensor)>(_)));

So the question is:
Is get_layoutC_TV() intended to support ThrLayoutVMNK with ThrK > 1, or is there an implicit invariant that ThrK == 1 for all valid uses of get_layoutC_TV()?

I could not find a static assertion enforcing ThrK == 1. Also, make_tiled_mma() seems to accept a 3D thread layout where the third mode is not forced to 1.

Here is a small reproducer showing that get_layoutC_TV() and partition_C() disagree when ThrK > 1:

#include <cute/tensor.hpp>
#include <cute/atom/mma_atom.hpp>
#include <cstdio>

using namespace cute;

template<int K>
void check() {
  auto mma = make_tiled_mma(
      UniversalFMA<float, float, float>{},
      Layout<Shape<_2, _3, Int<K>>>{});

  float data[6] = {};
  auto c = make_tensor(make_gmem_ptr(data), make_layout(Shape<_2, _3>{}));
  auto tv = mma.get_layoutC_TV();

  int mismatches = 0;

  for (int t = 0; t < size(mma); ++t) {
    auto part = mma.get_slice(t).partition_C(c);

    auto actual = &part(0, 0, 0) - data;
    auto reported = tv(t, 0);

    mismatches += actual != reported;

    if (t == 0 || t == 6) {
      std::printf(
          "K=%d thread=%d partition_C=%td get_layoutC_TV=%d\n",
          K, t, actual, int(reported));
    }
  }

  std::printf("K=%d mismatches=%d/%d\n", K, mismatches, int(size(mma)));
}

int main() {
  check<1>();
  check<4>();
}

Compiled with:

g++ -std=c++17 -Iinclude -I/usr/local/cuda/include repro.cpp -o repro

Observed output:

K=1 thread=0 partition_C=0 get_layoutC_TV=0
K=1 mismatches=0/6
K=4 thread=0 partition_C=0 get_layoutC_TV=0
K=4 thread=6 partition_C=0 get_layoutC_TV=6
K=4 mismatches=18/24

For ThrK == 1, get_layoutC_TV() matches partition_C(). For ThrK == 4, many threads report different C offsets from get_layoutC_TV() than the offsets selected by partition_C().

I also checked the current CUTLASS tree and the concrete call sites I found appear to use thread layouts with ThrK == 1, for example SM90 builder configurations such as (1,1,1) or (2,1,1), and tutorial/test layouts such as (2,2,1), (16,16,1), etc.

So maybe this is simply an undocumented invariant. But if ThrK > 1 is considered a valid TiledMMA configuration, should get_layoutC_TV() mask/ignore ThrK to be consistent with partition_C()?
Thanks.

Start with include/cute/atom/mma_atom.hpp, comparing get_layoutC_TV() with partition_C() and the existing A/B layout helpers. Compile and run the supplied reproducer for K=1 and K=4, then trace the relevant make_tiled_mma call sites. Done means documenting or testing the valid ThrK invariant and resolving the reported disagreement for supported configurations.