[BENCHMARK][GEMM] Re-enable CUTLASS's edge cases (#4181)

This PR updates the CUTLASS pin to be in sync with this [PR
364](intel/sycl-tla#364), that
fixes the edge cases encountered with CUTLASS in our integration.
Thanks to this update, we can now also re-enable the previously disabled
CUTLASS edge cases.

---------

Signed-off-by: Jefferson Le Quellec <[email protected]>

Author: jle-quel

benchmarks/cmake/FindCUTLASSLibrary.cmake

@@ -9,13 +9,13 @@ if (NOT CUTLASSLibrary_FOUND)
     set(CUTLASSLibrary_SOURCE_DIR
             "${CMAKE_CURRENT_BINARY_DIR}/CUTLASSLibrary")
     message(STATUS "CUTLASSLibrary is not specified. Will try to download
-                  CUTLASS library from https://github.com/codeplaysoftware/cutlass-fork into
+                  CUTLASS library from https://github.com/codeplaysoftware/cutlass-sycl.git into
                   ${CUTLASSLibrary_SOURCE_DIR}")
     file(READ cutlass_kernel/cutlass-library.conf CUTLASSLibrary_TAG)
     # Strip the potential trailing newline from tag
     string(STRIP "${CUTLASSLibrary_TAG}" CUTLASSLibrary_TAG)
     FetchContent_Declare(cutlass-library
-            GIT_REPOSITORY    https://github.com/codeplaysoftware/cutlass-fork
+            GIT_REPOSITORY    https://github.com/codeplaysoftware/cutlass-sycl.git
             GIT_TAG           ${CUTLASSLibrary_TAG}
             SOURCE_DIR ${CUTLASSLibrary_SOURCE_DIR}
             )

benchmarks/cutlass_kernel/cutlass-library.conf

@@ -1 +1 @@
-f36600cfc1c3a60d9b1d6f64f946a1e877ea33bd
+bb48e86d2fe7cb09eab2e719e78d5811d3da3131

benchmarks/cutlass_kernel/python_main.cpp

@@ -56,9 +56,9 @@ static auto gemm_run(const at::Tensor &A, const at::Tensor &B, at::Tensor &C,
 
   using CollectiveMainloop =
       typename cutlass::gemm::collective::CollectiveBuilder<
-          cutlass::arch::IntelPVC, cutlass::arch::OpClassTensorOp,
-          ElementInputA, LayoutA, AlignmentA, ElementInputB, LayoutB,
-          AlignmentB, ElementAccumulator, TileShape,
+          cutlass::arch::IntelXe, cutlass::arch::OpClassTensorOp, ElementInputA,
+          LayoutA, AlignmentA, ElementInputB, LayoutB, AlignmentB,
+          ElementAccumulator, TileShape,
           cute::Shape<cute::_1, cute::_1, cute::_1>,
           cutlass::gemm::collective::StageCountAuto,
           cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp;
@@ -71,7 +71,7 @@ static auto gemm_run(const at::Tensor &A, const at::Tensor &B, at::Tensor &C,
       cutlass::FloatRoundStyle::round_to_nearest>;
   using CollectiveEpilogue =
       typename cutlass::epilogue::collective::CollectiveBuilder<
-          cutlass::arch::IntelPVC, cutlass::arch::OpClassTensorOp, TileShape,
+          cutlass::arch::IntelXe, cutlass::arch::OpClassTensorOp, TileShape,
           cute::Shape<cute::_1, cute::_1, cute::_1>,
           cutlass::epilogue::collective::EpilogueTileAuto,
           ElementComputeEpilogue, ElementAccumulator, ElementAccumulator,
@@ -184,9 +184,8 @@ auto gemm(const at::Tensor &A, const at::Tensor &B, at::Tensor &C, const int M,
     return gemm_run<TileShape_RRR_3>(A, B, C, M, N, K, L);
   if (test_case == Dim{4096, 8, 128, 16384})
     return gemm_run<TileShape_RRR_4>(A, B, C, M, N, K, L);
-  /// FIXME: Getting a compile time error for RRR_5
-  // if (test_case == Dimension{4096, 8, 16384, 128})
-  //   return gemm_run<TileShape_RRR_5>(A, B, C, M, N, K, L);
+  if (test_case == Dim{4096, 8, 16384, 128})
+    return gemm_run<TileShape_RRR_5>(A, B, C, M, N, K, L);
 
   return gemm_run<TileShape_RRR_1>(A, B, C, M, N, K, L);
 }

benchmarks/triton_kernels_benchmark/gemm_benchmark.py

@@ -433,21 +433,14 @@ def cutlass_invoker():
             cutlass_fn = cutlass_invoker
             torch_fn = lambda: torch.matmul(torch_a, torch_b).to(torch.float32)
 
-            # FIXME: Remove temporary condition when https://github.com/codeplaysoftware/cutlass-fork/pull/313 will be merged
-            if (B, M, N, K) == (4096, 8, 128, 16384) or (B, M, N, K) == (4096, 8, 16384, 128):
-                min_ms = float('nan')
-                max_ms = float('nan')
-                mean_ms = float('nan')
-                cv = float('nan')
-            else:
-                rtol = 1e-2 if a.dtype == torch.bfloat16 else 1e-3
-                benchmark_suite.assert_close(cutlass_fn, torch_fn, atol=1e-4, rtol=rtol, err_msg='cutlass to torch')
-                _, min_ms, max_ms, mean_ms, cv = benchmark_suite.do_bench(
-                    cutlass_fn,
-                    n_warmup=10,
-                    n_repeat=10,
-                    quantiles=quantiles,
-                )
+            rtol = 1e-2 if a.dtype == torch.bfloat16 else 1e-3
+            benchmark_suite.assert_close(cutlass_fn, torch_fn, atol=1e-4, rtol=rtol, err_msg='cutlass to torch')
+            _, min_ms, max_ms, mean_ms, cv = benchmark_suite.do_bench(
+                cutlass_fn,
+                n_warmup=10,
+                n_repeat=10,
+                quantiles=quantiles,
+            )
 
         else:
             raise NotImplementedError(f'Unsupported provider {provider}')