mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-05-19 12:30:16 +00:00
87522d2ac9
* adding batched_gemm_and_reduction
* batched_gemm_reduce works with bactch_count=1
* fix a bug in grid_size; batched_gemm_reduce works for batch_count > 1
* adding profiler for batched_gemm_fp16
* fixed a bug in declaration of d1 and d0; both example and profiler work
* clang-format
* cleanup
* batched_gemm_reduce: add test
* minor change
* fixed some typo in function names
[ROCm/composable_kernel commit: 34c661e71c]
65 lines
2.8 KiB
C++
65 lines
2.8 KiB
C++
#include <iostream>
|
|
|
|
#include "profile_batched_gemm_reduce_impl.hpp"
|
|
|
|
int main()
|
|
{
|
|
using Row = ck::tensor_layout::gemm::RowMajor;
|
|
using Col = ck::tensor_layout::gemm::ColumnMajor;
|
|
|
|
int M = 512;
|
|
int N = 256;
|
|
int K = 128;
|
|
|
|
int BatchCount = 3;
|
|
|
|
bool pass = true;
|
|
|
|
pass = pass && ck::profiler::profile_batched_gemm_reduce_impl<ck::half_t,
|
|
ck::half_t,
|
|
ck::half_t,
|
|
float,
|
|
Row,
|
|
Row,
|
|
Row>(
|
|
true, 1, false, 1, M, N, K, K, N, N, BatchCount);
|
|
|
|
pass = pass && ck::profiler::profile_batched_gemm_reduce_impl<ck::half_t,
|
|
ck::half_t,
|
|
ck::half_t,
|
|
float,
|
|
Row,
|
|
Col,
|
|
Row>(
|
|
true, 1, false, 1, M, N, K, K, K, N, BatchCount);
|
|
|
|
pass = pass && ck::profiler::profile_batched_gemm_reduce_impl<ck::half_t,
|
|
ck::half_t,
|
|
ck::half_t,
|
|
float,
|
|
Col,
|
|
Row,
|
|
Row>(
|
|
true, 1, false, 1, M, N, K, M, N, N, BatchCount);
|
|
|
|
pass = pass && ck::profiler::profile_batched_gemm_reduce_impl<ck::half_t,
|
|
ck::half_t,
|
|
ck::half_t,
|
|
float,
|
|
Col,
|
|
Col,
|
|
Row>(
|
|
true, 1, false, 1, M, N, K, M, K, N, BatchCount);
|
|
|
|
if(pass)
|
|
{
|
|
std::cout << "test BatchedGEMM+Reduce fp16: Pass" << std::endl;
|
|
return 0;
|
|
}
|
|
else
|
|
{
|
|
std::cout << "test BatchedGEMM+Reduce fp16: Fail" << std::endl;
|
|
return -1;
|
|
}
|
|
}
|