ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-23 22:34:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add KBatch support for gemm_ab_scale (#2740)

Browse Source 
* Add KBatch support for gemm_ab_scale

* Revert kernel parameters change

* Remove printing

* fix formatting

* fix check

* Use {} in if

---------

Co-authored-by: Adam Osewski <19374865+aosewski@users.noreply.github.com>

[ROCm/composable_kernel commit: 9d4bfe3932]
This commit is contained in:
Sami Remes
2025-10-09 07:33:16 +01:00
committed by GitHub
parent e9ade69185
commit e7ef841a68
5 changed files with 34 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
profiler/src/profile_gemm_ab_scale.cpp
View File

@@ -40,7 +40,7 @@ enum struct ScaleBlockTile
int profile_gemm_ab_scale(int argc, char* argv[])
{
if(argc != 15 && argc != 18)
if(argc != 15 && argc != 16 && argc != 19)
{
printf("arg1: tensor operation (" OP_NAME ": " OP_DESC ")\n");
printf("arg2: data type (0: fp32; 1: fp16; 2: bf16; 3: int8; 4: f8@f16; 5: f16@f8; 6: "
@@ -57,10 +57,11 @@ int profile_gemm_ab_scale(int argc, char* argv[])
printf("arg7: print tensor value (0: no; 1: yes)\n");
printf("arg8: time kernel (0=no, 1=yes)\n");
printf("arg9 to 14: M, N, K, StrideA, StrideB, StrideE\n");
printf("arg15: KBatch (default: 1)\n");
printf("optional:\n");
printf("arg15: number of warm-up cycles (default 1)\n");
printf("arg16: number of iterations (default 10)\n");
printf("arg17: memory for rotating buffer (default 0, size in MB)\n");
printf("arg16: number of warm-up cycles (default 1)\n");
printf("arg17: number of iterations (default 10)\n");
printf("arg18: memory for rotating buffer (default 0, size in MB)\n");
exit(1);
}
@@ -79,15 +80,16 @@ int profile_gemm_ab_scale(int argc, char* argv[])
const int StrideA = std::stoi(argv[12]);
const int StrideB = std::stoi(argv[13]);
const int StrideE = std::stoi(argv[14]);
const int KBatch = (argc > 15) ? std::stoi(argv[15]) : 1;
int n_warmup = 1;
int n_iter = 10;
uint64_t rotating = 0;
if(argc == 18)
if(argc == 19)
{
n_warmup = std::stoi(argv[15]);
n_iter = std::stoi(argv[16]);
rotating = std::stoull(argv[17]) * 1024 * 1024;
n_warmup = std::stoi(argv[16]);
n_iter = std::stoi(argv[17]);
rotating = std::stoull(argv[18]) * 1024 * 1024;
}
using F32 = float;
@@ -149,6 +151,7 @@ int profile_gemm_ab_scale(int argc, char* argv[])
(StrideA < 0) ? DefaultStrideA : StrideA,
(StrideB < 0) ? DefaultStrideB : StrideB,
(StrideE < 0) ? DefaultStrideE : StrideE,
KBatch,
n_warmup,
n_iter,
rotating);