update aten bmm CK heuristic (pytorch#143294)

Summary: updates heuristic to use new instances based on ck profiling of LLM shapes

Differential Revision: D67280269

Pull Request resolved: pytorch#143294
Approved by: https://github.com/mxz297, https://github.com/xw285cornell

Author: bradleyhd

aten/src/ATen/native/hip/ck_bgemm_bfloat16.hip

@@ -32,38 +32,28 @@ static const std::unordered_map<
 
 // This is the heursitic to choose a kernel based on inputs
 BGEMMKernel_BFloat16 dispatch_bfloat16_bgemm(CUDABLAS_BGEMM_ARGTYPES(at::BFloat16)) {
-  // First check if there's a specific kernel for this shape.
+  // Optional/future use: directly lookup shape tuples to map to instances
+  /*
   auto it = lookup_dispatch.find({m, n, k});
   if (it != lookup_dispatch.end()) {
     return it->second;
   }
+  */
+
+  // B is A and A is B, so m<-->n
+  // std::cout << "dispatch_bfloat16_bgemm: m=" << m << " n=" << n << " k=" << k << " num_batches=" << num_batches << " transa=" << transa << " transb=" << transb << std::endl;
 
-  // Nout found, use heuristics.
-  // TN layout, so n, m, k
-  if (k == 8192) {
-    if (n <= 4) {
-      return bgemm_kernel_bf16bf16bf16_64_16x16x64_16x16_1x1_8x8x1_8x8x1_1x16x1x4_4_Intrawave_v1;
-    }
-    else if (n <= 32) {
-      return bgemm_kernel_bf16bf16bf16_64_16x16x64_16x16_1x1_8x8x1_8x8x1_1x16x1x4_4_Intrawave_v2;
-    }
-    else if (n <= 512) {
-      return bgemm_kernel_bf16bf16bf16_256_128x128x64_32x32_2x2_8x32x1_8x32x1_1x16x1x16_4_Intrawave_v3;
-    }
-    else {
-      return bgemm_kernel_bf16bf16bf16_256_256x256x32_32x32_4x4_8x32x1_8x32x1_1x16x1x16_4_Intrawave_v4;
-    }
+  if (m <= 5120) {
+    if (n <= 4) return bgemm_kernel_bf16bf16bf16_64_16x16x64_16x16_1x1_8x8x1_8x8x1_1x16x1x4_4_Intrawave_v1;
+    else if (n <= 32) return bgemm_kernel_bf16bf16bf16_128_16x64x64_16x16_1x2_8x16x1_8x16x1_1x16x1x8_4_Intrawave_v2;
+    else if (n <= 128) return bgemm_kernel_bf16bf16bf16_256_128x128x64_32x32_2x2_8x32x1_8x32x1_1x16x1x16_4_Intrawave_v3; // <512, <1024, <2048 missing
+    else if (n <= 4096) return bgemm_kernel_bf16bf16bf16_256_224x256x64_16x16_7x8_8x32x1_8x32x1_1x16x1x16_4_Intrawave_v3;
   }
-  else if (k == 5120) {
-    if (n <= 32) {
-      return bgemm_kernel_bf16bf16bf16_128_16x32x64_16x16_1x1_8x16x1_8x16x1_1x16x1x8_4_Intrawave_v1;
-    }
-    else if (n <= 512) {
-      return bgemm_kernel_bf16bf16bf16_256_128x128x64_32x32_2x2_8x32x1_8x32x1_1x16x1x16_4_Intrawave_v3;
-    }
-    else {
-      return bgemm_kernel_bf16bf16bf16_256_256x224x64_16x16_8x7_8x32x1_8x32x1_1x32x1x8_4_Intrawave_v3;
-    }
+  else if (m <= 8192) {
+    if (n <= 8) return bgemm_kernel_bf16bf16bf16_128_16x32x64_16x16_1x1_8x16x1_8x16x1_1x16x1x8_4_Intrawave_v1; // 3 options available, need to investigate
+    if (n <= 32) return bgemm_kernel_bf16bf16bf16_128_16x64x64_16x16_1x2_8x16x1_8x16x1_1x16x1x8_4_Intrawave_v2;
+    if (n <= 512) return bgemm_kernel_bf16bf16bf16_256_128x128x64_32x32_2x2_8x32x1_8x32x1_1x16x1x16_4_Intrawave_v3;
+    if (n <= 4096) return bgemm_kernel_bf16bf16bf16_256_256x224x64_16x16_8x7_8x32x1_8x32x1_1x32x1x8_4_Intrawave_v3;
   }
 
   // Default instance