mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-05-04 05:31:24 +00:00
9fcc1ee9fd
* Support wave32/wave64 in CK_TILE - Part 1 * remove blocksize in kernel launch * fix build error * fix clang format * fix clang format 2 * fix clang format 3 * fix fmha build error * fix fmha build 2 * fix fmha build 3 * fix build error 4 * address review comment * update change log * replace KernelBlockSize with kBlockSize * fix CI fail * fix clang format * address review comment and rebase code. * fix universal test fail --------- Co-authored-by: Lin, Qun <Quentin.Lin+amdeng@amd.com> Co-authored-by: Thomas Ning <Thomas.Ning@amd.com>
62 lines
2.4 KiB
C++
62 lines
2.4 KiB
C++
// Copyright © Advanced Micro Devices, Inc., or its affiliates.
|
|
// SPDX-License-Identifier: MIT
|
|
|
|
#include <ck_tile/core.hpp>
|
|
#include "smoothquant.hpp"
|
|
#include <iostream>
|
|
|
|
#pragma once
|
|
|
|
using S = ck_tile::stream_config;
|
|
using A = smoothquant_args;
|
|
|
|
template <typename DataType_,
|
|
ck_tile::index_t Repeat_M_, // each thread repeat along M
|
|
ck_tile::index_t Repeat_N_, // each thread repeat along N
|
|
ck_tile::index_t ThreadPerBlock_M_, // num threads along M
|
|
ck_tile::index_t ThreadPerBlock_N_, // num threads along N
|
|
ck_tile::index_t Vector_N_, // vector size along N
|
|
bool kPadN_,
|
|
bool kTwoPass_>
|
|
using trait_ = smoothquant_traits_<DataType_,
|
|
Repeat_M_,
|
|
Repeat_N_,
|
|
ThreadPerBlock_M_,
|
|
ThreadPerBlock_N_,
|
|
Vector_N_,
|
|
kPadN_,
|
|
kTwoPass_>;
|
|
|
|
template <typename Traits_>
|
|
float smoothquant_(const S& s, A a)
|
|
{
|
|
using DataType = typename Traits_::DataType;
|
|
|
|
using PipelineProblem = ck_tile::SmoothquantPipelineProblem<
|
|
typename SmoothquantTypeConfig<DataType>::XDataType,
|
|
typename SmoothquantTypeConfig<DataType>::SmoothScaleDataType,
|
|
typename SmoothquantTypeConfig<DataType>::ComputeDataType,
|
|
typename SmoothquantTypeConfig<DataType>::YScaleDataType,
|
|
typename SmoothquantTypeConfig<DataType>::QYDataType,
|
|
typename Traits_::Shape,
|
|
Traits_::kPadN,
|
|
Traits_::kTwoPass>;
|
|
|
|
using OnePassPipeline = ck_tile::SmoothquantPipelineOnePass<PipelineProblem>;
|
|
using TwoPassPipeline = ck_tile::SmoothquantPipelineTwoPass<PipelineProblem>;
|
|
using Pipeline = std::conditional_t<Traits_::kTwoPass, TwoPassPipeline, OnePassPipeline>;
|
|
|
|
using Kernel = ck_tile::Smoothquant<Pipeline>;
|
|
|
|
const dim3 grids = Kernel::GridSize(a);
|
|
constexpr dim3 blocks = Kernel::BlockSize();
|
|
constexpr ck_tile::index_t kBlockPerCu = 1;
|
|
|
|
auto kargs = Kernel::MakeKargs(a);
|
|
if(s.log_level_ > 0)
|
|
std::cout << ", " << Kernel::GetName() << std::flush;
|
|
|
|
return ck_tile::launch_kernel(
|
|
s, ck_tile::make_kernel<kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
|
|
}
|