[QST] Why do we only need the result of the last k-loop in `cute::gemm` dispatch-5?
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 10.5k
- Forks
- 2.1k
- Avg merge
- 3d 11h
- Merged PRs (30d)
- 7
Description
the original code is as follow
// Dispatch [5]: (V,M,K) x (V,N,K) => (V,M,N)
template <class MMA,
class TD, class DLayout,
class TA, class ALayout,
class TB, class BLayout,
class TC, class CLayout,
__CUTE_REQUIRES(DLayout::rank == 3 && is_rmem<TD>::value &&
ALayout::rank == 3 && is_rmem<TA>::value &&
BLayout::rank == 3 && is_rmem<TB>::value &&
CLayout::rank == 3 && is_rmem<TC>::value)>
CUTE_HOST_DEVICE
void
gemm(MMA_Atom<MMA> const& mma,
Tensor<TD, DLayout> & D, // (V,M,N) Logical data
Tensor<TA, ALayout> const& A, // (V,M,K) Logical data
Tensor<TB, BLayout> const& B, // (V,N,K) Logical data
Tensor<TC, CLayout> const& C) // (V,M,N) Logical data
{
CUTE_STATIC_ASSERT_V(size<1>(A) == size<1>(C)); // AM == CM
CUTE_STATIC_ASSERT_V(size<1>(B) == size<2>(C)); // BN == CN
CUTE_STATIC_ASSERT_V(size<2>(A) == size<2>(B)); // AK == BK
CUTE_STATIC_ASSERT_V(size<0>(C) == size<0>(D) && size<1>(C) == size<1>(D) && size<2>(C) == size<2>(D));
auto K = size<2>(A);
CUTE_UNROLL
for (int k = 0; k < K; ++k) {
gemm(mma, D, A(_,_,k), B(_,_,k), C);
}
}
In the for-loop of dim-k (D = Ak x Bk + C), the result of the last calculation will override the result of the previous one.
For example, the following code
auto tA = make_tensor<int>(make_layout(make_shape(_1{}, _1{}, _2{}))); // V=1, M=1, K=2
auto tB = make_tensor<int>(make_layout(make_shape(_1{}, _1{}, _2{}))); // V=1, N=1, K=2
auto tC = make_tensor<int>(make_layout(make_shape(_1{}, _1{}, _1{}))); // V=1, M=1, N=1
auto tD = make_tensor<int>(make_layout(make_shape(_1{}, _1{}, _1{}))); // V=1, M=1, N=1
fill(tA, 1); // A = [1, 1]
fill(tB, 1); // B = [1, 1]
fill(tC, 10); // C = [10]
gemm(tD, tA, tB, tC);
print_tensor(tD); // should be 1 x 1 + 1 x 1 + 10 = 12
will get
ptr[32b](0x7fff402f0840) o (_1,_1,_1):(_0,_0,_0):
11
instead of 12.
This is only correct when C and D point to the same register that result will be accumulated properly. Is this a restriction for calling this function(cute::gemm dispatch-5)?
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the dispatch-5 gemm implementation and its unrolled dim-k loop, then reproduce the issue with the provided 1x1x2 tensor example. Check how the loop passes C and D across iterations and determine whether separate C and D tensors are supported; done means the expected accumulated result is established or the restriction is documented.
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
- Mostly clear
- Newbie friendliness
- 35/100