Skip to content

GitLab

Commit 3cc57101 authored by Jing Zhang's avatar Jing Zhang
Browse files

wrap desc into a struct

parent 698573a9
Hide whitespace changes
Inline Side-by-side
Showing with 125 additions and 226 deletions
composable_kernel/include/config.hpp
View file @ 3cc57101
...@@ -176,7 +176,6 @@ struct gemm_desc ...@@ -176,7 +176,6 @@ struct gemm_desc
ck::index_t M, N, K; ck::index_t M, N, K;
ck::index_t StrideA, StrideB, StrideC; ck::index_t StrideA, StrideB, StrideC;
ck::index_t OffsetA, OffsetB, OffsetC; ck::index_t OffsetA, OffsetB, OffsetC;
ck::index_t BlockStart, BlockSize;
}; };
} // namespace ck } // namespace ck
......
composable_kernel/include/tensor_operation/gridwise_grouped_gemm_xdlops_v2r3.hpp
View file @ 3cc57101
...@@ -15,14 +15,10 @@ namespace ck { ...@@ -15,14 +15,10 @@ namespace ck {
template <typename GridwiseGemm, template <typename GridwiseGemm,
typename FloatAB, typename FloatAB,
typename FloatC, typename FloatC,
typename AGridDesc_K0_M_K1,
typename BGridDesc_K0_N_K1,
typename CGridDesc_M0_N0_M1_N1_M2_M3_M4_N2,
typename GemmDesc, typename GemmDesc,
typename AElementwiseOperation, typename AElementwiseOperation,
typename BElementwiseOperation, typename BElementwiseOperation,
typename CElementwiseOperation, typename CElementwiseOperation,
typename Block2CTileMap,
bool HasMainK0BlockLoop, bool HasMainK0BlockLoop,
index_t MaxGroupCount> index_t MaxGroupCount>
__global__ void __global__ void
...@@ -33,49 +29,43 @@ __global__ void ...@@ -33,49 +29,43 @@ __global__ void
const FloatAB* __restrict__ p_a_grid, const FloatAB* __restrict__ p_a_grid,
const FloatAB* __restrict__ p_b_grid, const FloatAB* __restrict__ p_b_grid,
FloatC* __restrict__ p_c_grid, FloatC* __restrict__ p_c_grid,
const StaticallyIndexedArray<AGridDesc_K0_M_K1, MaxGroupCount> a_grid_desc_k0_m_k1, const StaticallyIndexedArray<GemmDesc, MaxGroupCount> gemm_desc_,
const StaticallyIndexedArray<BGridDesc_K0_N_K1, MaxGroupCount> b_grid_desc_k0_n_k1,
const StaticallyIndexedArray<CGridDesc_M0_N0_M1_N1_M2_M3_M4_N2, MaxGroupCount>
c_grid_desc_m0_n0_m1_n1_m2_m3_m4_n2,
const StaticallyIndexedArray<GemmDesc, MaxGroupCount> gemm_shapes,
const index_t group_count, const index_t group_count,
const AElementwiseOperation a_element_op, const AElementwiseOperation a_element_op,
const BElementwiseOperation b_element_op, const BElementwiseOperation b_element_op,
const CElementwiseOperation c_element_op, const CElementwiseOperation c_element_op)
const StaticallyIndexedArray<Block2CTileMap, MaxGroupCount> block_2_ctile_map)
{ {
__shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()]; __shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()];
const index_t block_id = get_block_1d_id(); const index_t block_id = get_block_1d_id();
static_for<0, MaxGroupCount, 1>{}([&](auto i) { static_for<0, MaxGroupCount, 1>{}([&](auto i) {
if(block_id >= gemm_shapes[i].BlockStart && if(block_id >= gemm_desc_[i].BlockStart && block_id < gemm_desc_[i].BlockEnd)
block_id < (gemm_shapes[i].BlockStart + gemm_shapes[i].BlockSize))
{ {
const index_t group_id = i; const index_t group_id = i;
const index_t block_id_grp = block_id - gemm_shapes[i].BlockStart; const index_t block_id_grp = block_id - gemm_desc_[i].BlockStart;
const index_t a_offset_grp = gemm_shapes[i].OffsetA; const index_t a_offset_grp = gemm_desc_[i].OffsetA;
const index_t b_offset_grp = gemm_shapes[i].OffsetB; const index_t b_offset_grp = gemm_desc_[i].OffsetB;
const index_t c_offset_grp = gemm_shapes[i].OffsetC; const index_t c_offset_grp = gemm_desc_[i].OffsetC;
GridwiseGemm::template Run<HasMainK0BlockLoop>(p_a_grid + a_offset_grp, GridwiseGemm::template Run<HasMainK0BlockLoop>(
p_b_grid + b_offset_grp, p_a_grid + a_offset_grp,
p_c_grid + c_offset_grp, p_b_grid + b_offset_grp,
p_shared, p_c_grid + c_offset_grp,
a_grid_desc_k0_m_k1[i], p_shared,
b_grid_desc_k0_n_k1[i], gemm_desc_[i].a_grid_desc_k0_m_k1_,
c_grid_desc_m0_n0_m1_n1_m2_m3_m4_n2[i], gemm_desc_[i].b_grid_desc_k0_n_k1_,
a_element_op, gemm_desc_[i].c_grid_desc_m0_n0_m1_n1_m2_m3_m4_n2_,
b_element_op, a_element_op,
c_element_op, b_element_op,
block_2_ctile_map[i], c_element_op,
block_id_grp); gemm_desc_[i].block_2_ctile_map_,
block_id_grp);
return;
} }
}); });
} }
#if 0
template <typename GridwiseGemm, template <typename GridwiseGemm,
typename FloatAB, typename FloatAB,
typename FloatC, typename FloatC,
...@@ -159,6 +149,7 @@ __global__ void ...@@ -159,6 +149,7 @@ __global__ void
block_2_ctile_map_[group_id], block_2_ctile_map_[group_id],
block_id_grp); block_id_grp);
} }
#endif
template <index_t BlockSize, template <index_t BlockSize,
typename FloatAB, typename FloatAB,
......
device_operation/include/device_grouped_gemm_xdl.hpp
View file @ 3cc57101
...@@ -223,6 +223,20 @@ struct DeviceGroupedGemmXdl ...@@ -223,6 +223,20 @@ struct DeviceGroupedGemmXdl
CThreadTransferDstScalarPerVector, CThreadTransferDstScalarPerVector,
NumPrefetch>; NumPrefetch>;
struct GemmDesc
{
AGridDesc_K0_M_K1 a_grid_desc_k0_m_k1_;
BGridDesc_K0_N_K1 b_grid_desc_k0_n_k1_;
CGridDesc_M_N c_grid_desc_m_n_;
typename GridwiseGemm::CGridDesc_M0_N0_M1_N1_M2_M3_M4_N2
c_grid_desc_m0_n0_m1_n1_m2_m3_m4_n2_;
typename GridwiseGemm::DefaultBlock2CTileMap block_2_ctile_map_;
ck::index_t OffsetA, OffsetB, OffsetC;
ck::index_t BlockStart, BlockEnd;
};
// Argument // Argument
struct Argument : public BaseArgument struct Argument : public BaseArgument
{ {
...@@ -246,6 +260,8 @@ struct DeviceGroupedGemmXdl ...@@ -246,6 +260,8 @@ struct DeviceGroupedGemmXdl
c_element_op_{c_element_op} c_element_op_{c_element_op}
{ {
grid_size = 0;
static_for<0, MaxGroupCount, 1>{}([&](auto i) { static_for<0, MaxGroupCount, 1>{}([&](auto i) {
if(i < gemm_shapes_.size()) if(i < gemm_shapes_.size())
{ {
...@@ -257,27 +273,44 @@ struct DeviceGroupedGemmXdl ...@@ -257,27 +273,44 @@ struct DeviceGroupedGemmXdl
const index_t StrideB = gemm_shapes_[i].StrideB; const index_t StrideB = gemm_shapes_[i].StrideB;
const index_t StrideC = gemm_shapes_[i].StrideC; const index_t StrideC = gemm_shapes_[i].StrideC;
a_grid_desc_k0_m_k1_(i) = gemm_desc_(i).a_grid_desc_k0_m_k1_ =
DeviceGroupedGemmXdl::MakeAGridDescriptor_K0_M_K1(M, K, StrideA); DeviceGroupedGemmXdl::MakeAGridDescriptor_K0_M_K1(M, K, StrideA);
b_grid_desc_k0_n_k1_(i) = gemm_desc_(i).b_grid_desc_k0_n_k1_ =
DeviceGroupedGemmXdl::MakeBGridDescriptor_K0_N_K1(K, N, StrideB); DeviceGroupedGemmXdl::MakeBGridDescriptor_K0_N_K1(K, N, StrideB);
c_grid_desc_m_n_(i) = gemm_desc_(i).c_grid_desc_m_n_ =
DeviceGroupedGemmXdl::MakeCGridDescriptor_M_N(M, N, StrideC); DeviceGroupedGemmXdl::MakeCGridDescriptor_M_N(M, N, StrideC);
if(GridwiseGemm::CheckValidity(a_grid_desc_k0_m_k1_[i], const index_t grid_size_grp =
b_grid_desc_k0_n_k1_[i], GridwiseGemm::CalculateGridSize(gemm_desc_[i].c_grid_desc_m_n_);
c_grid_desc_m_n_[i],
gemm_desc_(i).BlockStart = grid_size;
gemm_desc_(i).BlockEnd = grid_size + grid_size_grp;
grid_size += grid_size_grp;
gemm_desc_(i).OffsetA = gemm_shapes_[i].OffsetA;
gemm_desc_(i).OffsetB = gemm_shapes_[i].OffsetB;
gemm_desc_(i).OffsetC = gemm_shapes_[i].OffsetC;
if(GridwiseGemm::CheckValidity(gemm_desc_[i].a_grid_desc_k0_m_k1_,
gemm_desc_[i].b_grid_desc_k0_n_k1_,
gemm_desc_[i].c_grid_desc_m_n_,
M01_, M01_,
N01_)) N01_))
{ {
c_grid_desc_m0_n0_m1_n1_m2_m3_m4_n2_(i) = gemm_desc_(i).c_grid_desc_m0_n0_m1_n1_m2_m3_m4_n2_ =
GridwiseGemm::MakeCGridDescriptor_M0_N0_M1_N1_M2_M3_M4_N2( GridwiseGemm::MakeCGridDescriptor_M0_N0_M1_N1_M2_M3_M4_N2(
c_grid_desc_m_n_[i]); gemm_desc_[i].c_grid_desc_m_n_);
block_2_ctile_map_(i) = gemm_desc_(i).block_2_ctile_map_ = GridwiseGemm::MakeDefaultBlock2CTileMap(
GridwiseGemm::MakeDefaultBlock2CTileMap(c_grid_desc_m_n_[i], M01, N01); gemm_desc_[i].c_grid_desc_m_n_, M01, N01);
} }
} }
else
{
gemm_desc_(i).BlockStart = -1;
gemm_desc_(i).BlockEnd = -1;
}
}); });
} }
...@@ -285,20 +318,16 @@ struct DeviceGroupedGemmXdl ...@@ -285,20 +318,16 @@ struct DeviceGroupedGemmXdl
const ADataType* p_a_grid_; const ADataType* p_a_grid_;
const BDataType* p_b_grid_; const BDataType* p_b_grid_;
CDataType* p_c_grid_; CDataType* p_c_grid_;
StaticallyIndexedArray<AGridDesc_K0_M_K1, MaxGroupCount> a_grid_desc_k0_m_k1_;
StaticallyIndexedArray<BGridDesc_K0_N_K1, MaxGroupCount> b_grid_desc_k0_n_k1_;
StaticallyIndexedArray<CGridDesc_M_N, MaxGroupCount> c_grid_desc_m_n_;
StaticallyIndexedArray<typename GridwiseGemm::CGridDesc_M0_N0_M1_N1_M2_M3_M4_N2,
MaxGroupCount>
c_grid_desc_m0_n0_m1_n1_m2_m3_m4_n2_;
StaticallyIndexedArray<typename GridwiseGemm::DefaultBlock2CTileMap, MaxGroupCount>
block_2_ctile_map_;
std::vector<gemm_desc> gemm_shapes_; std::vector<gemm_desc> gemm_shapes_;
index_t M01_; index_t M01_;
index_t N01_; index_t N01_;
AElementwiseOperation a_element_op_; AElementwiseOperation a_element_op_;
BElementwiseOperation b_element_op_; BElementwiseOperation b_element_op_;
CElementwiseOperation c_element_op_; CElementwiseOperation c_element_op_;
StaticallyIndexedArray<GemmDesc, MaxGroupCount> gemm_desc_;
index_t grid_size;
}; };
// Invoker // Invoker
...@@ -308,134 +337,102 @@ struct DeviceGroupedGemmXdl ...@@ -308,134 +337,102 @@ struct DeviceGroupedGemmXdl
float Run(const Argument& arg, int nrepeat = 1) float Run(const Argument& arg, int nrepeat = 1)
{ {
StaticallyIndexedArray<gemm_desc, MaxGroupCount> gemm_shapes;
index_t grid_size = 0;
static_for<0, MaxGroupCount, 1>{}([&](auto i) { static_for<0, MaxGroupCount, 1>{}([&](auto i) {
if(i < arg.gemm_shapes_.size()) if(i < arg.gemm_shapes_.size())
{ {
std::cout << "arg.a_grid_desc_k0_m_k1_{" std::cout << "arg.a_grid_desc_k0_m_k1_{"
<< arg.a_grid_desc_k0_m_k1_[i].GetLength(I0) << ", " << arg.gemm_desc_[i].a_grid_desc_k0_m_k1_.GetLength(I0) << ", "
<< arg.a_grid_desc_k0_m_k1_[i].GetLength(I1) << ", " << arg.gemm_desc_[i].a_grid_desc_k0_m_k1_.GetLength(I1) << ", "
<< arg.a_grid_desc_k0_m_k1_[i].GetLength(I2) << "}" << std::endl; << arg.gemm_desc_[i].a_grid_desc_k0_m_k1_.GetLength(I2) << "}"
<< std::endl;
std::cout << "arg.b_grid_desc_k0_n_k1_{" std::cout << "arg.b_grid_desc_k0_n_k1_{"
<< arg.b_grid_desc_k0_n_k1_[i].GetLength(I0) << ", " << arg.gemm_desc_[i].b_grid_desc_k0_n_k1_.GetLength(I0) << ", "
<< arg.b_grid_desc_k0_n_k1_[i].GetLength(I1) << ", " << arg.gemm_desc_[i].b_grid_desc_k0_n_k1_.GetLength(I1) << ", "
<< arg.b_grid_desc_k0_n_k1_[i].GetLength(I2) << "}" << std::endl; << arg.gemm_desc_[i].b_grid_desc_k0_n_k1_.GetLength(I2) << "}"
<< std::endl;
std::cout << "arg.c_grid_desc_m_n_{ " << arg.c_grid_desc_m_n_[i].GetLength(I0)
<< ", " << arg.c_grid_desc_m_n_[i].GetLength(I1) << "}" << std::endl; std::cout << "arg.c_grid_desc_m_n_{ "
<< arg.gemm_desc_[i].c_grid_desc_m_n_.GetLength(I0) << ", "
if(!GridwiseGemm::CheckValidity(arg.a_grid_desc_k0_m_k1_[i], << arg.gemm_desc_[i].c_grid_desc_m_n_.GetLength(I1) << "}"
arg.b_grid_desc_k0_n_k1_[i], << std::endl;
arg.c_grid_desc_m_n_[i],
std::cout << "Block: " << arg.gemm_desc_[i].BlockStart << ", "
<< arg.gemm_desc_[i].BlockEnd << std::endl;
if(!GridwiseGemm::CheckValidity(arg.gemm_desc_[i].a_grid_desc_k0_m_k1_,
arg.gemm_desc_[i].b_grid_desc_k0_n_k1_,
arg.gemm_desc_[i].c_grid_desc_m_n_,
arg.M01_, arg.M01_,
arg.N01_)) arg.N01_))
{ {
throw std::runtime_error( throw std::runtime_error(
"wrong! GridwiseGemm_k0mk1_k0nk1_mn_xdlops_v2r3 has invalid setting"); "wrong! GridwiseGemm_k0mk1_k0nk1_mn_xdlops_v2r3 has invalid setting");
} }
const index_t grid_size_grp =
GridwiseGemm::CalculateGridSize(arg.c_grid_desc_m_n_[i]);
gemm_shapes(i) = arg.gemm_shapes_[i];
gemm_shapes(i).BlockStart = grid_size;
gemm_shapes(i).BlockSize = grid_size_grp;
grid_size += grid_size_grp;
std::cout << "group_id " << i << " BlockStart " << gemm_shapes(i).BlockStart
<< " BlockSize " << gemm_shapes(i).BlockSize << std::endl;
}
else
{
gemm_shapes(i).BlockStart = -1;
gemm_shapes(i).BlockSize = -1;
} }
}); });
const auto K0 = arg.a_grid_desc_k0_m_k1_[Number<0>{}].GetLength(I0); const auto K0 = arg.gemm_desc_[I0].a_grid_desc_k0_m_k1_.GetLength(I0);
const bool has_main_k0_block_loop = GridwiseGemm::CalculateHasMainK0BlockLoop(K0); const bool has_main_k0_block_loop = GridwiseGemm::CalculateHasMainK0BlockLoop(K0);
float ave_time = 0; float ave_time = 0;
#if 1
if(has_main_k0_block_loop) if(has_main_k0_block_loop)
{ {
const auto kernel = kernel_grouped_gemm_xdlops_v2r3< const auto kernel =
GridwiseGemm, kernel_grouped_gemm_xdlops_v2r3<GridwiseGemm,
ADataType, // TODO: distiguish A/B datatype ADataType, // TODO: distiguish A/B datatype
CDataType, CDataType,
remove_reference_t<DeviceGroupedGemmXdl::AGridDesc_K0_M_K1>, remove_reference_t<GemmDesc>,
remove_reference_t<DeviceGroupedGemmXdl::BGridDesc_K0_N_K1>, AElementwiseOperation,
remove_reference_t<typename GridwiseGemm::CGridDesc_M0_N0_M1_N1_M2_M3_M4_N2>, BElementwiseOperation,
remove_reference_t<gemm_desc>, CElementwiseOperation,
AElementwiseOperation, true,
BElementwiseOperation, MaxGroupCount>;
CElementwiseOperation,
remove_reference_t<typename GridwiseGemm::DefaultBlock2CTileMap>,
true,
MaxGroupCount>;
ave_time = launch_and_time_kernel(kernel, ave_time = launch_and_time_kernel(kernel,
nrepeat, nrepeat,
dim3(grid_size), dim3(arg.grid_size),
dim3(BlockSize), dim3(BlockSize),
0, 0,
arg.p_a_grid_, arg.p_a_grid_,
arg.p_b_grid_, arg.p_b_grid_,
arg.p_c_grid_, arg.p_c_grid_,
arg.a_grid_desc_k0_m_k1_, arg.gemm_desc_,
arg.b_grid_desc_k0_n_k1_,
arg.c_grid_desc_m0_n0_m1_n1_m2_m3_m4_n2_,
gemm_shapes,
arg.gemm_shapes_.size(), arg.gemm_shapes_.size(),
arg.a_element_op_, arg.a_element_op_,
arg.b_element_op_, arg.b_element_op_,
arg.c_element_op_, arg.c_element_op_);
arg.block_2_ctile_map_);
} }
else else
{ {
const auto kernel = kernel_grouped_gemm_xdlops_v2r3< const auto kernel =
GridwiseGemm, kernel_grouped_gemm_xdlops_v2r3<GridwiseGemm,
ADataType, // TODO: distiguish A/B datatype ADataType, // TODO: distiguish A/B datatype
CDataType, CDataType,
remove_reference_t<DeviceGroupedGemmXdl::AGridDesc_K0_M_K1>, remove_reference_t<GemmDesc>,
remove_reference_t<DeviceGroupedGemmXdl::BGridDesc_K0_N_K1>, AElementwiseOperation,
remove_reference_t<typename GridwiseGemm::CGridDesc_M0_N0_M1_N1_M2_M3_M4_N2>, BElementwiseOperation,
remove_reference_t<gemm_desc>, CElementwiseOperation,
AElementwiseOperation, false,
BElementwiseOperation, MaxGroupCount>;
CElementwiseOperation,
remove_reference_t<typename GridwiseGemm::DefaultBlock2CTileMap>,
false,
MaxGroupCount>;
ave_time = launch_and_time_kernel(kernel, ave_time = launch_and_time_kernel(kernel,
nrepeat, nrepeat,
dim3(grid_size), dim3(arg.grid_size),
dim3(BlockSize), dim3(BlockSize),
0, 0,
arg.p_a_grid_, arg.p_a_grid_,
arg.p_b_grid_, arg.p_b_grid_,
arg.p_c_grid_, arg.p_c_grid_,
arg.a_grid_desc_k0_m_k1_, arg.gemm_desc_,
arg.b_grid_desc_k0_n_k1_,
arg.c_grid_desc_m0_n0_m1_n1_m2_m3_m4_n2_,
gemm_shapes,
arg.gemm_shapes_.size(), arg.gemm_shapes_.size(),
arg.a_element_op_, arg.a_element_op_,
arg.b_element_op_, arg.b_element_op_,
arg.c_element_op_, arg.c_element_op_);
arg.block_2_ctile_map_);
} }
#endif
return ave_time; return ave_time;
} }
...@@ -454,9 +451,9 @@ struct DeviceGroupedGemmXdl ...@@ -454,9 +451,9 @@ struct DeviceGroupedGemmXdl
static bool IsSupportedArgument(const Argument& arg) static bool IsSupportedArgument(const Argument& arg)
{ {
return GridwiseGemm::CheckValidity(arg.a_grid_desc_k0_m_k1_[Number<0>{}], return GridwiseGemm::CheckValidity(arg.gemm_desc_[Number<0>{}].a_grid_desc_k0_m_k1_,
arg.b_grid_desc_k0_n_k1_[Number<0>{}], arg.gemm_desc_[Number<0>{}].b_grid_desc_k0_n_k1_,
arg.c_grid_desc_m_n_[Number<0>{}], arg.gemm_desc_[Number<0>{}].c_grid_desc_m_n_,
arg.M01_, arg.M01_,
arg.N01_); arg.N01_);
} }
......
example/12_grouped_gemm_xdl/grouped_gemm_xdl.cpp
View file @ 3cc57101
...@@ -76,7 +76,7 @@ int main(int argc, char* argv[]) ...@@ -76,7 +76,7 @@ int main(int argc, char* argv[])
exit(0); exit(0);
} }
int group_count = 3; int group_count = 4;
// GEMM shape // GEMM shape
std::vector<ck::gemm_desc> gemm_shapes; std::vector<ck::gemm_desc> gemm_shapes;
...@@ -85,11 +85,11 @@ int main(int argc, char* argv[]) ...@@ -85,11 +85,11 @@ int main(int argc, char* argv[])
for(int i = 0; i < group_count; i++) for(int i = 0; i < group_count; i++)
{ {
int M = 2048 + 256 * i; int M = 3840;
int N = 2048 + 128 * i; int N = 1024;
int K = 256 + 128 * i; int K = 4096;
gemm_shapes.push_back({M, N, K, K, K, N, A_size, B_size, C_size, -1, -1}); gemm_shapes.push_back({M, N, K, K, K, N, A_size, B_size, C_size});
A_size += M * K; A_size += M * K;
B_size += N * K; B_size += N * K;
...@@ -236,94 +236,6 @@ int main(int argc, char* argv[]) ...@@ -236,94 +236,6 @@ int main(int argc, char* argv[])
check_error(c_host_tensors[i], c_device_tensors[i]); check_error(c_host_tensors[i], c_device_tensors[i]);
} }
} }
#if 0
Tensor<ADataType> a_m_k(f_host_tensor_descriptor(M, K, StrideA, ALayout{}));
Tensor<BDataType> b_k_n(f_host_tensor_descriptor(K, N, StrideB, BLayout{}));
Tensor<CDataType> c_m_n_host_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{}));
Tensor<CDataType> c_m_n_device_result(f_host_tensor_descriptor(M, N, StrideC, CLayout{}));
std::cout << "a_m_k: " << a_m_k.mDesc << std::endl;
std::cout << "b_k_n: " << b_k_n.mDesc << std::endl;
std::cout << "c_m_n: " << c_m_n_host_result.mDesc << std::endl;
switch(init_method)
{
case 0: break;
case 1:
a_m_k.GenerateTensorValue(GeneratorTensor_2<ADataType>{-5, 5});
b_k_n.GenerateTensorValue(GeneratorTensor_2<BDataType>{-5, 5});
break;
case 2:
a_m_k.GenerateTensorValue(GeneratorTensor_3<ADataType>{0.0, 1.0});
b_k_n.GenerateTensorValue(GeneratorTensor_3<BDataType>{-0.5, 0.5});
break;
default:
a_m_k.GenerateTensorValue(GeneratorTensor_Sequential<0>{});
b_k_n.GenerateTensorValue(GeneratorTensor_Sequential<1>{});
}
DeviceMem a_m_k_device_buf(sizeof(ADataType) * a_m_k.mDesc.GetElementSpace());
DeviceMem b_k_n_device_buf(sizeof(BDataType) * b_k_n.mDesc.GetElementSpace());
DeviceMem c_m_n_device_buf(sizeof(CDataType) * c_m_n_device_result.mDesc.GetElementSpace());
a_m_k_device_buf.ToDevice(a_m_k.mData.data());
b_k_n_device_buf.ToDevice(b_k_n.mData.data());
auto a_element_op = AElementOp{};
auto b_element_op = BElementOp{};
auto c_element_op = CElementOp{};
// do GEMM
auto gemm = DeviceGemmInstance{};
auto invoker = gemm.MakeInvoker();
auto argument = gemm.MakeArgument(static_cast<ADataType*>(a_m_k_device_buf.GetDeviceBuffer()),
static_cast<BDataType*>(b_k_n_device_buf.GetDeviceBuffer()),
static_cast<CDataType*>(c_m_n_device_buf.GetDeviceBuffer()),
M,
N,
K,
StrideA,
StrideB,
StrideC,
a_element_op,
b_element_op,
c_element_op);
if(!gemm.IsSupportedArgument(argument))
{
throw std::runtime_error(
"wrong! device_gemm with the specified compilation parameters does "
"not support this GEMM problem");
}
float ave_time = invoker.Run(argument, nrepeat);
std::size_t flop = std::size_t(2) * M * N * K;
std::size_t num_btype =
sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + sizeof(CDataType) * M * N;
float tflops = static_cast<float>(flop) / 1.E9 / ave_time;
float gb_per_sec = num_btype / 1.E6 / ave_time;
std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s, "
<< gemm.GetTypeString() << std::endl;
c_m_n_device_buf.FromDevice(c_m_n_device_result.mData.data());
if(do_verification)
{
auto ref_gemm = ReferenceGemmInstance{};
auto ref_invoker = ref_gemm.MakeInvoker();
auto ref_argument = ref_gemm.MakeArgument(
a_m_k, b_k_n, c_m_n_host_result, a_element_op, b_element_op, c_element_op);
ref_invoker.Run(ref_argument);
check_error(c_m_n_host_result, c_m_n_device_result);
}
#endif
return 0; return 0;
} }
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment