namespace deep_gemm {

if (num_former_iters == kNumFormerIters) {

inner_launch_k_iterations(func, cute::Int<kNumFormerIters>{});

return;

DG_STATIC_ASSERT(BLOCK_M % WGMMA::M == 0, "Invalid block size");

// Shared memory

static constexpr uint32_t SMEM_D_SIZE = BLOCK_M * (BLOCK_N + BLOCK_N_PADDING) * sizeof(__nv_bfloat16);

static constexpr uint32_t SMEM_A_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(__nv_fp8_e4m3);

static constexpr uint32_t SMEM_B_SIZE_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(__nv_fp8_e4m3);

// Fill shared memory pointers

#pragma unroll

smem_a[i] = reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + SMEM_D_SIZE + i * SMEM_A_SIZE_PER_STAGE);

smem_b[i] = reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + SMEM_D_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE);

smem_scales_a[i] = reinterpret_cast<float*>(smem_buffer + SMEM_D_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE) + i * SMEM_SCALES_A_SIZE_PER_STAGE);

// Fill barriers

auto barrier_start_ptr = reinterpret_cast<Barrier*>(reinterpret_cast<uint8_t*>(smem_scales_b) + SMEM_SCALES_B_SIZE);

#pragma unroll

full_barriers[i] = barrier_start_ptr + i;

empty_barriers[i] = barrier_start_ptr + kNumStages + i;

}

// NOTES: we always use `lane_idx` to arrive for the `lane_idx`-th CTA in the cluster,

// even with TMA multicast disabled, we want to make the behavior aligned

#pragma unroll

full_barriers[i]->init(1);

empty_barriers[i]->init(kNumTMAMulticast * kNumMathThreads / 32);

}

// For pipeline unrolling

struct DivisibleK {};

struct NotDivisibleK {};

constexpr bool kShouldOptimize = BLOCK_K / constexpr_gcd(BLOCK_K, BLOCK_N) <= 4 and not kMustUseUniformedScaleB;

// NOTES: for too-many branches (> 5), we disable this optimization

// Otherwise, the compiler must know the dynamic variable `num_former_iters`'s real value

} else {

}

}, func, kShouldOptimize ? num_former_iters : 0);

};

// Register reconfigurations

// Block scheduler

uint32_t m_block_idx, n_block_idx;

if (threadIdx.x == kNumMathThreads) {

// Persistently schedule over blocks

while (scheduler.get_next_block(m_block_idx, n_block_idx)) {

// Assign TMA multicast number into A and B

// NOTES: there may be additional odd rows/columns or cases where multicast is not possible.

// Issue TMA A

auto& full_barrier = *full_barriers[s];

tma_copy(&tensor_map_a, reinterpret_cast<uint64_t*>(&full_barrier),

smem_a[s], k_idx, scheduler.get_global_idx(shape_m, BLOCK_M, m_block_idx),

num_tma_multicast_a);

empty_barriers[s]->wait((scheduler.current_iter * kNumIterations + k_iter + 1) & 1);

full_barriers[s]->arrive();

}

}

// To safely deconstruct distributed shared barriers, we need another round of empty waits

cutlass::arch::NamedBarrier(kNumMathThreads).sync();

// Accumulation for WGMMA or CUDA promotion

DG_STATIC_ASSERT(BLOCK_M % WAVE_BLOCK_M == 0, "Invalid block sizes");

float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0};

// Empty barrier arrival

if constexpr (kNumTMAMulticast == 1) {

lane_idx == 0 ? empty_barriers[s]->arrive() : void();

} else {

};

// Launch MMAs

#pragma unroll

// Read B scales

float scale_b_0 = ld_shared(smem_scales_b + k_iter * kNumStages + s), scale_b_1;

// NOTES: even some blocks do not need to read the second row, but we still load one to align with other blocks

// Commit WGMMA instructions

#pragma unroll

warpgroup_fence_operand(accum[i]);

warpgroup_arrive();

#pragma unroll

auto desc_a = make_smem_desc(smem_a[s] + (math_wg_idx * WGMMA::M + m_offset) * BLOCK_K + k * WGMMA::K, 1);

auto desc_b = make_smem_desc(smem_b[s] + k * WGMMA::K, 1);

WGMMA::wgmma(desc_a, desc_b, accum, k);

}

warpgroup_commit_batch();

#pragma unroll

warpgroup_fence_operand(accum[i]);

warpgroup_wait<0>();

auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx;

#pragma unroll

// NOTES: for unrolled `num_former_iters` cases, we expect the compiler to automatically make it a constant

bool predicate = kMustUseUniformedScaleB or i < num_former_iters;

shifted_accum[i * 4 + 0] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 0];

full_barriers[s]->wait((scheduler.current_iter * kNumIterations + k_iter) & 1);

empty_barrier_arrive(s);

}

// TMA checks

constexpr uint32_t kNumElemBytes = sizeof(nv_bfloat16);

DG_STATIC_ASSERT(BLOCK_N % TMA_D_BLOCK_N == 0 and BLOCK_N / TMA_D_BLOCK_N <= 32,

"Unaligned TMA store or too many TMA store instructions");

DG_STATIC_ASSERT(TMA_D_BLOCK_N % 8 == 0, "Invalid TMA block N");

"Swizzling and padding are not compatible");

// Wait last TMA store to be finished

uint8_t* smem_ptr = nullptr;

if constexpr (kSwizzleDMode > 0) {

// Calculate the swizzling atom offset and in-atom offset

auto atom_offset = i / (TMA_D_BLOCK_N / 8), in_atom_offset = i % (TMA_D_BLOCK_N / 8);

// Calculate the index of the bank group to be written in the atom

}; // namespace deep_gemm

// Only used for masked layout

uint32_t curr_group_idx, curr_cumsum;

int* grouped_layout = nullptr) {

num_aligned_m_blocks = ceil_div(shape_m, BLOCK_M);

if constexpr (kGemmType == GemmType::Normal) {

}

}

__device__ __forceinline__ bool is_tma_multicast_valid(const uint32_t& m_block_idx) const {

if (num_blocks_in_group == 1)

return false;

}

}

uint32_t& m_block_idx, uint32_t& n_block_idx) {

DG_STATIC_ASSERT(kNum1DBlocksPerGroup % kNumTMAMulticast == 0, "Invalid group size");

}

template <bool kIgnoreGroupedForGroupedContiguous=true>

const uint32_t& block_idx, const uint32_t& m_block_idx=0) {

if constexpr (kGemmType == GemmType::Normal) {

return block_idx * block_size;

'--ptxas-options=--register-usage-level=10' +

(',--verbose' if 'DG_JIT_PTXAS_VERBOSE' in os.environ else ''),

# Suppress some unnecessary warnings, such as unused variables for certain `constexpr` branch cases

@staticmethod

def include_dirs() -> List[str]: