From f396cdd15ccf873d5c92ae73b73ed680c77f4400 Mon Sep 17 00:00:00 2001
From: dan_the_3rd <43445237+danthe3rd@users.noreply.github.com>
Date: Wed, 1 Mar 2023 13:13:51 +0100
Subject: [PATCH] ex24[gemm_grouped]: Allow to change layout/dtype (#841)

* ex24[gemm_grouped]: Allow to change layout/dtype

* Address suggestion from @jackkosaian

---------

Co-authored-by: danthe3rd <danthe3rd>
---
 examples/24_gemm_grouped/gemm_grouped.cu          | 8 ++++----
 include/cutlass/gemm/device/gemm_universal_base.h | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/examples/24_gemm_grouped/gemm_grouped.cu b/examples/24_gemm_grouped/gemm_grouped.cu
index f70878d5..71b77b04 100644
--- a/examples/24_gemm_grouped/gemm_grouped.cu
+++ b/examples/24_gemm_grouped/gemm_grouped.cu
@@ -1487,8 +1487,8 @@ int main(int argc, char const **args) {
 
   // Gemm operator cutlass_tensorop_f16_s16816gemm_f16_128x128_32x4_nt_align8
   using GemmBatched = cutlass::gemm::device::GemmUniversal<
-    cutlass::half_t, LayoutA,
-    cutlass::half_t, LayoutB,
+    ElementA, LayoutA,
+    ElementB, LayoutB,
     ElementOutput,   LayoutC,
     ElementAccumulator,
     cutlass::arch::OpClassTensorOp,
@@ -1510,11 +1510,11 @@ int main(int argc, char const **args) {
   // for scheduling mode. This will be used as the template for all scheduling
   // modes executed.
   using GemmKernel = typename cutlass::gemm::kernel::DefaultGemmGrouped<
-    cutlass::half_t, 
+    ElementA,
     LayoutA,
     cutlass::ComplexTransform::kNone,
     8,
-    cutlass::half_t,
+    ElementB,
     LayoutB,
     cutlass::ComplexTransform::kNone,
     8,
diff --git a/include/cutlass/gemm/device/gemm_universal_base.h b/include/cutlass/gemm/device/gemm_universal_base.h
index cca768ae..6fb1753e 100644
--- a/include/cutlass/gemm/device/gemm_universal_base.h
+++ b/include/cutlass/gemm/device/gemm_universal_base.h
@@ -306,7 +306,7 @@ public:
   /// Initializes GEMM state from arguments and workspace memory
   Status initialize(
     Arguments const &args,
-    void *workspace,
+    void *workspace = nullptr,
     cudaStream_t stream = nullptr)
   {
     CUTLASS_TRACE_HOST("GemmUniversalBase::initialize() - workspace "