Can TMA support copy part of data? like copy_if.

[defei-coder]

I want to use TMA copy part of data.
For example:
host code below
// bM=64, bK=128 auto tileShapeQ = make_shape(bM{}, bK{}); Layout gmemLayoutQ = make_layout(make_shape(M, K, H, B), make_stride(K * H, 1, K, H * M * K)); Tensor gQ = make_tensor(ptrQ, gmemLayoutQ); auto tmaQ = make_tma_copy(SM90_TMA_LOAD{}, gQ, smemLayoutQ, tileShapeQ, Int<1>{}); // smemLayoutQ is created with tileShapeQ
Device code:
Tensor mQ = tmaLoadQ.get_tma_tensor(shape(gmemLayoutQ)); Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), smemLayoutQ); auto cta_tmaQ = tmaQ.get_slice(0); Tensor tQgQX = cta_tmaQ.partition_S(gQ); Tensor tQsQX = cta_tmaQ.partition_D(sQ); if (warp_idx == 0 and lane_predicate) { cute::set_barrier_transaction_bytes(tma_load_mbar, kTmaTransactionBytes); copy(tma_load_x.with(tma_load_mbar, mcast_mask_a), tQgQX, tQsQX); }
By above code, I can copy copy bMbK datas by TMA_LOAD.
But now I want to copy mbK datas which m less than bM. I try some ways but failed. I confused whether TMA can do this.
I want to do like the below code(cp.async)
// S is a shape of [n, m, k] tensor, just copy part of S if (get<0>(identity_MN(0, m, 0)) < max_MN) { #pragma unroll for (int k = 0; k < size<2>(S); ++k) { if (Is_even_K || predicate_K(k)) { cute::copy(tiled_copy, S(_, m, k), D(_, m, k)); } } }

[thakkarV]

TMA supports this (predication for OOB reads) out of the box. you should not have to do anything special

[defei-coder]

Thanks for your reply! You mean TMA default support OOB, actually I want to do like this: the full shape is 64*64, but I want to load 32*64 elements in block0 and load 12*64 elements in block1 and others in block2. Also I notice TMA split the global tensor divided by number of blocks, the offset of each block is blockIdx.x * ele_per_block(which ele_per_block = all_elememts / gridDims.x), but I want to process different length of elements in each block.

[thakkarV]

the full shape is 6464, but I want to load 3264 elements in block0 and load 12*64 elements in block1 and others in block2.

Okay, create an affine layout to represent your tensors. you clearly have some layout in mind that is not 2D. For TMA, the gmem layout of your tensor must be affine. If it is not, then TMA cannot help you. If you have OOB elements, they will be zero filled by TMA but each CTA must have a consistent view of the smem "vector" it is going to load from gmem