Skip to content

Commit 446ad8c

Browse files
xiaolil1xiaolil1
committed
refine code
1 parent 40b1b68 commit 446ad8c

1 file changed

Lines changed: 83 additions & 102 deletions

File tree

‎csrc/xpu_cutlass_fusion.cpp‎

Lines changed: 83 additions & 102 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ using TiledMma =
6666
typename TiledMMAHelper<MMA_Atom<XE_8x16x16_F32BF16BF16F32_TT>, Layout<TileShape>,
6767
Layout<Shape<_1, _8, _1>, Stride<_8, _1, _0>>>::TiledMMA;
6868
using GmemTiledCopyA = XE_2D_U16x32x32_LD_N;
69+
using GmemTiledCopyB = XE_2D_U4x32x16_LD_T;
6970
constexpr int PipelineStages = 4;
7071

7172
using MmaAtomShape = typename TiledMma::AtomShape_MNK;
@@ -96,6 +97,10 @@ static constexpr uint32_t MaxThreadsPerBlock = size(TiledMma{});
9697
using DispatchPolicy = cutlass::gemm::MainloopIntelPVCMixedPrecision<PipelineStages>;
9798
static constexpr int SubgroupSize = DispatchPolicy::SubgroupSize;
9899

100+
static constexpr auto FragsM = get<0>(SubgroupTileShape{}) / get<0>(MmaAtomShape());
101+
static constexpr auto FragsN = get<1>(SubgroupTileShape{}) / get<1>(MmaAtomShape());
102+
static constexpr auto FragmentSize = (get<0>(MmaAtomShape()) * get<1>(MmaAtomShape())) / SubgroupSize;
103+
99104
// Design Scheduler
100105
using TileScheduler_ = PersistentScheduler;
101106
static_assert(cute::is_void_v<TileScheduler_> or cute::is_same_v<TileScheduler_, PersistentScheduler>, "Intel PVC does not support specializing the tile scheduler.");
@@ -116,7 +121,6 @@ using atom_load_A = Copy_Atom<traits_load_A, ElementA>;
116121
using val_layout_load_A = decltype(make_layout(shape_div(typename traits_load_A::BlockShape{}, CopyThreadShape{})));
117122
using Copy_A = decltype(make_tiled_copy(atom_load_A{}, Layout<CopyThreadShape>{}, val_layout_load_A{}));
118123

119-
using GmemTiledCopyB = XE_2D_U4x32x16_LD_T;
120124
using StrideB = cutlass::gemm::TagToStrideB_t<cutlass::layout::ColumnMajor>;
121125
using traits_load_B = Copy_Traits<GmemTiledCopyB, StrideB>;
122126
using atom_load_B = Copy_Atom<traits_load_B, ElementB>;
@@ -176,25 +180,22 @@ class gemm_4bit_cutlass_kernel {
176180
CUTLASS_DEVICE
177181
void operator()(Params const& params, char* smem_buf) {
178182
int thread_idx = int(ThreadIdxX());
179-
180-
// Load Dequatize LUT and save to SLM, 16 for 4bits
181-
float* quant_map = reinterpret_cast<float*>(smem_buf);
182-
if (thread_idx < 16) {
183-
quant_map[thread_idx] = params.datatype[thread_idx];
184-
}
185-
barrier_arrive(3);
186-
187-
int m_coord, n_coord, l_coord;
188-
if (params.scheduler.raster_order_ == TileScheduler::RasterOrder::AlongN) {
189-
m_coord = BlockIdxY();
190-
n_coord = BlockIdxX();
191-
l_coord = BlockIdxZ();
192-
} else {
193-
m_coord = BlockIdxX();
194-
n_coord = BlockIdxY();
195-
l_coord = BlockIdxZ();
196-
}
197-
183+
const int m_coord = (params.scheduler.raster_order_ == TileScheduler::RasterOrder::AlongN)
184+
? BlockIdxY() : BlockIdxX();
185+
const int n_coord = (params.scheduler.raster_order_ == TileScheduler::RasterOrder::AlongN)
186+
? BlockIdxX() : BlockIdxY();
187+
const int l_coord = BlockIdxZ();
188+
189+
float* quant_map;
190+
{
191+
// Load Dequatize LUT and save to SLM, 16 for 4bits
192+
quant_map = reinterpret_cast<float*>(smem_buf);
193+
if (thread_idx < 16) {
194+
quant_map[thread_idx] = params.datatype[thread_idx];
195+
}
196+
barrier_arrive(3);
197+
}
198+
198199
Tensor mA_mkl = cute::get_pvc_tensor(make_shape(params.m, params.k, params.l));
199200
Tensor mB_nkl = cute::get_pvc_tensor(make_shape(params.n, params.k,1));
200201

@@ -206,7 +207,6 @@ class gemm_4bit_cutlass_kernel {
206207
clear(accumulators);
207208

208209
auto k_tile_iter = cute::make_coord_iterator(idx2crd(0, make_shape(params.k)), make_shape(params.k));
209-
int k_tile_count = ceil_div(params.k, get<2>(WorkgroupTileShape{}));
210210

211211
auto thr_copy_A = params.tiled_copy_a.get_slice(thread_idx);
212212
auto thr_copy_B = params.tiled_copy_b.get_slice(thread_idx);
@@ -247,6 +247,7 @@ class gemm_4bit_cutlass_kernel {
247247
auto pAgA = thr_prefetch_A.partition_S(gA);
248248
auto pBgB = thr_prefetch_B.partition_S(gB);
249249

250+
const int k_tile_count = ceil_div(params.k, get<2>(WorkgroupTileShape{}));
250251
const int k_reload_factor = ceil_div(params.group_size, BLK_K);
251252

252253
auto tSgS = [&](){
@@ -256,77 +257,49 @@ class gemm_4bit_cutlass_kernel {
256257

257258
}();
258259

259-
static constexpr int FragsM = get<0>(SubgroupTileShape{}) / get<0>(MmaAtomShape()); // A frags per sub_group
260-
static constexpr int FragsN = get<1>(SubgroupTileShape{}) / get<1>(MmaAtomShape()); // B frags per sub_group
261-
262-
static constexpr int FragmentSize = (get<0>(MmaAtomShape()) * get<1>(MmaAtomShape())) / SubgroupSize;
263-
264-
auto m_sg = get_sub_group_id() / ATOM_N;
265-
auto n_sg = get_sub_group_id() % ATOM_N;
266-
267-
// Represent the full output tensor
268-
Tensor mD_mnl = cute::get_pvc_tensor(make_shape(params.m, params.n, params.l));
269-
270-
// Tile the output tensor per WG and select the tile for current WG
271-
Tensor g_wg_D = local_tile(mD_mnl, take<0,2>(WorkgroupTileShape{}), make_coord(m_coord,n_coord,l_coord)); // (BLK_M,BLK_N)
272-
273-
// Tile the output tensor per SG and select tile for the current SG
274-
Tensor gD = local_tile(g_wg_D, take<0,2>(SubgroupTileShape{}), make_coord(m_sg,n_sg)); // (SG_M,SG_N)
275-
276-
auto thread_xe_store_d = params.tiled_store_d.get_thread_slice(thread_idx);
277-
Tensor tCgD = thread_xe_store_d.partition_D(gD);
278-
279260
const int k_start_idx = crd2idx((*k_tile_iter), make_shape(params.k));
280261
int prefetch_k = k_start_idx;
281262

282-
auto dequant = [](auto const& in, auto& out, auto& tCrS_input, const float* quant_map_) {
283-
constexpr auto N = decltype(cute::size<1>(in))::value;
284-
constexpr auto K = decltype(cute::size(out))::value / N;
285-
286-
using compress_type = uint32_t;
287-
constexpr auto compress_size = cute::sizeof_bits_v<compress_type> / cute::sizeof_bits_v<ElementB>;
288-
static_assert((compress_size % N) == 0);
289-
290-
constexpr auto vec_size = K / compress_size;
291-
using VecSrcType = cute::array<compress_type, vec_size>;
292-
using VecDstElemType = cute::array<ElementMMA, compress_size>;
293-
using VecDstType = cute::array<VecDstElemType, vec_size>;
294-
295-
auto s_tensor = cute::make_tensor(
296-
(VecSrcType*)(cute::raw_pointer_cast(in.data())),
297-
cute::make_shape(cute::Int<K / (compress_size * vec_size)>{}, cute::Int<N>{})
298-
);
299-
300-
auto d_tensor = cute::make_tensor(
301-
(VecDstType*)(cute::raw_pointer_cast(out.data())),
302-
cute::make_shape(cute::Int<K / (compress_size * vec_size)>{}, cute::Int<N>{})
303-
);
304-
263+
auto dequant = [&] {
264+
constexpr int N = decltype(cute::size<1>(mma_B))::value;
265+
constexpr int K = decltype(cute::size(mma_B))::value / N;
266+
//if(cute::thread0()) printf("K = %d, N = %d\n", K, N);
267+
268+
using compress_type = uint32_t;
269+
constexpr int compress_size = cute::sizeof_bits_v<compress_type> / cute::sizeof_bits_v<ElementB>;
270+
constexpr auto vec_size = K / compress_size;
271+
272+
using VecSrcType = cute::array<compress_type, vec_size>;
273+
using VecDstElemType = cute::array<ElementMMA, compress_size>;
274+
using VecDstType = cute::array<VecDstElemType, vec_size>;
275+
276+
auto s_tensor = cute::make_tensor((VecSrcType*)(cute::raw_pointer_cast(dequant_frag.data())), cute::make_shape(cute::Int<K / (compress_size * vec_size)>{}, cute::Int<N>{}));
277+
auto d_tensor = cute::make_tensor((VecDstType*)(cute::raw_pointer_cast(mma_B.data())), cute::make_shape(cute::Int<K / (compress_size * vec_size)>{}, cute::Int<N>{}));
278+
279+
//auto src_ = *(cute::array<VecSrcType, K / (compress_size * vec_size, N)>*)(s_tensor.data());
280+
//auto dst_ = *(cute::array<VecDstType, K / (compress_size * vec_size, N)>*)(d_tensor.data());
281+
#pragma unroll
282+
for (int n = 0; n < N; n++) {
283+
float scale_value = fragment_scale(n);
284+
auto src = *(cute::array<VecSrcType, K / (compress_size * vec_size)>*)(s_tensor(_, n).data());
285+
auto& dst = *(cute::array<VecDstType, K / (compress_size * vec_size)>*)(d_tensor(_, n).data());
286+
//auto& src = *(cute::array<VecSrcType, K / (compress_size * vec_size)>*)(src_[n]);
287+
//auto& dst = *(cute::array<VecDstType, K / (compress_size * vec_size)>*)(dst_[n]);
305288
#pragma unroll
306-
for (int n = 0; n < N; n++) {
307-
float ts = tCrS_input(n);
308-
auto& src = *(cute::array<VecSrcType, K / (compress_size * vec_size)>*)(s_tensor(_, n).data());
309-
auto& dst = *(cute::array<VecDstType, K / (compress_size * vec_size)>*)(d_tensor(_, n).data());
310-
311-
#pragma unroll
312-
for (int k = 0; k < K / (compress_size * vec_size); k++) {
313-
VecDstType dst_val;
314-
315-
#pragma unroll
316-
for (int i = 0; i < vec_size; i++) {
317-
VecDstElemType dst_elem;
318-
319-
#pragma unroll
320-
for (int j = 0; j < compress_size; j++) {
321-
dst_elem[j] = static_cast<ElementMMA>(
322-
quant_map_[(src[k][i] >> (4 * ((j+1)%2 + (j/2)*2))) & 0xf] * ts
323-
);
324-
}
325-
dst_val[i] = dst_elem;
326-
}
327-
dst[k] = dst_val;
328-
}
289+
for (int k = 0; k < K / (compress_size * vec_size); k++) {
290+
VecDstType dst_val;
291+
#pragma unroll
292+
for (int i = 0; i < vec_size; i++) {
293+
VecDstElemType dst_elem;
294+
#pragma unroll
295+
for (int j = 0; j < compress_size; j++) {
296+
dst_elem[j] = static_cast<ElementMMA>(quant_map[(src[k][i] >> (4 * ((j+1)%2 + (j/2)*2))) & 0xf] * scale_value);
297+
}
298+
dst_val[i] = dst_elem;
299+
}
300+
dst[k] = dst_val;
329301
}
302+
}
330303
};
331304

332305
CUTLASS_PRAGMA_UNROLL
@@ -335,30 +308,38 @@ class gemm_4bit_cutlass_kernel {
335308
prefetch(tiled_prefetch_b, pBgB(_,_,_,prefetch_k));
336309
}
337310

338-
for (int k_tile = k_start_idx, k_s = 0; k_tile < k_tile_count + k_start_idx; k_tile++, prefetch_k++, k_s++) {
311+
for (int k_tile = k_start_idx, k_s = 0; k_tile < k_tile_count; k_tile++, k_s++) {
339312
copy(params.tiled_copy_b, tBgB(_,_,_,k_tile), frag_copy_B);
340-
341313
copy(params.tiled_copy_scale, tSgS(_, _, _, (k_start_idx + k_s) / k_reload_factor), frag_copy_Scale);
342-
343-
dequant(dequant_frag, mma_B, fragment_scale, quant_map);
344-
314+
//barrier_wait(3);
315+
dequant();
345316
copy(params.tiled_copy_a, tAgA(_,_,_,k_tile), frag_copy_A);
346-
347-
if(prefetch_k < k_tile_count) {
317+
318+
if (prefetch_k < k_tile_count) {
348319
prefetch(tiled_prefetch_a, pAgA(_,_,_,prefetch_k));
349320
prefetch(tiled_prefetch_b, pBgB(_,_,_,prefetch_k));
321+
prefetch_k++;
350322
}
351-
323+
352324
cute::gemm(tiled_mma, mma_A, mma_B, accumulators);
353325
barrier_wait(3);
354326
}
355327

356-
CUTLASS_PRAGMA_UNROLL
357-
for (int epi_n = 0; epi_n < FragsN; ++epi_n) {
358-
CUTLASS_PRAGMA_UNROLL
359-
for (int epi_m = 0; epi_m < FragsM; ++epi_m) {
360-
copy(params.tiled_store_d, accumulators(_, epi_m, epi_n), tCgD(_, epi_m, epi_n));
361-
}
328+
Tensor mD_mnl = cute::get_pvc_tensor(make_shape(params.m, params.n, params.l));
329+
Tensor g_wg_D = local_tile(mD_mnl, take<0,2>(WorkgroupTileShape{}), make_coord(m_coord,n_coord,l_coord));
330+
Tensor gD = local_tile(g_wg_D, take<0,2>(SubgroupTileShape{}), make_coord(
331+
get_sub_group_id() / ATOM_N,
332+
get_sub_group_id() % ATOM_N
333+
));
334+
335+
auto thread_xe_store_d = params.tiled_store_d.get_thread_slice(thread_idx);
336+
Tensor tCgD = thread_xe_store_d.partition_D(gD);
337+
338+
#pragma unroll
339+
for (int epi = 0; epi < FragsM * FragsN; ++epi) {
340+
int epi_m = epi / FragsN;
341+
int epi_n = epi % FragsN;
342+
copy(params.tiled_store_d, accumulators(_, epi_m, epi_n), tCgD(_, epi_m, epi_n));
362343
}
363344
}
364345
};
@@ -412,7 +393,7 @@ void gemm_4bit_cutlass(int m, int n, int k, int l, T *A, unsigned char *B,
412393

413394
StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(m, n, l));
414395
auto mD = make_tensor(make_gmem_ptr(out), make_layout(make_shape(m, n, l), stride_D));
415-
Copy_D tiled_store_d = {tiled_store_d.with(mD)};
396+
Copy_D tiled_store_d = {Copy_D{}.with(mD)};
416397
params.tiled_store_d = tiled_store_d;
417398

418399
params.hw_info = hw_info;

0 commit comments

Comments
 (0)