Use macro to decide mfma cluster size

a9255313 · Po-Yen, Chen · 3fea2678 · a9255313
Commit a9255313 authored Dec 09, 2022 by Po-Yen, Chen
Show whitespace changes
Inline Side-by-side

Showing with 20 additions and 22 deletions

include/ck/tensor_operation/gpu/grid/gridwise_gemm_pipeline_v2.hpp ...k/tensor_operation/gpu/grid/gridwise_gemm_pipeline_v2.hpp +20 -22

No files found.
--- a/include/ck/tensor_operation/gpu/grid/gridwise_gemm_pipeline_v2.hpp
+++ b/include/ck/tensor_operation/gpu/grid/gridwise_gemm_pipeline_v2.hpp
@@ -100,36 +100,34 @@ struct GridwiseGemmPipeline_v2
                // global read i + 2
                b_blockwise_copy.RunRead(b_grid_desc, b_grid_buf);

-#if defined(ENABLE_PIPELINE_V2_OPT)
-                __builtin_amdgcn_sched_group_barrier(0x020, 2, 0); // VMEM read
-                __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA
-
-                __builtin_amdgcn_sched_group_barrier(0x020, 1, 0); // VMEM read
-                __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA
+#if !defined(NUM_MFMA_PER_CLUSTER)
+#define NUM_MFMA_PER_CLUSTER 1
+#endif

+#if defined(ENABLE_PIPELINE_V2_OPT)
                __builtin_amdgcn_sched_group_barrier(0x020, 2, 0);                    // VMEM read
-                __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA
+                __builtin_amdgcn_sched_group_barrier(0x008, NUM_MFMA_PER_CLUSTER, 0); // MFMA

                __builtin_amdgcn_sched_group_barrier(0x020, 1, 0);                    // VMEM read
-                __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA
+                __builtin_amdgcn_sched_group_barrier(0x008, NUM_MFMA_PER_CLUSTER, 0); // MFMA

                __builtin_amdgcn_sched_group_barrier(0x020, 2, 0);                    // VMEM read
-                __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA
+                __builtin_amdgcn_sched_group_barrier(0x008, NUM_MFMA_PER_CLUSTER, 0); // MFMA

                __builtin_amdgcn_sched_group_barrier(0x020, 1, 0);                    // VMEM read
-                __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA
+                __builtin_amdgcn_sched_group_barrier(0x008, NUM_MFMA_PER_CLUSTER, 0); // MFMA

                __builtin_amdgcn_sched_group_barrier(0x020, 2, 0);                    // VMEM read
-                __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA
+                __builtin_amdgcn_sched_group_barrier(0x008, NUM_MFMA_PER_CLUSTER, 0); // MFMA

                __builtin_amdgcn_sched_group_barrier(0x020, 1, 0);                    // VMEM read
-                __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA
+                __builtin_amdgcn_sched_group_barrier(0x008, NUM_MFMA_PER_CLUSTER, 0); // MFMA

                __builtin_amdgcn_sched_group_barrier(0x020, 2, 0);                    // VMEM read
-                __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA
+                __builtin_amdgcn_sched_group_barrier(0x008, NUM_MFMA_PER_CLUSTER, 0); // MFMA

                __builtin_amdgcn_sched_group_barrier(0x020, 1, 0);                    // VMEM read
-                __builtin_amdgcn_sched_group_barrier(0x008, 1, 0); // MFMA
+                __builtin_amdgcn_sched_group_barrier(0x008, NUM_MFMA_PER_CLUSTER, 0); // MFMA
 #endif

                ++i;