Any example code for SM90_TMA_REDUCE_ADD?

[ziyuhuang123]

Any example code for SM90_TMA_REDUCE_ADD? Thanks!

I realized that, several SMEM blocks will be reduced to one block in global? How to realize it?

[AdityaKane2001]

I believe the intended application is to have a single CTA write out multiple tensors to GMEM (i.e. reduce multiple tensors). If I understand correctly, multiple CTAs cannot write to the same location in GMEM, or more specifically if the example in the tma/ directory can support that without considerable changes.

On that note, one quick question, if I had to load multiple tensor tiles from GMEM using the TMA (for example in an addition kernel, where one must copy two input tiles and store to one output tile), what is the preferred way? Is it simply creating two barriers, and issuing two loads from thread 0? Or would it be to load the two tiles from two different threads/warps (say warp_idx == 0 and warp_idx == 1)?

using SharedStorage = TwoTensorSharedStorageTMA<Element, SmemLayout>; // Note the updated definition.

extern __shared__ char shared_memory[];
SharedStorage &shared_storage =
    *reinterpret_cast<SharedStorage *>(shared_memory);

// Define smem tensor
Tensor sS1 =
    make_tensor(make_smem_ptr(shared_storage.smem1.data()), smemLayout);
Tensor sS2 =
    make_tensor(make_smem_ptr(shared_storage.smem2.data()), smemLayout);

// Get mbarrier object and its value type
auto &mbarrier1 = shared_storage.mbarrier1;
auto &mbarrier2 = shared_storage.mbarrier2;
using BarrierType = cutlass::arch::ClusterTransactionBarrier::ValueType;
static_assert(cute::is_same_v<BarrierType, uint64_t>,
              "Value type of mbarrier is uint64_t.");

// Constants used for TMA
const int warp_idx = cutlass::canonical_warp_idx_sync();
const bool lane_predicate = cute::elect_one_sync();
constexpr int kTmaTransactionBytes =
    sizeof(ArrayEngine<Element, size(SmemLayout{})>);

// Prefetch TMA descriptors for load and store
if (warp_idx == 0 && lane_predicate) {
  prefetch_tma_descriptor(tmaLoad1.get_tma_descriptor());
  prefetch_tma_descriptor(tmaLoad2.get_tma_descriptor());
  prefetch_tma_descriptor(tmaStore.get_tma_descriptor());
}

// Get CTA view of gmem tensor
auto blkCoord = make_coord(blockIdx.x, blockIdx.y);

Tensor mS1 = tmaLoad1.get_tma_tensor(shape(gmemLayout));
Tensor gS1 = local_tile(mS1, tileShape, blkCoord);
auto cta_tmaS1 = tmaLoad1.get_slice(Int<0>{});

Tensor mS2 = tmaLoad2.get_tma_tensor(shape(gmemLayout));
Tensor gS2 = local_tile(mS2, tileShape, blkCoord);
auto cta_tmaS2 = tmaLoad2.get_slice(Int<0>{});

if (warp_idx == 0 and lane_predicate) {
  mbarrier1.init(1 /* arrive count */);
  mbarrier1.arrive_and_expect_tx(kTmaTransactionBytes);
  copy(tmaLoad1.with(reinterpret_cast<BarrierType &>(mbarrier1)),
       cta_tmaS1.partition_S(gS1), cta_tmaS1.partition_D(sS1));

  mbarrier2.init(1 /* arrive count */);
  mbarrier2.arrive_and_expect_tx(kTmaTransactionBytes);
  copy(tmaLoad2.with(reinterpret_cast<BarrierType &>(mbarrier2)),
       cta_tmaS2.partition_S(gS2), cta_tmaS2.partition_D(sS2));
}
__syncthreads();

mbarrier1.wait(0 /* phase */);
mbarrier2.wait(0 /* phase */);

/cc @jayhshah