CUTLASS 3.5.0 (#1411)

include/cutlass/gemm/kernel/default_gemm.h

@@ -218,6 +218,84 @@ struct DefaultGemm<ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignment
 
 ////////////////////////////////////////////////////////////////////////////////
 
+/// Partial specialization for Ada Architecture
+template <
+    /// Element type for A matrix operand
+    typename ElementA,
+    /// Layout type for A matrix operand
+    typename LayoutA,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentA,
+    /// Element type for B matrix operand
+    typename ElementB,
+    /// Layout type for B matrix operand
+    typename LayoutB,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentB,
+    /// Element type for C and D matrix operands
+    typename ElementC,
+    /// Element type for internal accumulation
+    typename ElementAccumulator,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename InstructionShape,
+    /// Epilogue output operator
+    typename EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle,
+    /// Number of stages used in the pipelined mainloop
+    int Stages,
+    /// If true, kernel is configured to support serial reduction in the
+    /// epilogue
+    bool SplitKSerial,
+    /// Operation performed by GEMM
+    typename Operator,
+    /// Use zfill or predicate for out-of-bound cp.async
+    SharedMemoryClearOption SharedMemoryClear,
+    /// Gather operand A by using an index array
+    bool GatherA,
+    /// Gather operand B by using an index array
+    bool GatherB,
+    /// Scatter result D by using an index array
+    bool ScatterD,
+    /// Permute result D
+    typename PermuteDLayout,
+    /// Permute operand A
+    typename PermuteALayout,
+    /// Permute operand B
+    typename PermuteBLayout
+>
+struct DefaultGemm<ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, ElementC,
+                   layout::RowMajor, ElementAccumulator, arch::OpClassTensorOp,
+                   arch::Sm89, ThreadblockShape, WarpShape, InstructionShape,
+                   EpilogueOutputOp, ThreadblockSwizzle, Stages, SplitKSerial,
+                   Operator, SharedMemoryClear, GatherA, GatherB, ScatterD, 
+                   PermuteDLayout, PermuteALayout, PermuteBLayout> {
+  /// Define the threadblock-scoped matrix multiply-accumulate
+  using Mma = typename cutlass::gemm::threadblock::DefaultMma<
+      ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB,
+      ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm89,
+      ThreadblockShape, WarpShape, InstructionShape, Stages,
+      Operator, false, SharedMemoryClear, GatherA, GatherB,
+      PermuteALayout, PermuteBLayout>::ThreadblockMma;
+
+  static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK;
+
+  /// Define the epilogue
+  using Epilogue =
+      typename cutlass::epilogue::threadblock::DefaultEpilogueTensorOp<
+          ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp,
+          EpilogueOutputOp::kCount, ScatterD, PermuteDLayout>::Epilogue;
+
+  /// Define the kernel-level GEMM operator.
+  using GemmKernel = kernel::Gemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
 /// Partial specialization for Ampere Architecture
 template <
     /// Element type for A matrix operand

include/cutlass/gemm/kernel/default_gemm_sparse.h

@@ -183,6 +183,67 @@ struct DefaultSparseGemm<ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAli
   using GemmKernel = kernel::SparseGemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
 };
 
+///////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for Ada Architecture
+template <
+    /// Element type for A matrix operand
+    typename ElementA,
+    /// Layout type for A matrix operand
+    typename LayoutA,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentA,
+    /// Element type for B matrix operand
+    typename ElementB,
+    /// Layout type for B matrix operand
+    typename LayoutB,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentB,
+    /// Element type for C and D matrix operands
+    typename ElementC,
+    /// Element type for internal accumulation
+    typename ElementAccumulator,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename InstructionShape,
+    /// Epilogue output operator
+    typename EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle,
+    /// Number of stages used in the pipelined mainloop
+    int Stages,
+    /// If true, kernel is configured to support serial reduction in the
+    /// epilogue
+    bool SplitKSerial,
+    /// Operation performed by GEMM
+    typename Operator>
+struct DefaultSparseGemm<ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB, ElementC,
+                   layout::RowMajor, ElementAccumulator, arch::OpClassTensorOp,
+                   arch::Sm89, ThreadblockShape, WarpShape, InstructionShape,
+                   EpilogueOutputOp, ThreadblockSwizzle, Stages, SplitKSerial,
+                   Operator> {
+  /// Define the threadblock-scoped matrix multiply-accumulate
+  using Mma = typename cutlass::gemm::threadblock::DefaultSparseMma<
+      ElementA, LayoutA, kAlignmentA, ElementB, LayoutB, kAlignmentB,
+      ElementAccumulator, layout::RowMajor, arch::OpClassTensorOp, arch::Sm89,
+      ThreadblockShape, WarpShape, InstructionShape, Stages,
+      Operator>::ThreadblockMma;
+
+  static const int kPartitionsK = ThreadblockShape::kK / WarpShape::kK;
+
+  /// Define the epilogue
+  using Epilogue =
+      typename cutlass::epilogue::threadblock::DefaultEpilogueTensorOp<
+          ThreadblockShape, typename Mma::Operator, kPartitionsK, EpilogueOutputOp,
+          EpilogueOutputOp::kCount>::Epilogue;
+
+  /// Define the kernel-level GEMM operator.
+  using GemmKernel = kernel::SparseGemm<Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
+};
+
 ////////////////////////////////////////////////////////////////////////////////
 
 }  // namespace kernel

include/cutlass/gemm/kernel/default_gemm_sparse_with_absmax.h

@@ -0,0 +1,157 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief
+    Default configuration for a sparse GEMM with fused absolute-maximum calculations and scaling
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+
+#include "cutlass/layout/matrix.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/arch/wmma.h"
+
+#include "cutlass/epilogue/threadblock/epilogue.h"
+#include "cutlass/epilogue/thread/linear_combination.h"
+
+#include "cutlass/gemm/gemm.h"
+#include "cutlass/gemm/kernel/gemm.h"
+#include "cutlass/gemm/kernel/sparse_gemm_with_absmax.h"
+#include "cutlass/gemm/kernel/default_gemm_sparse.h"
+#include "cutlass/gemm/kernel/gemm_pipelined.h"
+#include "cutlass/gemm/threadblock/default_mma_core_sm75.h"
+#include "cutlass/gemm/threadblock/default_mma_core_sm70.h"
+#include "cutlass/gemm/threadblock/default_mma_core_sm80.h"
+#include "cutlass/gemm/threadblock/default_mma_core_sparse_sm80.h"
+#include "cutlass/gemm/threadblock/default_sparse_mma.h"
+#include "cutlass/gemm/threadblock/default_mma_core_simt.h"
+#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
+
+#include "cutlass/epilogue/threadblock/default_epilogue_with_absmax.h"
+#include "cutlass/epilogue/threadblock/default_epilogue_tensor_op.h"
+#include "cutlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h"
+#include "cutlass/epilogue/threadblock/default_epilogue_simt.h"
+#include "cutlass/transform/threadblock/predicated_tile_iterator.h"
+
+#if defined(CUTLASS_ARCH_WMMA_ENABLED)
+#include "cutlass/epilogue/threadblock/default_epilogue_wmma_tensor_op.h"
+#endif //CUTLASS_ARCH_WMMA_ENABLED
+
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace kernel {
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <
+    /// Element type for A matrix operand
+    typename ElementA_,
+    /// Layout type for A matrix operand
+    typename LayoutA_,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentA,
+    /// Element type for B matrix operand
+    typename ElementB_,
+    /// Layout type for B matrix operand
+    typename LayoutB_,
+    /// Access granularity of B matrix in units of elements
+    int kAlignmentB,
+    /// Element type for C and D matrix operands
+    typename ElementC_,
+    /// Layout type for C and D matrix operands
+    typename LayoutC_,
+    /// Element type for internal accumulation
+    typename ElementAccumulator,
+    /// Operator class tag
+    typename OperatorClass,
+    /// Tag indicating architecture to tune for
+    typename ArchTag,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename InstructionShape,
+    /// Epilogue output operator
+    typename EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle,
+    /// Number of stages used in the pipelined mainloop
+    int Stages,
+    /// If true, kernel is configured to support serial reduction in the
+    /// epilogue
+    bool SplitKSerial,
+    /// Operation performed by GEMM
+    typename Operator>
+struct DefaultSparseGemmWithAbsmax {
+
+  using GemmBase = typename DefaultSparseGemm<
+    ElementA_, LayoutA_, kAlignmentA,
+    ElementB_, LayoutB_, kAlignmentB,
+    ElementC_, LayoutC_, ElementAccumulator,
+    OperatorClass,
+    ArchTag,
+    ThreadblockShape,
+    WarpShape,
+    InstructionShape,
+    EpilogueOutputOp,
+    ThreadblockSwizzle,
+    Stages,
+    SplitKSerial,
+    Operator
+  >::GemmKernel;
+
+  // Define epilogue
+  using Epilogue = typename cutlass::epilogue::threadblock::DefaultEpilogueWithAbsMax<
+    typename GemmBase::Epilogue::Shape,
+    typename GemmBase::Epilogue::WarpMmaOperator,
+    GemmBase::Epilogue::kPartitionsK,
+    ElementC_,
+    typename EpilogueOutputOp::ElementAuxOutput,
+    ElementC_,
+    EpilogueOutputOp,
+    GemmBase::Epilogue::kElementsPerAccess
+  >::Epilogue;
+
+  /// Define the kernel-level GEMM operator.
+  using GemmKernel = kernel::SparseGemmWithAbsmax<typename GemmBase::Mma, Epilogue, ThreadblockSwizzle, SplitKSerial>;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace kernel
+}  // namespace gemm
+}  // namespace cutlass

include/cutlass/gemm/kernel/default_gemm_with_absmax.h

@@ -0,0 +1,143 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+/*! \file
+  \brief
+    Default configuration for a GEMM with fused absolute-maximum calculations and scaling
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+
+#include "cutlass/gemm/kernel/gemm_with_absmax.h"
+#include "cutlass/gemm/kernel/default_gemm_universal.h"
+
+#include "cutlass/epilogue/threadblock/default_epilogue_with_absmax.h"
+#include "cutlass/epilogue/threadblock/epilogue_with_absmax.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace kernel {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  /// Element type for A matrix operand
+  typename ElementA_,
+  /// Layout type for A matrix operand
+  typename LayoutA_,
+  /// Complex elementwise transformation on A operand
+  ComplexTransform TransformA,
+  /// Access granularity of A matrix in units of elements
+  int kAlignmentA,
+  /// Element type for B matrix operand
+  typename ElementB_,
+  /// Layout type for B matrix operand
+  typename LayoutB_,
+  /// Complex elementwise transformation on B operand
+  ComplexTransform TransformB,
+  /// Access granularity of B matrix in units of elements
+  int kAlignmentB,
+  /// Element type for C and D matrix operands
+  typename ElementC_,
+  /// Layout type for C and D matrix operands
+  typename LayoutC_,
+  /// Element type for internal accumulation
+  typename ElementAccumulator,
+  /// Operator class tag
+  typename OperatorClass,
+  /// Tag indicating architecture to tune for
+  typename ArchTag,
+  /// Threadblock-level tile size (concept: GemmShape)
+  typename ThreadblockShape,
+  /// Warp-level tile size (concept: GemmShape)
+  typename WarpShape,
+  /// Warp-level tile size (concept: GemmShape)
+  typename InstructionShape,
+  /// Epilogue output operator
+  typename EpilogueOutputOp,
+  /// Threadblock-level swizzling operator
+  typename ThreadblockSwizzle,
+  /// Number of stages used in the pipelined mainloop
+  int Stages,
+  /// Operation performed by GEMM
+  typename Operator,
+  ///
+  typename Enable = void
+>
+struct DefaultGemmWithAbsMax {
+
+  using GemmBase = typename DefaultGemmUniversal<
+    ElementA_, LayoutA_, TransformA, kAlignmentA,
+    ElementB_, LayoutB_, TransformB, kAlignmentB,
+    ElementC_, LayoutC_, ElementAccumulator,
+    OperatorClass,
+    ArchTag,
+    ThreadblockShape,
+    WarpShape,
+    InstructionShape,
+    EpilogueOutputOp,
+    ThreadblockSwizzle,
+    Stages,
+    Operator
+  >::GemmKernel;
+
+  // Define epilogue
+  using Epilogue = typename cutlass::epilogue::threadblock::DefaultEpilogueWithAbsMax<
+    typename GemmBase::Epilogue::Shape,
+    typename GemmBase::Epilogue::WarpMmaOperator,
+    GemmBase::Epilogue::kPartitionsK,
+    ElementC_,
+    typename EpilogueOutputOp::ElementAuxOutput,
+    ElementC_,
+    EpilogueOutputOp,
+    GemmBase::Epilogue::kElementsPerAccess
+  >::Epilogue;
+
+  // Compose the GEMM kernel
+  using GemmKernel = GemmWithAbsMax<
+    typename GemmBase::Mma,
+    Epilogue,
+    ThreadblockSwizzle
+  >;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace kernel
+}  // namespace gemm
+}  // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

include/cutlass/gemm/kernel/default_rank_2k_grouped.h

@@ -165,7 +165,7 @@ struct DefaultRank2KGrouped<ElementA, LayoutA, TransformA, kAlignmentA,
           FillModeC, ElementAccumulator, OperatorClass, ArchTag, ThreadblockShape,
           WarpShape, InstructionShape, EpilogueOutputOp,
           ThreadblockSwizzle, Stages, Operator, BlasMode_, GroupScheduleMode_,
-          typename std::enable_if< ! cutlass::is_complex<ElementAccumulator>::value>::type
+          typename platform::enable_if< ! cutlass::is_complex<ElementAccumulator>::value>::type
 > {
   // If true, we must construct a 'transposed-and-exchanged' Rank2K operator.
   static bool const kInternalTranspose = platform::is_same<LayoutC, layout::ColumnMajor>::value;
@@ -283,7 +283,7 @@ struct DefaultRank2KGrouped<ElementA, LayoutA, TransformA, kAlignmentA,
           FillModeC, ElementAccumulator, OperatorClass, ArchTag, ThreadblockShape,
           WarpShape, InstructionShape, EpilogueOutputOp,
           ThreadblockSwizzle, Stages, Operator, BlasMode_, GroupScheduleMode_,
-          typename std::enable_if<cutlass::is_complex<ElementAccumulator>::value>::type
+          typename platform::enable_if<cutlass::is_complex<ElementAccumulator>::value>::type
 > {
   // If true, we must construct a 'transposed-and-exchanged' Rank2K operator.
   static bool const kInternalTranspose = platform::is_same<LayoutC, layout::ColumnMajor>::value;

include/cutlass/gemm/kernel/default_rank_2k_universal.h

@@ -183,7 +183,7 @@ struct DefaultRank2KUniversal<
   SplitKSerial,
   Operator,
   BlasMode::kSymmetric,
-  typename std::enable_if< ! cutlass::is_complex<ElementAccumulator>::value>::type
+  typename platform::enable_if< ! cutlass::is_complex<ElementAccumulator>::value>::type
 > {
 
   using DefaultRank2Kkernel = typename kernel::DefaultRank2K<
@@ -301,7 +301,7 @@ struct DefaultRank2KUniversal<
   SplitKSerial,
   Operator,
   kBlasMode,
-  typename std::enable_if<cutlass::is_complex<ElementAccumulator>::value>::type
+  typename platform::enable_if<cutlass::is_complex<ElementAccumulator>::value>::type
 > {
 
   using DefaultRank2Kkernel = typename kernel::DefaultRank2KComplex<

include/cutlass/gemm/kernel/default_rank_k_universal.h

@@ -165,7 +165,7 @@ struct DefaultRankKUniversal<
   SplitKSerial,
   Operator,
   BlasMode::kSymmetric,
-  typename std::enable_if< ! cutlass::is_complex<ElementAccumulator>::value>::type
+  typename platform::enable_if< ! cutlass::is_complex<ElementAccumulator>::value>::type
 > {
 
   using DefaultRankKkernel = typename kernel::DefaultRankK<
@@ -265,7 +265,7 @@ struct DefaultRankKUniversal<
   SplitKSerial,
   Operator,
   kBlasMode,
-  typename std::enable_if<cutlass::is_complex<ElementAccumulator>::value>::type
+  typename platform::enable_if<cutlass::is_complex<ElementAccumulator>::value>::type
 > {
 
   using DefaultRankKkernel = typename kernel::DefaultRankKComplex<

include/cutlass/gemm/kernel/default_symm_universal.h

@@ -182,7 +182,7 @@ struct DefaultSymmUniversal<
   SplitKSerial,
   Operator,
   BlasMode::kSymmetric,
-  typename std::enable_if< ! cutlass::is_complex<ElementAccumulator>::value>::type
+  typename platform::enable_if< ! cutlass::is_complex<ElementAccumulator>::value>::type
 > {
 
   using DefaultSymmkernel = typename kernel::DefaultSymm<
@@ -298,7 +298,7 @@ struct DefaultSymmUniversal<
   SplitKSerial,
   Operator,
   kBlasMode,
-  typename std::enable_if<cutlass::is_complex<ElementAccumulator>::value>::type
+  typename platform::enable_if<cutlass::is_complex<ElementAccumulator>::value>::type
 > {
 
   using DefaultSymmkernel = typename kernel::DefaultSymmComplex<

include/cutlass/gemm/kernel/default_trmm_universal.h

@@ -190,7 +190,7 @@ struct DefaultTrmmUniversal<
   Stages,
   SplitKSerial,
   Operator,
-  typename std::enable_if< ! cutlass::is_complex<ElementAccumulator>::value>::type
+  typename platform::enable_if< ! cutlass::is_complex<ElementAccumulator>::value>::type
 > {
 
   using DefaultTrmmKernel = typename kernel::DefaultTrmm<
@@ -311,7 +311,7 @@ struct DefaultTrmmUniversal<
   Stages,
   SplitKSerial,
   Operator,
-  typename std::enable_if<cutlass::is_complex<ElementAccumulator>::value>::type
+  typename platform::enable_if<cutlass::is_complex<ElementAccumulator>::value>::type
 > {
 
   using DefaultTrmmKernel = typename kernel::DefaultTrmmComplex<

include/cutlass/gemm/kernel/ell_gemm.h

@@ -73,32 +73,30 @@ struct EllGemm {
 
   /// Parameters structure
   struct Params {
-    cutlass::gemm::GemmCoord problem_size;
-    cutlass::gemm::GemmCoord grid_tiled_shape;
-    int swizzle_log_tile;
-    typename Mma::IteratorA::Params params_A;
-    typename Mma::IteratorA::TensorRef ref_A;
-    typename Mma::IteratorB::Params params_B;
-    typename Mma::IteratorB::TensorRef ref_B;
-    typename Epilogue::OutputTileIterator::Params params_C;
-    typename Epilogue::OutputTileIterator::TensorRef ref_C;
-    typename Epilogue::OutputTileIterator::Params params_D;
-    typename Epilogue::OutputTileIterator::TensorRef ref_D;
-    typename OutputOp::Params output_op;
-    int *semaphore;
-    int gemm_k_iterations;
-    int gemm_k_size;
-    const int* ell_idx;
-    int ell_ncol;
-    int ell_blocksize;
-    int ell_base_idx;
+    cutlass::gemm::GemmCoord problem_size{};
+    cutlass::gemm::GemmCoord grid_tiled_shape{};
+    int swizzle_log_tile{0};
+    typename Mma::IteratorA::Params params_A{};
+    typename Mma::IteratorA::TensorRef ref_A{};
+    typename Mma::IteratorB::Params params_B{};
+    typename Mma::IteratorB::TensorRef ref_B{};
+    typename Epilogue::OutputTileIterator::Params params_C{};
+    typename Epilogue::OutputTileIterator::TensorRef ref_C{};
+    typename Epilogue::OutputTileIterator::Params params_D{};
+    typename Epilogue::OutputTileIterator::TensorRef ref_D{};
+    typename OutputOp::Params output_op{};
+    int *semaphore = nullptr;
+    int gemm_k_iterations{0};
+    int gemm_k_size{0};
+    const int* ell_idx = nullptr;
+    int ell_ncol{0};
+    int ell_blocksize{0};
+    int ell_base_idx{0};
 
     //
     // Methods
     //
-
-    CUTLASS_HOST_DEVICE
-    Params(): swizzle_log_tile(0), semaphore(0), gemm_k_iterations(0), gemm_k_size(0) { }
+   Params() = default;
 
     CUTLASS_HOST_DEVICE
     Params(
@@ -154,9 +152,7 @@ struct EllGemm {
   //
   // Methods
   //
-
-  CUTLASS_HOST_DEVICE
-  EllGemm() { }
+  EllGemm() = default;
 
   /// Determines whether kernel satisfies alignment
     static Status can_implement(
@@ -458,32 +454,30 @@ struct EllGemm<Mma_, Epilogue_, ThreadblockSwizzle_, SplitKSerial, false> {
 
   /// Parameters structure
   struct Params {
-    cutlass::gemm::GemmCoord problem_size;
-    cutlass::gemm::GemmCoord grid_tiled_shape;
-    int swizzle_log_tile;
-    typename Mma::IteratorA::Params params_A;
-    typename Mma::IteratorA::TensorRef ref_A;
-    typename Mma::IteratorB::Params params_B;
-    typename Mma::IteratorB::TensorRef ref_B;
-    typename Epilogue::OutputTileIterator::Params params_C;
-    typename Epilogue::OutputTileIterator::TensorRef ref_C;
-    typename Epilogue::OutputTileIterator::Params params_D;
-    typename Epilogue::OutputTileIterator::TensorRef ref_D;
-    typename OutputOp::Params output_op;
-    int *semaphore;
-    int gemm_k_iterations;
-    int gemm_k_size;
-    const int* ell_idx;
-    int ell_ncol;
-    int ell_blocksize;
-    int ell_base_idx;
+    cutlass::gemm::GemmCoord problem_size{};
+    cutlass::gemm::GemmCoord grid_tiled_shape{};
+    int swizzle_log_tile{0};
+    typename Mma::IteratorA::Params params_A{};
+    typename Mma::IteratorA::TensorRef ref_A{};
+    typename Mma::IteratorB::Params params_B{};
+    typename Mma::IteratorB::TensorRef ref_B{};
+    typename Epilogue::OutputTileIterator::Params params_C{};
+    typename Epilogue::OutputTileIterator::TensorRef ref_C{};
+    typename Epilogue::OutputTileIterator::Params params_D{};
+    typename Epilogue::OutputTileIterator::TensorRef ref_D{};
+    typename OutputOp::Params output_op{};
+    int *semaphore = nullptr;
+    int gemm_k_iterations{0};
+    int gemm_k_size{0};
+    const int* ell_idx = nullptr;
+    int ell_ncol{0};
+    int ell_blocksize{0};
+    int ell_base_idx{0};
 
     //
     // Methods
     //
-
-    CUTLASS_HOST_DEVICE
-    Params(): swizzle_log_tile(0), semaphore(0), gemm_k_iterations(0), gemm_k_size(0) { }
+    Params() = default;
 
     CUTLASS_HOST_DEVICE
     Params(

include/cutlass/gemm/kernel/gemm_batched.h

@@ -65,31 +65,29 @@ struct GemmBatched {
 
   /// Parameters structure
   struct Params {
-    cutlass::gemm::GemmCoord problem_size;
-    cutlass::gemm::GemmCoord grid_tiled_shape;
-    int swizzle_log_tile;
-    typename Mma::IteratorA::Params params_A;
-    typename Mma::IteratorA::TensorRef ref_A;
-    int64_t stride_A;
-    typename Mma::IteratorB::Params params_B;
-    typename Mma::IteratorB::TensorRef ref_B;
-    int64_t stride_B;
-    typename Epilogue::OutputTileIterator::Params params_C;
-    typename Epilogue::OutputTileIterator::TensorRef ref_C;
-    int64_t stride_C;
-    typename Epilogue::OutputTileIterator::Params params_D;
-    typename Epilogue::OutputTileIterator::TensorRef ref_D;
-    int64_t stride_D;
-    typename OutputOp::Params epilogue;
-    int batch_count;
-    int gemm_k_iterations;
+    cutlass::gemm::GemmCoord problem_size{};
+    cutlass::gemm::GemmCoord grid_tiled_shape{};
+    int swizzle_log_tile{0};
+    typename Mma::IteratorA::Params params_A{};
+    typename Mma::IteratorA::TensorRef ref_A{};
+    int64_t stride_A{0};
+    typename Mma::IteratorB::Params params_B{};
+    typename Mma::IteratorB::TensorRef ref_B{};
+    int64_t stride_B{0};
+    typename Epilogue::OutputTileIterator::Params params_C{};
+    typename Epilogue::OutputTileIterator::TensorRef ref_C{};
+    int64_t stride_C{0};
+    typename Epilogue::OutputTileIterator::Params params_D{};
+    typename Epilogue::OutputTileIterator::TensorRef ref_D{};
+    int64_t stride_D{0};
+    typename OutputOp::Params epilogue{};
+    int batch_count{1};
+    int gemm_k_iterations{0};
 
     //
     // Methods
     //
-
-    CUTLASS_HOST_DEVICE
-    Params() : swizzle_log_tile(0) { }
+    Params() = default;
 
     CUTLASS_HOST_DEVICE
     Params(
@@ -123,9 +121,7 @@ struct GemmBatched {
       stride_D(stride_D_),
       epilogue(epilogue_),
       batch_count(batch_count_),
-      gemm_k_iterations((problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK) {
-
-    }
+      gemm_k_iterations((problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK) {}
   };
 
   /// Shared memory storage structure
@@ -137,9 +133,7 @@ struct GemmBatched {
   //
   // Methods
   //
-
-  CUTLASS_HOST_DEVICE
-  GemmBatched() { } 
+  GemmBatched() = default;
 
   /// Executes one GEMM
   CUTLASS_DEVICE

include/cutlass/gemm/kernel/gemm_grouped.h

@@ -133,46 +133,32 @@ public:
     // Data members
     //
 
-    GemmCoord *problem_sizes;
-    int problem_count;
-    int threadblock_count;
+    GemmCoord *problem_sizes{nullptr};
+    int problem_count{0};
+    int threadblock_count{0};
 
-    typename EpilogueOutputOp::Params output_op;
+    typename EpilogueOutputOp::Params output_op{};
 
-    ElementA ** ptr_A;
-    ElementB ** ptr_B;
-    ElementC ** ptr_C;
-    ElementC ** ptr_D;
+    ElementA ** ptr_A{nullptr};
+    ElementB ** ptr_B{nullptr};
+    ElementC ** ptr_C{nullptr};
+    ElementC ** ptr_D{nullptr};
 
-    typename LayoutA::Stride::LongIndex *lda;
-    typename LayoutB::Stride::LongIndex *ldb;
-    typename LayoutC::Stride::LongIndex *ldc;
-    typename LayoutC::Stride::LongIndex *ldd;
+    typename LayoutA::Stride::LongIndex *lda{nullptr};
+    typename LayoutB::Stride::LongIndex *ldb{nullptr};
+    typename LayoutC::Stride::LongIndex *ldc{nullptr};
+    typename LayoutC::Stride::LongIndex *ldd{nullptr};
 
     // Only used by device-level operator
-    GemmCoord *host_problem_sizes;
+    GemmCoord *host_problem_sizes{nullptr};
+
 
     //
     // Methods
     //
 
     /// Default ctor
-    CUTLASS_HOST_DEVICE
-    Arguments(): 
-      problem_count(0),
-      threadblock_count(0), 
-      ptr_A(nullptr), 
-      ptr_B(nullptr), 
-      ptr_C(nullptr), 
-      ptr_D(nullptr), 
-      lda(nullptr),
-      ldb(nullptr),
-      ldc(nullptr),
-      ldd(nullptr),
-      host_problem_sizes(nullptr)
-    {
-
-    }
+    Arguments() = default;
 
     /// Ctor
     CUTLASS_HOST_DEVICE
@@ -216,36 +202,26 @@ public:
   /// Parameters structure
   struct Params {
 
-    typename ProblemVisitor::Params problem_visitor;
-    int threadblock_count;
+    typename ProblemVisitor::Params problem_visitor{};
+    int threadblock_count{0};
 
-    typename EpilogueOutputOp::Params output_op;
+    typename EpilogueOutputOp::Params output_op{};
 
-    ElementA ** ptr_A;
-    ElementB ** ptr_B;
-    ElementC ** ptr_C;
-    ElementC ** ptr_D;
+    ElementA ** ptr_A{nullptr};
+    ElementB ** ptr_B{nullptr};
+    ElementC ** ptr_C{nullptr};
+    ElementC ** ptr_D{nullptr};
 
-    typename LayoutA::Stride::LongIndex *lda;
-    typename LayoutB::Stride::LongIndex *ldb;
-    typename LayoutC::Stride::LongIndex *ldc;
-    typename LayoutC::Stride::LongIndex *ldd;
+    typename LayoutA::Stride::LongIndex *lda{nullptr};
+    typename LayoutB::Stride::LongIndex *ldb{nullptr};
+    typename LayoutC::Stride::LongIndex *ldc{nullptr};
+    typename LayoutC::Stride::LongIndex *ldd{nullptr};
 
     //
     // Methods
     //
 
-    CUTLASS_HOST_DEVICE
-    Params():
-      ptr_A(nullptr),
-      ptr_B(nullptr),
-      ptr_C(nullptr),
-      ptr_D(nullptr),
-      lda(nullptr),
-      ldb(nullptr),
-      ldc(nullptr),
-      ldd(nullptr)
-    { }
+    Params() = default;
 
     CUTLASS_HOST_DEVICE
     Params(Arguments const &args,

include/cutlass/gemm/kernel/gemm_grouped_softmax_mainloop_fusion.h

@@ -135,50 +135,33 @@ public:
     // Data members
     //
 
-    GemmCoord *problem_sizes;
-    int problem_count;
-    int threadblock_count;
+    GemmCoord *problem_sizes{nullptr};
+    int problem_count{0};
+    int threadblock_count{0};
 
-    typename EpilogueOutputOp::Params output_op;
+    typename EpilogueOutputOp::Params output_op{};
 
-    ElementA ** ptr_A;
-    ElementB ** ptr_B;
-    ElementC ** ptr_C;
-    ElementC ** ptr_D;
-    void ** ptr_norm;
-    void ** ptr_sum;
+    ElementA ** ptr_A{nullptr};
+    ElementB ** ptr_B{nullptr};
+    ElementC ** ptr_C{nullptr};
+    ElementC ** ptr_D{nullptr};
+    void ** ptr_norm{nullptr};
+    void ** ptr_sum{nullptr};
 
-    typename LayoutA::Stride::LongIndex *lda;
-    typename LayoutB::Stride::LongIndex *ldb;
-    typename LayoutC::Stride::LongIndex *ldc;
-    typename LayoutC::Stride::LongIndex *ldd;
+    typename LayoutA::Stride::LongIndex *lda{nullptr};
+    typename LayoutB::Stride::LongIndex *ldb{nullptr};
+    typename LayoutC::Stride::LongIndex *ldc{nullptr};
+    typename LayoutC::Stride::LongIndex *ldd{nullptr};
 
     // Only used by device-level operator
-    GemmCoord *host_problem_sizes;
+    GemmCoord *host_problem_sizes{nullptr};
 
     //
     // Methods
     //
 
     /// Default ctor
-    CUTLASS_HOST_DEVICE
-    Arguments():
-      problem_count(0),
-      threadblock_count(0),
-      ptr_A(nullptr),
-      ptr_B(nullptr),
-      ptr_C(nullptr),
-      ptr_D(nullptr),
-      ptr_norm(nullptr),
-      ptr_sum(nullptr),
-      lda(nullptr),
-      ldb(nullptr),
-      ldc(nullptr),
-      ldd(nullptr),
-      host_problem_sizes(nullptr)
-    {
-
-    }
+    Arguments() = default;
 
     /// Ctor
     CUTLASS_HOST_DEVICE
@@ -226,41 +209,29 @@ public:
   /// Parameters structure
   struct Params {
 
-    typename ProblemVisitor::Params problem_visitor;
-    int threadblock_count;
+    typename ProblemVisitor::Params problem_visitor{};
+    int threadblock_count{0};
 
-    typename EpilogueOutputOp::Params output_op;
+    typename EpilogueOutputOp::Params output_op{};
 
-    ElementA ** ptr_A;
-    ElementB ** ptr_B;
-    ElementC ** ptr_C;
-    ElementC ** ptr_D;
+    ElementA ** ptr_A{nullptr};
+    ElementB ** ptr_B{nullptr};
+    ElementC ** ptr_C{nullptr};
+    ElementC ** ptr_D{nullptr};
 
-    void ** ptr_norm;
-    void ** ptr_sum;
+    void ** ptr_norm{nullptr};
+    void ** ptr_sum{nullptr};
 
-    typename LayoutA::Stride::LongIndex *lda;
-    typename LayoutB::Stride::LongIndex *ldb;
-    typename LayoutC::Stride::LongIndex *ldc;
-    typename LayoutC::Stride::LongIndex *ldd;
+    typename LayoutA::Stride::LongIndex *lda{nullptr};
+    typename LayoutB::Stride::LongIndex *ldb{nullptr};
+    typename LayoutC::Stride::LongIndex *ldc{nullptr};
+    typename LayoutC::Stride::LongIndex *ldd{nullptr};
 
     //
     // Methods
     //
 
-    CUTLASS_HOST_DEVICE
-    Params():
-      ptr_A(nullptr),
-      ptr_B(nullptr),
-      ptr_C(nullptr),
-      ptr_D(nullptr),
-      ptr_norm(nullptr),
-      ptr_sum(nullptr),
-      lda(nullptr),
-      ldb(nullptr),
-      ldc(nullptr),
-      ldd(nullptr)
-    { }
+    Params() = default;
 
     CUTLASS_HOST_DEVICE
     Params(Arguments const &args,

include/cutlass/gemm/kernel/gemm_layernorm_mainloop_fusion.h

@@ -111,59 +111,52 @@ public:
     // Data members
     //
 
-    typename EpilogueOutputOp::Params epilogue;
+    typename EpilogueOutputOp::Params epilogue{};
 
-    void const * ptr_A;
-    void const * ptr_B;
-    void const * ptr_var;
-    void const * ptr_mean;
-    void const * ptr_gamma;
-    void const * ptr_beta;
-    void const * ptr_C;
-    void * ptr_D;
+    void const * ptr_A{nullptr};
+    void const * ptr_B{nullptr};
+    void const * ptr_var{nullptr};
+    void const * ptr_mean{nullptr};
+    void const * ptr_gamma{nullptr};
+    void const * ptr_beta{nullptr};
+    void const * ptr_C{nullptr};
+    void * ptr_D{nullptr};
 
-    int64_t batch_stride_A;
-    int64_t batch_stride_B;
-    int64_t batch_stride_var;
-    int64_t batch_stride_mean;
-    int64_t batch_stride_gamma;
-    int64_t batch_stride_beta;
-    int64_t batch_stride_C;
+    int64_t batch_stride_A{0};
+    int64_t batch_stride_B{0};
+    int64_t batch_stride_var{0};
+    int64_t batch_stride_mean{0};
+    int64_t batch_stride_gamma{0};
+    int64_t batch_stride_beta{0};
+    int64_t batch_stride_C{0};
 
-    typename LayoutA::Stride stride_a;
-    typename LayoutB::Stride stride_b;
-    typename LayoutScaleBias::Stride stride_var;
-    typename LayoutScaleBias::Stride stride_mean;
-    typename LayoutScaleBias::Stride stride_gamma;
-    typename LayoutScaleBias::Stride stride_beta;
-    typename LayoutC::Stride stride_c;
-    typename LayoutC::Stride stride_d;
+    typename LayoutA::Stride stride_a{};
+    typename LayoutB::Stride stride_b{};
+    typename LayoutScaleBias::Stride stride_var{};
+    typename LayoutScaleBias::Stride stride_mean{};
+    typename LayoutScaleBias::Stride stride_gamma{};
+    typename LayoutScaleBias::Stride stride_beta{};
+    typename LayoutC::Stride stride_c{};
+    typename LayoutC::Stride stride_d{};
 
-    typename LayoutA::Stride::LongIndex lda;
-    typename LayoutB::Stride::LongIndex ldb;
-    typename LayoutScaleBias::Stride::LongIndex ld_var;
-    typename LayoutScaleBias::Stride::LongIndex ld_mean;
-    typename LayoutScaleBias::Stride::LongIndex ld_gamma;
-    typename LayoutScaleBias::Stride::LongIndex ld_beta;
-    typename LayoutC::Stride::LongIndex ldc;
-    typename LayoutC::Stride::LongIndex ldd;
+    typename LayoutA::Stride::LongIndex lda{};
+    typename LayoutB::Stride::LongIndex ldb{};
+    typename LayoutScaleBias::Stride::LongIndex ld_var{};
+    typename LayoutScaleBias::Stride::LongIndex ld_mean{};
+    typename LayoutScaleBias::Stride::LongIndex ld_gamma{};
+    typename LayoutScaleBias::Stride::LongIndex ld_beta{};
+    typename LayoutC::Stride::LongIndex ldc{};
+    typename LayoutC::Stride::LongIndex ldd{};
 
-    int const * ptr_gather_A_indices;
-    int const * ptr_gather_B_indices;
-    int const * ptr_scatter_D_indices;
+    int const * ptr_gather_A_indices{nullptr};
+    int const * ptr_gather_B_indices{nullptr};
+    int const * ptr_scatter_D_indices{nullptr};
 
     //
     // Methods
     //
     
-    Arguments(): 
-      ptr_A(nullptr), ptr_B(nullptr), ptr_C(nullptr), ptr_D(nullptr),
-      ptr_var(nullptr), ptr_mean(nullptr),
-      ptr_gamma(nullptr), ptr_beta(nullptr),
-      ptr_gather_A_indices(nullptr),
-      ptr_gather_B_indices(nullptr),
-      ptr_scatter_D_indices(nullptr)
-    {}
+    Arguments() = default;
 
     /// constructs an arguments structure
     Arguments(

include/cutlass/gemm/kernel/gemm_params.h

@@ -70,46 +70,36 @@ struct GemmParams {
   // Data members
   //
 
-  cutlass::gemm::GemmCoord problem_size;
-  cutlass::gemm::GemmCoord grid_tiled_shape;
-  int swizzle_log_tile;
+  cutlass::gemm::GemmCoord problem_size{};
+  cutlass::gemm::GemmCoord grid_tiled_shape{};
+  int swizzle_log_tile{};
 
-  // Data members for Mma::Iterator::Params
-  MmaIteratorParams params_itr_a;
-  MmaIteratorParams params_itr_b;  
+  GemmUniversalMode mode{GemmUniversalMode::kGemm};
+  int batch_count{1};
+  int gemm_k_size{0};
 
-  // Data member for Epilogue::OutputTileIterator::Params 
-  EpilogueIteratorParams params_itr_c;
-  EpilogueIteratorParams params_itr_d;
+  void * ptr_A{nullptr};
+  void * ptr_B{nullptr};
+  void * ptr_C{nullptr};
+  void * ptr_D{nullptr};
 
+  LongIndex lda{0};
+  LongIndex ldb{0};
+  LongIndex ldc{0};
+  LongIndex ldd{0};
 
-  GemmUniversalMode mode;
-  int batch_count;
-  int gemm_k_size;
+  LongIndex batch_stride_A{0};
+  LongIndex batch_stride_B{0};
+  LongIndex batch_stride_C{0};
+  LongIndex batch_stride_D{0};
 
-  void * ptr_A;
-  void * ptr_B;
-  void * ptr_C;
-  void * ptr_D;
-
-  LongIndex lda; 
-  LongIndex ldb; 
-  LongIndex ldc; 
-  LongIndex ldd;
-
-  LongIndex batch_stride_A;
-  LongIndex batch_stride_B;
-  LongIndex batch_stride_C;
-  LongIndex batch_stride_D;
-
-  int *semaphore;
+  int *semaphore{nullptr};
 
   //
   // Methods
   //
 
-  CUTLASS_HOST_DEVICE
-  GemmParams()  {}
+  GemmParams() = default;
 
   CUTLASS_HOST_DEVICE
   GemmParams(

include/cutlass/gemm/kernel/gemm_planar_complex.h

@@ -112,51 +112,39 @@ public:
     // Data members
     //
 
-    typename EpilogueOutputOp::Params epilogue;
+    typename EpilogueOutputOp::Params epilogue{};
 
-    void const * ptr_A_real;
-    void const * ptr_A_imag;
+    void const * ptr_A_real{nullptr};
+    void const * ptr_A_imag{nullptr};
+    void const * ptr_B_real{nullptr};
+    void const * ptr_B_imag{nullptr};
+    void const * ptr_C_real{nullptr};
+    void const * ptr_C_imag{nullptr};
+    void * ptr_D_real{nullptr};
+    void * ptr_D_imag{nullptr};
 
-    void const * ptr_B_real;
-    void const * ptr_B_imag;
-
-    void const * ptr_C_real;
-    void const * ptr_C_imag;
-
-    void * ptr_D_real;
-    void * ptr_D_imag;
-
-    typename LayoutA::Stride::Index lda_real;
-    typename LayoutA::Stride::Index lda_imag;
-    typename LayoutB::Stride::Index ldb_real;
-    typename LayoutB::Stride::Index ldb_imag;
-    typename LayoutC::Stride::Index ldc_real;
-    typename LayoutC::Stride::Index ldc_imag;
-    typename LayoutC::Stride::Index ldd_real;
-    typename LayoutC::Stride::Index ldd_imag;
+    typename LayoutA::Stride::Index lda_real{};
+    typename LayoutA::Stride::Index lda_imag{};
+    typename LayoutB::Stride::Index ldb_real{};
+    typename LayoutB::Stride::Index ldb_imag{};
+    typename LayoutC::Stride::Index ldc_real{};
+    typename LayoutC::Stride::Index ldc_imag{};
+    typename LayoutC::Stride::Index ldd_real{};
+    typename LayoutC::Stride::Index ldd_imag{};
     
-    int64_t batch_stride_A;
-    int64_t batch_stride_A_imag;
-    int64_t batch_stride_B;
-    int64_t batch_stride_B_imag;
-    int64_t batch_stride_C;
-    int64_t batch_stride_C_imag;
-    int64_t batch_stride_D_imag;
+    int64_t batch_stride_A{0};
+    int64_t batch_stride_A_imag{0};
+    int64_t batch_stride_B{0};
+    int64_t batch_stride_B_imag{0};
+    int64_t batch_stride_C{0};
+    int64_t batch_stride_C_imag{0};
+    int64_t batch_stride_D_imag{0};
 
     //
     // Methods
     //
-    
-    Arguments() :
-      ptr_A_real(nullptr), 
-      ptr_A_imag(nullptr), 
-      ptr_B_real(nullptr), 
-      ptr_B_imag(nullptr), 
-      ptr_C_real(nullptr), 
-      ptr_C_imag(nullptr), 
-      ptr_D_real(nullptr),
-      ptr_D_imag(nullptr)
-    {}
+
+    Arguments() = default;
 
     /// constructs an arguments structure
     Arguments(
@@ -260,34 +248,34 @@ public:
     // Data members
     //
 
-    typename Mma::IteratorA::Params params_A_real;
-    typename Mma::IteratorA::Params params_A_imag;
-    typename Mma::IteratorB::Params params_B_real;
-    typename Mma::IteratorB::Params params_B_imag;
-    typename Epilogue::OutputTileIterator::Params params_C_real;
-    typename Epilogue::OutputTileIterator::Params params_C_imag;
-    typename Epilogue::OutputTileIterator::Params params_D_real;
-    typename Epilogue::OutputTileIterator::Params params_D_imag;
+    typename Mma::IteratorA::Params params_A_real{};
+    typename Mma::IteratorA::Params params_A_imag{};
+    typename Mma::IteratorB::Params params_B_real{};
+    typename Mma::IteratorB::Params params_B_imag{};
+    typename Epilogue::OutputTileIterator::Params params_C_real{};
+    typename Epilogue::OutputTileIterator::Params params_C_imag{};
+    typename Epilogue::OutputTileIterator::Params params_D_real{};
+    typename Epilogue::OutputTileIterator::Params params_D_imag{};
     
-    typename EpilogueOutputOp::Params output_op;
+    typename EpilogueOutputOp::Params output_op{};
 
-    void * ptr_A_real;
-    void * ptr_A_imag;
-    void * ptr_B_real;
-    void * ptr_B_imag;
-    void * ptr_C_real;
-    void * ptr_C_imag;
-    void * ptr_D_real;
-    void * ptr_D_imag;
+    void * ptr_A_real{nullptr};
+    void * ptr_A_imag{nullptr};
+    void * ptr_B_real{nullptr};
+    void * ptr_B_imag{nullptr};
+    void * ptr_C_real{nullptr};
+    void * ptr_C_imag{nullptr};
+    void * ptr_D_real{nullptr};
+    void * ptr_D_imag{nullptr};
 
-    int64_t batch_stride_A;
-    int64_t batch_stride_B;
-    int64_t batch_stride_C;
+    int64_t batch_stride_A{0};
+    int64_t batch_stride_B{0};
+    int64_t batch_stride_C{0};
 
-    int64_t batch_stride_A_imag;
-    int64_t batch_stride_B_imag;
-    int64_t batch_stride_C_imag;
-    int64_t batch_stride_D_imag;
+    int64_t batch_stride_A_imag{0};
+    int64_t batch_stride_B_imag{0};
+    int64_t batch_stride_C_imag{0};
+    int64_t batch_stride_D_imag{0};
 
     //
     // Host dispatch API

include/cutlass/gemm/kernel/gemm_planar_complex_array.h

@@ -112,50 +112,38 @@ public:
     // Data members
     //
 
-    typename EpilogueOutputOp::Params epilogue;
+    typename EpilogueOutputOp::Params epilogue{};
 
-    int const *ptr_M;
-    int const *ptr_N;
-    int const *ptr_K;
+    int const *ptr_M{nullptr};
+    int const *ptr_N{nullptr};
+    int const *ptr_K{nullptr};
 
-    void const * const * ptr_A_real;
-    void const * const * ptr_A_imag;
+    void const * const * ptr_A_real{nullptr};
+    void const * const * ptr_A_imag{nullptr};
 
-    void const * const * ptr_B_real;
-    void const * const * ptr_B_imag;
+    void const * const * ptr_B_real{nullptr};
+    void const * const * ptr_B_imag{nullptr};
 
-    void const * const * ptr_C_real;
-    void const * const * ptr_C_imag;
+    void const * const * ptr_C_real{nullptr};
+    void const * const * ptr_C_imag{nullptr};
 
-    void * const * ptr_D_real;
-    void * const * ptr_D_imag;
+    void * const * ptr_D_real{nullptr};
+    void * const * ptr_D_imag{nullptr};
 
-    typename LayoutA::Stride::Index lda_real;
-    typename LayoutA::Stride::Index lda_imag;
-    typename LayoutB::Stride::Index ldb_real;
-    typename LayoutB::Stride::Index ldb_imag;
-    typename LayoutC::Stride::Index ldc_real;
-    typename LayoutC::Stride::Index ldc_imag;
-    typename LayoutC::Stride::Index ldd_real;
-    typename LayoutC::Stride::Index ldd_imag;
+    typename LayoutA::Stride::Index lda_real{};
+    typename LayoutA::Stride::Index lda_imag{};
+    typename LayoutB::Stride::Index ldb_real{};
+    typename LayoutB::Stride::Index ldb_imag{};
+    typename LayoutC::Stride::Index ldc_real{};
+    typename LayoutC::Stride::Index ldc_imag{};
+    typename LayoutC::Stride::Index ldd_real{};
+    typename LayoutC::Stride::Index ldd_imag{};
 
     //
     // Methods
     //
-    
-    Arguments(): 
-      ptr_M(nullptr),
-      ptr_N(nullptr),
-      ptr_K(nullptr),
-      ptr_A_real(nullptr), 
-      ptr_A_imag(nullptr), 
-      ptr_B_real(nullptr), 
-      ptr_B_imag(nullptr), 
-      ptr_C_real(nullptr), 
-      ptr_C_imag(nullptr), 
-      ptr_D_real(nullptr),
-      ptr_D_imag(nullptr)
-    {}
+
+    Arguments() = default;
 
     /// constructs an arguments structure
     Arguments(
@@ -248,29 +236,29 @@ public:
     // Data members
     //
 
-    typename Mma::IteratorA::Params params_A_real;
-    typename Mma::IteratorA::Params params_A_imag;
-    typename Mma::IteratorB::Params params_B_real;
-    typename Mma::IteratorB::Params params_B_imag;
-    typename Epilogue::OutputTileIterator::Params params_C_real;
-    typename Epilogue::OutputTileIterator::Params params_C_imag;
-    typename Epilogue::OutputTileIterator::Params params_D_real;
-    typename Epilogue::OutputTileIterator::Params params_D_imag;
+    typename Mma::IteratorA::Params params_A_real{};
+    typename Mma::IteratorA::Params params_A_imag{};
+    typename Mma::IteratorB::Params params_B_real{};
+    typename Mma::IteratorB::Params params_B_imag{};
+    typename Epilogue::OutputTileIterator::Params params_C_real{};
+    typename Epilogue::OutputTileIterator::Params params_C_imag{};
+    typename Epilogue::OutputTileIterator::Params params_D_real{};
+    typename Epilogue::OutputTileIterator::Params params_D_imag{};
 
-    typename EpilogueOutputOp::Params output_op;
+    typename EpilogueOutputOp::Params output_op{};
 
-    int const *ptr_M;
-    int const *ptr_N;
-    int const *ptr_K;
+    int const *ptr_M{nullptr};
+    int const *ptr_N{nullptr};
+    int const *ptr_K{nullptr};
 
-    void const * const * ptr_A_real;
-    void const * const * ptr_A_imag;
-    void const * const * ptr_B_real;
-    void const * const * ptr_B_imag;
-    void const * const * ptr_C_real;
-    void const * const * ptr_C_imag;
-    void * const * ptr_D_real;
-    void * const * ptr_D_imag;
+    void const * const * ptr_A_real{nullptr};
+    void const * const * ptr_A_imag{nullptr};
+    void const * const * ptr_B_real{nullptr};
+    void const * const * ptr_B_imag{nullptr};
+    void const * const * ptr_C_real{nullptr};
+    void const * const * ptr_C_imag{nullptr};
+    void * const * ptr_D_real{nullptr};
+    void * const * ptr_D_imag{nullptr};
 
     //
     // Host dispatch API

include/cutlass/gemm/kernel/gemm_streamk_with_fused_epilogue.h

@@ -29,7 +29,7 @@
  *
  **************************************************************************************************/
 /*! \file
-    \brief Stream-K Gemm kernel compatible with fused epilogues 
+    \brief Stream-K Gemm kernel compatible with fused epilogues
     that broadcast a bias vector over the MMA output.
 */
 
@@ -65,7 +65,7 @@ struct GemmStreamkWithFusedEpilogue;
 
 // GemmStreamkWithFusedEpilogue with two sources
 template <
-  typename Mma_,                  ///! Threadblock-scoped matrix multiply-accumulate 
+  typename Mma_,                  ///! Threadblock-scoped matrix multiply-accumulate
   typename Epilogue_,             ///! Epilogue
   typename ThreadblockSwizzle_    ///! Threadblock swizzling function
 >
@@ -126,55 +126,46 @@ struct GemmStreamkWithFusedEpilogue<Mma_, Epilogue_, ThreadblockSwizzle_, false>
     // Data members
     //
 
-    GemmUniversalMode mode;
-    GemmCoord problem_size;
-    int batch_count;        // Either (mode == GemmUniversalMode::kBatched) the batch count, or (mode == GemmUniversalMode::kGemm) the tile-splitting factor
+    GemmUniversalMode mode{GemmUniversalMode::kGemm};
+    GemmCoord problem_size{};
+    int batch_count{1};        // Either (mode == GemmUniversalMode::kBatched) the batch count, or (mode == GemmUniversalMode::kGemm) the tile-splitting factor
 
-    typename EpilogueOutputOp::Params epilogue;
+    typename EpilogueOutputOp::Params epilogue{};
 
-    void const * ptr_A;
-    void const * ptr_B;
-    void const * ptr_C1;
-    void const * ptr_C2;
-    void * ptr_D;
+    void const * ptr_A{nullptr};
+    void const * ptr_B{nullptr};
+    void const * ptr_C1{nullptr};
+    void const * ptr_C2{nullptr};
+    void * ptr_D{nullptr};
 
     void * ptr_Vector;
     void * ptr_Tensor;
 
-    int64_t batch_stride_A;
-    int64_t batch_stride_B;
-    int64_t batch_stride_C1;
-    int64_t batch_stride_C2;
-    int64_t batch_stride_D;
-    int64_t batch_stride_Vector;
-    int64_t batch_stride_Tensor;
+    int64_t batch_stride_A{0};
+    int64_t batch_stride_B{0};
+    int64_t batch_stride_C1{0};
+    int64_t batch_stride_C2{0};
+    int64_t batch_stride_D{0};
+    int64_t batch_stride_Vector{0};
+    int64_t batch_stride_Tensor{0};
 
-    typename LayoutA::Stride::Index lda;
-    typename LayoutB::Stride::Index ldb;
-    typename LayoutC::Stride::Index ldc1;
-    typename LayoutC::Stride::Index ldc2;
-    typename LayoutC::Stride::Index ldd;
-    typename LayoutC::Stride::Index ldr;
-    typename LayoutC::Stride::Index ldt;
+    typename LayoutA::Stride::Index lda{};
+    typename LayoutB::Stride::Index ldb{};
+    typename LayoutC::Stride::Index ldc1{};
+    typename LayoutC::Stride::Index ldc2{};
+    typename LayoutC::Stride::Index ldd{};
+    typename LayoutC::Stride::Index ldr{};
+    typename LayoutC::Stride::Index ldt{};
 
-    int avail_sms;          /// The number of SMs that StreamK dispatch heuristics will attempt to load-balance across (-1 defaults to device width, 1 implies classic data-parallel scheduling)
+    int avail_sms{-1};          /// The number of SMs that StreamK dispatch heuristics will attempt to load-balance across (-1 defaults to device width, 1 implies classic data-parallel scheduling)
 
 
     //
     // Methods
     //
-    
+
     /// Default Constructor
-    Arguments():
-      mode(GemmUniversalMode::kGemm),
-      batch_count(1),
-      ptr_A(nullptr),
-      ptr_B(nullptr),
-      ptr_C1(nullptr),
-      ptr_C2(nullptr),
-      ptr_D(nullptr),
-      avail_sms(-1)
-    {}
+    Arguments() = default;
 
     /// constructs an arguments structure
     Arguments(
@@ -208,14 +199,14 @@ struct GemmStreamkWithFusedEpilogue<Mma_, Epilogue_, ThreadblockSwizzle_, false>
       mode(mode),
       problem_size(problem_size),
       batch_count(batch_split),
-      epilogue(epilogue), 
-      ptr_A(ptr_A), ptr_B(ptr_B), ptr_C1(ptr_C1), ptr_C2(ptr_C2), ptr_D(ptr_D), 
-      ptr_Vector(ptr_Vector), 
+      epilogue(epilogue),
+      ptr_A(ptr_A), ptr_B(ptr_B), ptr_C1(ptr_C1), ptr_C2(ptr_C2), ptr_D(ptr_D),
+      ptr_Vector(ptr_Vector),
       ptr_Tensor(ptr_Tensor),
-      batch_stride_A(batch_stride_A), 
-      batch_stride_B(batch_stride_B), 
-      batch_stride_C1(batch_stride_C1), 
-      batch_stride_C2(batch_stride_C2), 
+      batch_stride_A(batch_stride_A),
+      batch_stride_B(batch_stride_B),
+      batch_stride_C1(batch_stride_C1),
+      batch_stride_C2(batch_stride_C2),
       batch_stride_Vector(batch_stride_Vector),
       batch_stride_Tensor(batch_stride_Tensor),
       lda(lda), ldb(ldb), ldc1(ldc1), ldc2(ldc2), ldd(ldd), ldr(ldr), ldt(ldt), avail_sms(avail_sms)
@@ -251,42 +242,42 @@ struct GemmStreamkWithFusedEpilogue<Mma_, Epilogue_, ThreadblockSwizzle_, false>
     // Data members
     //
 
-    void * ptr_A;
-    void * ptr_B;
+    void * ptr_A{nullptr};
+    void * ptr_B{nullptr};
 
-    typename Mma::IteratorA::Params params_A;
-    typename Mma::IteratorB::Params params_B;
+    typename Mma::IteratorA::Params params_A{};
+    typename Mma::IteratorB::Params params_B{};
 
-    int64_t batch_stride_A;
-    int64_t batch_stride_B;
+    int64_t batch_stride_A{0};
+    int64_t batch_stride_B{0};
 
-    GemmUniversalMode mode;
+    GemmUniversalMode mode{GemmUniversalMode::kGemm};
 
-    ThreadblockSwizzle block_mapping;
+    ThreadblockSwizzle block_mapping{};
 
-    void *barrier_workspace;
-    void *partials_workspace;
+    void *barrier_workspace{nullptr};
+    void *partials_workspace{nullptr};
 
-    typename EpilogueOutputOp::Params output_op;
+    typename EpilogueOutputOp::Params output_op{};
 
-    void * ptr_C1;
-    void * ptr_C2;
-    void * ptr_D;
-    void * ptr_Tensor;
-    void * ptr_Vector;
+    void * ptr_C1{nullptr};
+    void * ptr_C2{nullptr};
+    void * ptr_D{nullptr};
+    void * ptr_Tensor{nullptr};
+    void * ptr_Vector{nullptr};
 
-    typename Epilogue::OutputTileIterator::Params params_C1;
-    typename Epilogue::OutputTileIterator::Params params_C2;
-    typename Epilogue::OutputTileIterator::Params params_D;
-    typename Epilogue::TensorTileIterator::Params params_Tensor;
+    typename Epilogue::OutputTileIterator::Params params_C1{};
+    typename Epilogue::OutputTileIterator::Params params_C2{};
+    typename Epilogue::OutputTileIterator::Params params_D{};
+    typename Epilogue::TensorTileIterator::Params params_Tensor{};
 
-    int64_t batch_stride_C1;
-    int64_t batch_stride_C2;
-    int64_t batch_stride_D;
-    int64_t batch_stride_Vector;
-    int64_t batch_stride_Tensor;
+    int64_t batch_stride_C1{0};
+    int64_t batch_stride_C2{0};
+    int64_t batch_stride_D{0};
+    int64_t batch_stride_Vector{0};
+    int64_t batch_stride_Tensor{0};
 
-    typename LayoutC::Stride::Index ldr;
+    typename LayoutC::Stride::Index ldr{};
 
   protected:
 
@@ -361,17 +352,17 @@ struct GemmStreamkWithFusedEpilogue<Mma_, Epilogue_, ThreadblockSwizzle_, false>
       barrier_workspace(nullptr),
       partials_workspace(nullptr)
     {
-      CUTLASS_TRACE_HOST("GemmStreamkWithFusedEpilogue::Params::Params() - problem_size: " << problem_size);
+      CUTLASS_TRACE_HOST("GemmStreamkWithFusedEpilogue::Params::Params()");
       CUTLASS_TRACE_HOST("  ptr_Vector: " << (void *)this->ptr_Vector);
       CUTLASS_TRACE_HOST("  ptr_Tensor: " << (void *)this->ptr_Tensor);
       CUTLASS_TRACE_HOST("  ldr: " << this->ldr);
       CUTLASS_TRACE_HOST("  ldt: " << args.ldt);
-      CUTLASS_TRACE_HOST("  avail_sms: " << avail_sms);
 
       // Number of SMs to make available for StreamK decomposition
       int avail_sms = (args.avail_sms == -1) ?
                         device_sms :
                         fast_min(args.avail_sms, device_sms);
+      CUTLASS_TRACE_HOST("  avail_sms: " << avail_sms);
 
       // Initialize the block mapping structure
       block_mapping = ThreadblockSwizzle(
@@ -404,7 +395,6 @@ struct GemmStreamkWithFusedEpilogue<Mma_, Epilogue_, ThreadblockSwizzle_, false>
     {
       uint8_t *ptr = static_cast<uint8_t*>(workspace);
 
-
       // Establish partials workspace
       partials_workspace = nullptr;
       size_t partials_workspace_bytes = get_partials_workspace_size();
@@ -840,7 +830,7 @@ protected:
     typename Epilogue::ElementTensor *ptr_Tensor = static_cast<typename Epilogue::ElementTensor *>(params.ptr_Tensor);
 
     // Define the reduction output pointer and move to the appropriate place
-    typename Epilogue::ElementVector *ptr_Vector = 
+    typename Epilogue::ElementVector *ptr_Vector =
       static_cast<typename Epilogue::ElementVector *>(params.ptr_Vector);
 
     // Update pointers for batched/array mode(s)
@@ -969,7 +959,7 @@ protected:
     typename Epilogue::ElementTensor *ptr_Tensor = static_cast<typename Epilogue::ElementTensor *>(params.ptr_Tensor);
 
     // Define the reduction output pointer and move to the appropriate place
-    typename Epilogue::ElementVector *ptr_Vector = 
+    typename Epilogue::ElementVector *ptr_Vector =
       static_cast<typename Epilogue::ElementVector *>(params.ptr_Vector);
 
     // Tile iterator loading from residual1.
@@ -1256,7 +1246,7 @@ public:
 
 // GemmStreamkWithFusedEpilogue with one source
 template <
-  typename Mma_,                  ///! Threadblock-scoped matrix multiply-accumulate 
+  typename Mma_,                  ///! Threadblock-scoped matrix multiply-accumulate
   typename Epilogue_,             ///! Epilogue
   typename ThreadblockSwizzle_    ///! Threadblock swizzling function
 >
@@ -1318,51 +1308,43 @@ struct GemmStreamkWithFusedEpilogue<Mma_, Epilogue_, ThreadblockSwizzle_, true>
     // Data members
     //
 
-    GemmUniversalMode mode;
-    GemmCoord problem_size;
-    int batch_count;        // Either (mode == GemmUniversalMode::kBatched) the batch count, or (mode == GemmUniversalMode::kGemm) the tile-splitting factor
+    GemmUniversalMode mode{GemmUniversalMode::kGemm};
+    GemmCoord problem_size{};
+    int batch_count{1};        // Either (mode == GemmUniversalMode::kBatched) the batch count, or (mode == GemmUniversalMode::kGemm) the tile-splitting factor
 
-    typename EpilogueOutputOp::Params epilogue;
+    typename EpilogueOutputOp::Params epilogue{};
 
-    void const * ptr_A;
-    void const * ptr_B;
-    void const * ptr_C;
-    void * ptr_D;
+    void const * ptr_A{nullptr};
+    void const * ptr_B{nullptr};
+    void const * ptr_C{nullptr};
+    void * ptr_D{nullptr};
 
-    void * ptr_Vector;
-    void * ptr_Tensor;
+    void * ptr_Vector{nullptr};
+    void * ptr_Tensor{nullptr};
 
-    int64_t batch_stride_A;
-    int64_t batch_stride_B;
-    int64_t batch_stride_C;
-    int64_t batch_stride_D;
-    int64_t batch_stride_Vector;
-    int64_t batch_stride_Tensor;
+    int64_t batch_stride_A{0};
+    int64_t batch_stride_B{0};
+    int64_t batch_stride_C{0};
+    int64_t batch_stride_D{0};
+    int64_t batch_stride_Vector{0};
+    int64_t batch_stride_Tensor{0};
 
-    typename LayoutA::Stride::Index lda;
-    typename LayoutB::Stride::Index ldb;
-    typename LayoutC::Stride::Index ldc;
-    typename LayoutC::Stride::Index ldd;
-    typename LayoutC::Stride::Index ldr;
-    typename LayoutC::Stride::Index ldt;
+    typename LayoutA::Stride::Index lda{};
+    typename LayoutB::Stride::Index ldb{};
+    typename LayoutC::Stride::Index ldc{};
+    typename LayoutC::Stride::Index ldd{};
+    typename LayoutC::Stride::Index ldr{};
+    typename LayoutC::Stride::Index ldt{};
 
-    int avail_sms;          /// The number of SMs that StreamK dispatch heuristics will attempt to load-balance across (-1 defaults to device width, 1 implies classic data-parallel scheduling)
+    int avail_sms{-1};          /// The number of SMs that StreamK dispatch heuristics will attempt to load-balance across (-1 defaults to device width, 1 implies classic data-parallel scheduling)
 
 
     //
     // Methods
     //
-    
+
     /// Default Constructor
-    Arguments(): 
-      mode(GemmUniversalMode::kGemm),
-      batch_count(1),
-      ptr_A(nullptr),
-      ptr_B(nullptr),
-      ptr_C(nullptr),
-      ptr_D(nullptr),
-      avail_sms(-1)
-    {}
+    Arguments() = default;
 
     /// constructs an arguments structure
     Arguments(
@@ -1393,13 +1375,13 @@ struct GemmStreamkWithFusedEpilogue<Mma_, Epilogue_, ThreadblockSwizzle_, true>
       mode(mode),
       problem_size(problem_size),
       batch_count(batch_split),
-      epilogue(epilogue), 
-      ptr_A(ptr_A), ptr_B(ptr_B), ptr_C(ptr_C), ptr_D(ptr_D), 
-      ptr_Vector(ptr_Vector), 
+      epilogue(epilogue),
+      ptr_A(ptr_A), ptr_B(ptr_B), ptr_C(ptr_C), ptr_D(ptr_D),
+      ptr_Vector(ptr_Vector),
       ptr_Tensor(ptr_Tensor),
-      batch_stride_A(batch_stride_A), 
-      batch_stride_B(batch_stride_B), 
-      batch_stride_C(batch_stride_C), 
+      batch_stride_A(batch_stride_A),
+      batch_stride_B(batch_stride_B),
+      batch_stride_C(batch_stride_C),
       batch_stride_Vector(batch_stride_Vector),
       batch_stride_Tensor(batch_stride_Tensor),
       lda(lda), ldb(ldb), ldc(ldc), ldd(ldd), ldr(ldr), ldt(ldt), avail_sms(avail_sms)
@@ -1415,7 +1397,7 @@ struct GemmStreamkWithFusedEpilogue<Mma_, Epilogue_, ThreadblockSwizzle_, true>
     /// Returns arguments for the transposed problem
     Arguments transposed_problem() const {
       Arguments args(*this);
-      
+
       std::swap(args.problem_size.m(), args.problem_size.n());
       std::swap(args.ptr_A, args.ptr_B);
       std::swap(args.lda, args.ldb);
@@ -1436,40 +1418,39 @@ struct GemmStreamkWithFusedEpilogue<Mma_, Epilogue_, ThreadblockSwizzle_, true>
     // Data members
     //
 
-    void * ptr_A;
-    void * ptr_B;
+    void * ptr_A{nullptr};
+    void * ptr_B{nullptr};
 
-    typename Mma::IteratorA::Params params_A;
-    typename Mma::IteratorB::Params params_B;
+    typename Mma::IteratorA::Params params_A{};
+    typename Mma::IteratorB::Params params_B{};
 
-    int64_t batch_stride_A;
-    int64_t batch_stride_B;
+    int64_t batch_stride_A{0};
+    int64_t batch_stride_B{0};
 
-    GemmUniversalMode mode;
+    GemmUniversalMode mode{GemmUniversalMode::kGemm};
 
-    ThreadblockSwizzle block_mapping;
+    ThreadblockSwizzle block_mapping{};
 
-    void *barrier_workspace;
-    void *partials_workspace;
+    void *barrier_workspace{nullptr};
+    void *partials_workspace{nullptr};
 
-    typename EpilogueOutputOp::Params output_op;
+    typename EpilogueOutputOp::Params output_op{};
 
-    void * ptr_C;
-    void * ptr_D;
-    void * ptr_Tensor;
-    void * ptr_Vector;
+    void * ptr_C{nullptr};
+    void * ptr_D{nullptr};
+    void * ptr_Tensor{nullptr};
+    void * ptr_Vector{nullptr};
 
-    typename Epilogue::OutputTileIterator::Params params_C;
-    typename Epilogue::OutputTileIterator::Params params_D;
-    typename Epilogue::TensorTileIterator::Params params_Tensor;
+    typename Epilogue::OutputTileIterator::Params params_C{};
+    typename Epilogue::OutputTileIterator::Params params_D{};
+    typename Epilogue::TensorTileIterator::Params params_Tensor{};
 
-    int64_t batch_stride_C;
-    int64_t batch_stride_D;
-    int64_t batch_stride_Vector;
-    int64_t batch_stride_Tensor;
+    int64_t batch_stride_C{0};
+    int64_t batch_stride_D{0};
+    int64_t batch_stride_Vector{0};
+    int64_t batch_stride_Tensor{0};
 
-
-    typename LayoutC::Stride::Index ldr;
+    typename LayoutC::Stride::Index ldr{};
 
   protected:
 
@@ -1540,17 +1521,17 @@ struct GemmStreamkWithFusedEpilogue<Mma_, Epilogue_, ThreadblockSwizzle_, true>
       barrier_workspace(nullptr),
       partials_workspace(nullptr)
     {
-      CUTLASS_TRACE_HOST("GemmStreamkWithFusedEpilogue::Params::Params() - problem_size: " << problem_size);
+      CUTLASS_TRACE_HOST("GemmStreamkWithFusedEpilogue::Params::Params()");
       CUTLASS_TRACE_HOST("  ptr_Vector: " << (void *)this->ptr_Vector);
       CUTLASS_TRACE_HOST("  ptr_Tensor: " << (void *)this->ptr_Tensor);
       CUTLASS_TRACE_HOST("  ldr: " << this->ldr);
       CUTLASS_TRACE_HOST("  ldt: " << args.ldt);
-      CUTLASS_TRACE_HOST("  avail_sms: " << avail_sms);
 
       // Number of SMs to make available for StreamK decomposition
       int avail_sms = (args.avail_sms == -1) ?
                         device_sms :
                         fast_min(args.avail_sms, device_sms);
+      CUTLASS_TRACE_HOST("  avail_sms: " << avail_sms);
 
       // Initialize the block mapping structure
       block_mapping = ThreadblockSwizzle(
@@ -2018,7 +1999,7 @@ protected:
     typename Epilogue::ElementTensor *ptr_Tensor = static_cast<typename Epilogue::ElementTensor *>(params.ptr_Tensor);
 
     // Define the reduction output pointer and move to the appropriate place
-    typename Epilogue::ElementVector *ptr_Vector = 
+    typename Epilogue::ElementVector *ptr_Vector =
       static_cast<typename Epilogue::ElementVector *>(params.ptr_Vector);
 
     // Update pointers for batched/array mode(s)
@@ -2131,7 +2112,7 @@ protected:
     typename Epilogue::ElementTensor *ptr_Tensor = static_cast<typename Epilogue::ElementTensor *>(params.ptr_Tensor);
 
     // Define the reduction output pointer and move to the appropriate place
-    typename Epilogue::ElementVector *ptr_Vector = 
+    typename Epilogue::ElementVector *ptr_Vector =
       static_cast<typename Epilogue::ElementVector *>(params.ptr_Vector);
 
     // Tile iterator loading from source tensor.

include/cutlass/gemm/kernel/gemm_universal_streamk.h

@@ -122,33 +122,33 @@ public:
     // Data members
     //
 
-    GemmUniversalMode mode;
-    GemmCoord problem_size;
-    int batch_count;        // Either (mode == GemmUniversalMode::kBatched) the batch count, or (mode == GemmUniversalMode::kGemm) the tile-splitting factor
+    GemmUniversalMode mode = GemmUniversalMode::kGemm;
+    GemmCoord problem_size {};
+    int batch_count {1};        // Either (mode == GemmUniversalMode::kBatched) the batch count, or (mode == GemmUniversalMode::kGemm) the tile-splitting factor
 
-    typename EpilogueOutputOp::Params epilogue;
+    typename EpilogueOutputOp::Params epilogue{};
 
-    void const * ptr_A;
-    void const * ptr_B;
-    void const * ptr_C;
-    void * ptr_D;
+    void const * ptr_A = nullptr;
+    void const * ptr_B = nullptr;
+    void const * ptr_C = nullptr;
+    void * ptr_D = nullptr;
 
-    int64_t batch_stride_A;
-    int64_t batch_stride_B;
-    int64_t batch_stride_C;
-    int64_t batch_stride_D;
+    int64_t batch_stride_A{0};
+    int64_t batch_stride_B{0};
+    int64_t batch_stride_C{0};
+    int64_t batch_stride_D{0};
 
-    typename LayoutA::Stride stride_a;
-    typename LayoutB::Stride stride_b;
-    typename LayoutC::Stride stride_c;
-    typename LayoutC::Stride stride_d;
+    typename LayoutA::Stride stride_a{0};
+    typename LayoutB::Stride stride_b{0};
+    typename LayoutC::Stride stride_c{0};
+    typename LayoutC::Stride stride_d{0};
 
-    typename LayoutA::Stride::LongIndex lda;
-    typename LayoutB::Stride::LongIndex ldb;
-    typename LayoutC::Stride::LongIndex ldc;
-    typename LayoutC::Stride::LongIndex ldd;
+    typename LayoutA::Stride::LongIndex lda{0};
+    typename LayoutB::Stride::LongIndex ldb{0};
+    typename LayoutC::Stride::LongIndex ldc{0};
+    typename LayoutC::Stride::LongIndex ldd{0};
 
-    int avail_sms;          /// The number of SMs that StreamK dispatch heuristics will attempt to load-balance across (-1 defaults to device width, 1 implies classic data-parallel scheduling)
+    int avail_sms{-1};          /// The number of SMs that StreamK dispatch heuristics will attempt to load-balance across (-1 defaults to device width, 1 implies classic data-parallel scheduling)
 
 
     //
@@ -156,15 +156,7 @@ public:
     //
 
     /// Default Constructor
-    Arguments():
-      mode(GemmUniversalMode::kGemm),
-      batch_count(1),
-      ptr_A(nullptr),
-      ptr_B(nullptr),
-      ptr_C(nullptr),
-      ptr_D(nullptr),
-      avail_sms(-1)
-    {}
+    Arguments() = default;
 
     /// Constructor
     Arguments(
@@ -257,32 +249,32 @@ public:
     // Data members
     //
 
-    void * ptr_A;
-    void * ptr_B;
+    void * ptr_A = nullptr;
+    void * ptr_B = nullptr;
 
-    typename Mma::IteratorA::Params params_A;
-    typename Mma::IteratorB::Params params_B;
+    typename Mma::IteratorA::Params params_A{};
+    typename Mma::IteratorB::Params params_B{};
 
-    int64_t batch_stride_A;
-    int64_t batch_stride_B;
+    int64_t batch_stride_A{0};
+    int64_t batch_stride_B{0};
 
-    GemmUniversalMode mode;
+    GemmUniversalMode mode = GemmUniversalMode::kGemm;
 
-    ThreadblockSwizzle block_mapping;
+    ThreadblockSwizzle block_mapping{};
 
-    void *barrier_workspace;
-    void *partials_workspace;
+    void *barrier_workspace = nullptr;
+    void *partials_workspace = nullptr;
 
-    typename EpilogueOutputOp::Params output_op;
+    typename EpilogueOutputOp::Params output_op{};
 
-    void * ptr_D;
-    void * ptr_C;
+    void * ptr_D = nullptr;
+    void * ptr_C = nullptr;
 
-    typename Epilogue::OutputTileIterator::Params params_D;
-    typename Epilogue::OutputTileIterator::Params params_C;
+    typename Epilogue::OutputTileIterator::Params params_D{};
+    typename Epilogue::OutputTileIterator::Params params_C{};
 
-    int64_t batch_stride_D;
-    int64_t batch_stride_C;
+    int64_t batch_stride_D{0};
+    int64_t batch_stride_C{0};
 
 
   protected:
@@ -326,7 +318,6 @@ public:
     /// Default constructor
     Params() = default;
 
-
     /// Constructor
     Params(
       Arguments const &args,  /// GEMM application arguments

include/cutlass/gemm/kernel/gemm_universal_with_visitor_streamk.h

@@ -129,35 +129,35 @@ public:
     //
     // Data members
     //
-    cute::Shape<int32_t,int32_t,int32_t> problem_shape;
+    cute::Shape<int32_t,int32_t,int32_t> problem_shape{};
 
-    void * ptr_A;
-    void * ptr_B;
+    void * ptr_A{nullptr};
+    void * ptr_B{nullptr};
 
-    typename Mma::IteratorA::Params params_A;
-    typename Mma::IteratorB::Params params_B;
+    typename Mma::IteratorA::Params params_A{};
+    typename Mma::IteratorB::Params params_B{};
 
-    int64_t batch_stride_A;
-    int64_t batch_stride_B;
+    int64_t batch_stride_A{0};
+    int64_t batch_stride_B{0};
 
-    GemmUniversalMode mode;
+    GemmUniversalMode mode{GemmUniversalMode::kGemm};
 
-    ThreadblockSwizzle block_mapping;
+    ThreadblockSwizzle block_mapping{};
 
-    void *barrier_workspace;
-    void *partials_workspace;
+    void *barrier_workspace{nullptr};
+    void *partials_workspace{nullptr};
 
-    typename FusionCallbacks::Params output_op;
+    typename FusionCallbacks::Params output_op{};
 
 
-    void * ptr_D;
-    void * ptr_C;
+    void * ptr_D{nullptr};
+    void * ptr_C{nullptr};
 
-    typename Epilogue::OutputTileIterator::Params params_D;
-    typename Epilogue::OutputTileIterator::Params params_C;
+    typename Epilogue::OutputTileIterator::Params params_D{};
+    typename Epilogue::OutputTileIterator::Params params_C{};
 
-    int64_t batch_stride_D;
-    int64_t batch_stride_C;
+    int64_t batch_stride_D{0};
+    int64_t batch_stride_C{0};
 
 
   protected:

include/cutlass/gemm/kernel/gemm_with_absmax.h

@@ -0,0 +1,759 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+/*! \file
+    \brief Gemm kernel with an epilogue that computes the absolute maximum value of the output
+    and a pre-activation-function auxiliary output. The auxiliary output is also (optionally)
+    stored to global memory.
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/fast_math.h"
+#include "cutlass/layout/layout.h"
+#include "cutlass/gemm/gemm.h"
+#include "cutlass/matrix_coord.h"
+#include "cutlass/complex.h"
+#include "cutlass/semaphore.h"
+#include "cutlass/gemm/kernel/params_universal_base.h"
+
+#include "cutlass/trace.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace kernel {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Gemm that computes the absolute maximum value of the output and a pre-activation-function
+// auxiliary output.
+template <
+  typename Mma_,                  ///! Threadblock-scoped matrix multiply-accumulate
+  typename Epilogue_,             ///! Epilogue
+  typename ThreadblockSwizzle_    ///! Threadblock swizzling function
+>
+struct GemmWithAbsMax {
+public:
+
+  using Mma = Mma_;
+  using Epilogue = Epilogue_;
+  using EpilogueOutputOp = typename Epilogue::OutputOp;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
+
+  using ElementA = typename Mma::IteratorA::Element;
+  using LayoutA = typename Mma::IteratorA::Layout;
+  using ElementB = typename Mma::IteratorB::Element;
+  using LayoutB = typename Mma::IteratorB::Layout;
+  using ElementC = typename Epilogue::OutputTileIterator::Element;
+  using LayoutC = typename Epilogue::OutputTileIterator::Layout;
+
+  static ComplexTransform const kTransformA = Mma::kTransformA;
+  static ComplexTransform const kTransformB = Mma::kTransformB;
+  using Operator = typename Mma::Operator;
+
+  using OperatorClass = typename Mma::Operator::OperatorClass;
+  using ThreadblockShape = typename Mma::Shape;
+  using WarpShape = typename Mma::Operator::Shape;
+  using InstructionShape = typename Mma::Policy::Operator::InstructionShape;
+  using ArchTag = typename Mma::ArchTag;
+
+  static int const kStages = Mma::kStages;
+  static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
+  static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
+  static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess;
+
+  /// Warp count (concept: GemmShape)
+  using WarpCount = typename Mma::WarpCount;
+  static int const kThreadCount = 32 * WarpCount::kCount;
+
+  /// Split-K preserves splits that are 128b aligned
+  static int const kSplitKAlignment = const_max(
+    128 / sizeof_bits<ElementA>::value,
+    128 / sizeof_bits<ElementB>::value
+  );
+
+  //
+  // Structures
+  //
+
+  /// Argument structure
+  struct Arguments : UniversalArgumentsBase
+  {
+    //
+    // Data members
+    //
+
+    typename EpilogueOutputOp::Params epilogue;
+
+    void const * ptr_A;
+    void const * ptr_B;
+    void const * ptr_C;
+    void * ptr_D;
+    void * ptr_Aux;
+
+    void * ptr_Vector;
+
+    int64_t batch_stride_A;
+    int64_t batch_stride_B;
+    int64_t batch_stride_C;
+    int64_t batch_stride_Vector;
+
+    typename LayoutA::Stride::Index lda;
+    typename LayoutB::Stride::Index ldb;
+    typename LayoutC::Stride::Index ldc;
+    typename LayoutC::Stride::Index ldd;
+    typename LayoutC::Stride::Index ldaux;
+    typename LayoutC::Stride::Index ldr;
+
+    //
+    // Methods
+    //
+
+    Arguments():
+      ptr_A(nullptr),
+      ptr_B(nullptr),
+      ptr_C(nullptr),
+      ptr_D(nullptr),
+      ptr_Aux(nullptr)
+    {}
+
+    /// Constructs an arguments structure with ldaux
+    Arguments(
+      GemmUniversalMode mode,
+      GemmCoord problem_size,
+      int batch_count,
+      typename EpilogueOutputOp::Params epilogue,
+      void const * ptr_A,
+      void const * ptr_B,
+      void const * ptr_C,
+      void * ptr_D,
+      void * ptr_Aux,
+      void * ptr_Vector,
+      int64_t batch_stride_A,
+      int64_t batch_stride_B,
+      int64_t batch_stride_C,
+      int64_t batch_stride_D,
+      int64_t batch_stride_Vector,
+      typename LayoutA::Stride::Index lda,
+      typename LayoutB::Stride::Index ldb,
+      typename LayoutC::Stride::Index ldc,
+      typename LayoutC::Stride::Index ldd,
+      typename LayoutC::Stride::Index ldr,
+      typename LayoutC::Stride::Index ldaux)
+    :
+      UniversalArgumentsBase(mode, problem_size, batch_count, batch_stride_D),
+      epilogue(epilogue),
+      ptr_A(ptr_A), ptr_B(ptr_B), ptr_C(ptr_C), ptr_D(ptr_D), ptr_Aux(ptr_Aux),
+      ptr_Vector(ptr_Vector),
+      batch_stride_A(batch_stride_A),
+      batch_stride_B(batch_stride_B),
+      batch_stride_C(batch_stride_C),
+      batch_stride_Vector(batch_stride_Vector),
+      lda(lda), ldb(ldb), ldc(ldc), ldd(ldd), ldaux(ldaux), ldr(ldr)
+    {
+    }
+
+    /// Constructs an Arguments structure without ldaux.
+    /// These parameters are overridden with D batch stride and ldd.
+    Arguments(
+      GemmUniversalMode mode,
+      GemmCoord problem_size,
+      int batch_count,
+      typename EpilogueOutputOp::Params epilogue,
+      void const * ptr_A,
+      void const * ptr_B,
+      void const * ptr_C,
+      void * ptr_D,
+      void * ptr_Aux,
+      void * ptr_Vector,
+      int64_t batch_stride_A,
+      int64_t batch_stride_B,
+      int64_t batch_stride_C,
+      int64_t batch_stride_D,
+      int64_t batch_stride_Vector,
+      typename LayoutA::Stride::Index lda,
+      typename LayoutB::Stride::Index ldb,
+      typename LayoutC::Stride::Index ldc,
+      typename LayoutC::Stride::Index ldd,
+      typename LayoutC::Stride::Index ldr)
+    : Arguments(mode, problem_size, batch_count, epilogue, ptr_A, ptr_B, ptr_C, ptr_D, ptr_Aux, ptr_Vector,
+               batch_stride_A, batch_stride_B, batch_stride_C, batch_stride_D, batch_stride_Vector,
+               lda, ldb, ldc, ldd, ldr, ldd)
+    {
+    }
+
+    /// Returns arguments for the transposed problem
+    Arguments transposed_problem() const {
+      Arguments args(*this);
+
+      std::swap(args.problem_size.m(), args.problem_size.n());
+      std::swap(args.ptr_A, args.ptr_B);
+      std::swap(args.lda, args.ldb);
+      std::swap(args.batch_stride_A, args.batch_stride_B);
+
+      return args;
+    }
+  };
+
+
+  //
+  // Structure for precomputing values in host memory and passing to kernels
+  //
+
+  /// Parameters structure
+  struct Params : UniversalParamsBase<
+    ThreadblockSwizzle,
+    ThreadblockShape,
+    ElementA,
+    ElementB,
+    ElementC,
+    LayoutA,
+    LayoutB>
+  {
+    using ParamsBase = UniversalParamsBase<
+      ThreadblockSwizzle,
+      ThreadblockShape,
+      ElementA,
+      ElementB,
+      ElementC,
+      LayoutA,
+      LayoutB>;
+
+    //
+    // Data members
+    //
+
+    typename Mma::IteratorA::Params params_A;
+    typename Mma::IteratorB::Params params_B;
+    typename Epilogue::OutputTileIterator::Params params_C;
+    typename Epilogue::OutputTileIterator::Params params_D;
+    typename Epilogue::AuxOutputTileIterator::Params params_Aux;
+
+    typename EpilogueOutputOp::Params output_op;
+
+    void * ptr_A;
+    void * ptr_B;
+    void * ptr_C;
+    void * ptr_D;
+    void * ptr_Aux;
+
+    void * ptr_Vector;
+    typename LayoutC::Stride::Index ldr;
+
+    int64_t batch_stride_A;
+    int64_t batch_stride_B;
+    int64_t batch_stride_C;
+    int64_t batch_stride_Vector;
+
+    //
+    // Host dispatch API
+    //
+
+    /// Default constructor
+    Params() = default;
+
+    /// Constructor
+    Params(
+      Arguments const &args,  /// GEMM application arguments
+      int device_sms,         /// Number of SMs on the device
+      int sm_occupancy)       /// Kernel SM occupancy (in thread blocks)
+    :
+      ParamsBase(args, device_sms, sm_occupancy),
+      params_A(args.lda),
+      params_B(args.ldb),
+      params_C(args.ldc),
+      params_D(args.ldd),
+      params_Aux(args.ldaux),
+      output_op(args.epilogue),
+      ptr_A(const_cast<void *>(args.ptr_A)),
+      ptr_B(const_cast<void *>(args.ptr_B)),
+      ptr_C(const_cast<void *>(args.ptr_C)),
+      ptr_D(args.ptr_D),
+      ptr_Aux(args.ptr_Aux),
+      ptr_Vector(args.ptr_Vector),
+      ldr(args.ldr),
+      batch_stride_A(args.batch_stride_A),
+      batch_stride_B(args.batch_stride_B),
+      batch_stride_C(args.batch_stride_C),
+      batch_stride_Vector(args.batch_stride_Vector)
+    {
+
+    }
+
+    /// Lightweight update given a subset of arguments.
+    CUTLASS_HOST_DEVICE
+    void update(Arguments const &args)
+    {
+      ptr_A = const_cast<void *>(args.ptr_A);
+      ptr_B = const_cast<void *>(args.ptr_B);
+      ptr_C = const_cast<void *>(args.ptr_C);
+      ptr_D = args.ptr_D;
+      ptr_Aux = args.ptr_Aux;
+
+      ptr_Vector = args.ptr_Vector;
+      ldr = args.ldr;
+
+      batch_stride_A = args.batch_stride_A;
+      batch_stride_B = args.batch_stride_B;
+      batch_stride_C = args.batch_stride_C;
+      this->batch_stride_D = args.batch_stride_D;
+      batch_stride_Vector = args.batch_stride_Vector;
+
+      output_op = args.epilogue;
+    }
+  };
+
+
+  /// Shared memory storage structure
+  union SharedStorage {
+    typename Mma::SharedStorage main_loop;
+    typename Epilogue::SharedStorage epilogue;
+  };
+
+public:
+
+  //
+  // Host dispatch API
+  //
+
+  /// Determines whether kernel satisfies alignment
+  static Status can_implement(
+    cutlass::gemm::GemmCoord const & problem_size) {
+
+    static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
+    static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
+    static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess;
+
+    bool isAMisaligned = false;
+    bool isBMisaligned = false;
+    bool isCMisaligned = false;
+
+    if (platform::is_same<LayoutA, layout::RowMajor>::value) {
+      isAMisaligned = problem_size.k() % kAlignmentA;
+    } else if (platform::is_same<LayoutA, layout::ColumnMajor>::value) {
+      isAMisaligned = problem_size.m() % kAlignmentA;
+    } else if (platform::is_same<LayoutA, layout::ColumnMajorInterleaved<32>>::value
+            || platform::is_same<LayoutA, layout::ColumnMajorInterleaved<64>>::value) {
+      isAMisaligned = problem_size.k() % kAlignmentA;
+    }
+
+    if (platform::is_same<LayoutB, layout::RowMajor>::value) {
+      isBMisaligned = problem_size.n() % kAlignmentB;
+    } else if (platform::is_same<LayoutB, layout::ColumnMajor>::value) {
+      isBMisaligned = problem_size.k() % kAlignmentB;
+    } else if (platform::is_same<LayoutB, layout::RowMajorInterleaved<32>>::value
+            || platform::is_same<LayoutB, layout::RowMajorInterleaved<64>>::value) {
+      isBMisaligned = problem_size.k() % kAlignmentB;
+    }
+
+    if (platform::is_same<LayoutC, layout::RowMajor>::value) {
+      isCMisaligned = problem_size.n() % kAlignmentC;
+    } else if (platform::is_same<LayoutC, layout::ColumnMajor>::value) {
+      isCMisaligned = problem_size.m() % kAlignmentC;
+    } else if (platform::is_same<LayoutC, layout::ColumnMajorInterleaved<32>>::value
+            || platform::is_same<LayoutC, layout::ColumnMajorInterleaved<64>>::value) {
+      isCMisaligned = problem_size.n() % kAlignmentC;
+    }
+
+    if (isAMisaligned) {
+      CUTLASS_TRACE_HOST("  returning kErrorMisalignedOperand for A operand");
+      return Status::kErrorMisalignedOperand;
+    }
+
+    if (isBMisaligned) {
+      CUTLASS_TRACE_HOST("  returning kErrorMisalignedOperand for B operand");
+      return Status::kErrorMisalignedOperand;
+    }
+
+    if (isCMisaligned) {
+      CUTLASS_TRACE_HOST("  returning kErrorMisalignedOperand for C operand");
+      return Status::kErrorMisalignedOperand;
+    }
+
+    CUTLASS_TRACE_HOST("  returning kSuccess");
+
+    return Status::kSuccess;
+  }
+
+  static Status can_implement(Arguments const &args) {
+    return can_implement(args.problem_size);
+  }
+
+public:
+
+  //
+  // Device-only API
+  //
+
+  // Factory invocation
+  CUTLASS_DEVICE
+  static void invoke(
+    Params const &params,
+    SharedStorage &shared_storage)
+  {
+    GemmWithAbsMax op;
+    op(params, shared_storage);
+  }
+
+  /// Executes one GEMM
+  CUTLASS_DEVICE
+  void operator()(Params const &params, SharedStorage &shared_storage) {
+
+    // Compute threadblock location
+    ThreadblockSwizzle threadblock_swizzle;
+
+    cutlass::gemm::GemmCoord threadblock_tile_offset = threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
+
+    // Early exit if CTA is out of range
+    if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() ||
+      params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) {
+
+      return;
+    }
+
+    int offset_k = 0;
+    int problem_size_k = params.problem_size.k();
+
+    ElementA *ptr_A = static_cast<ElementA *>(params.ptr_A);
+    ElementB *ptr_B = static_cast<ElementB *>(params.ptr_B);
+
+    //
+    // Fetch pointers based on mode.
+    //
+    if (params.mode == GemmUniversalMode::kGemm ||
+      params.mode == GemmUniversalMode::kGemmSplitKParallel) {
+
+      if (threadblock_tile_offset.k() + 1 < params.grid_tiled_shape.k()) {
+
+        problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size;
+      }
+
+      offset_k = threadblock_tile_offset.k() * params.gemm_k_size;
+    }
+    else if (params.mode == GemmUniversalMode::kBatched) {
+      ptr_A += threadblock_tile_offset.k() * params.batch_stride_A;
+      ptr_B += threadblock_tile_offset.k() * params.batch_stride_B;
+    }
+    else if (params.mode == GemmUniversalMode::kArray) {
+      ptr_A = static_cast<ElementA * const *>(params.ptr_A)[threadblock_tile_offset.k()];
+      ptr_B = static_cast<ElementB * const *>(params.ptr_B)[threadblock_tile_offset.k()];
+    }
+
+    __syncthreads();
+
+    // Compute initial location in logical coordinates
+    cutlass::MatrixCoord tb_offset_A{
+      threadblock_tile_offset.m() * Mma::Shape::kM,
+      offset_k,
+    };
+
+    cutlass::MatrixCoord tb_offset_B{
+      offset_k,
+      threadblock_tile_offset.n() * Mma::Shape::kN
+    };
+
+    // Compute position within threadblock
+    int thread_idx = threadIdx.x;
+
+    // Construct iterators to A and B operands
+    typename Mma::IteratorA iterator_A(
+      params.params_A,
+      ptr_A,
+      {params.problem_size.m(), problem_size_k},
+      thread_idx,
+      tb_offset_A);
+
+    typename Mma::IteratorB iterator_B(
+      params.params_B,
+      ptr_B,
+      {problem_size_k, params.problem_size.n()},
+      thread_idx,
+      tb_offset_B);
+
+    // Broadcast the warp_id computed by lane 0 to ensure dependent code
+    // is compiled as warp-uniform.
+    int warp_idx = canonical_warp_idx_sync();
+
+    int lane_idx = threadIdx.x % 32;
+
+    //
+    // Main loop
+    //
+
+    // Construct thread-scoped matrix multiply
+    Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx);
+
+    typename Mma::FragmentC accumulators;
+
+    accumulators.clear();
+
+    // Compute threadblock-scoped matrix multiply-add
+    int gemm_k_iterations = (problem_size_k - offset_k + Mma::Shape::kK - 1) / Mma::Shape::kK;
+
+    // Compute threadblock-scoped matrix multiply-add
+    mma(
+      gemm_k_iterations,
+      accumulators,
+      iterator_A,
+      iterator_B,
+      accumulators);
+
+    //
+    // Epilogue
+    //
+
+    EpilogueOutputOp output_op(params.output_op);
+
+    //
+    // Masked tile iterators constructed from members
+    //
+
+    threadblock_tile_offset = threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
+
+    //assume identity swizzle
+    MatrixCoord threadblock_offset(
+      threadblock_tile_offset.m() * Mma::Shape::kM,
+      threadblock_tile_offset.n() * Mma::Shape::kN
+    );
+
+    int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * params.grid_tiled_shape.m();
+
+    ElementC *ptr_C = static_cast<ElementC *>(params.ptr_C);
+    ElementC *ptr_D = static_cast<ElementC *>(params.ptr_D);
+    typename Epilogue::ElementAuxOutput *ptr_Aux = static_cast<typename Epilogue::ElementAuxOutput *>(params.ptr_Aux);
+    typename Epilogue::ElementVector *ptr_Vector = static_cast<typename Epilogue::ElementVector *>(params.ptr_Vector);
+
+    //
+    // Fetch pointers based on mode.
+    //
+
+    //
+    // Special path when split-K not enabled.
+    //
+
+    if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() == 1) {
+
+      // Tile iterators loading from source tensors.
+      typename Epilogue::OutputTileIterator iterator_C(
+        params.params_C,
+        ptr_C,
+        params.problem_size.mn(),
+        thread_idx,
+        threadblock_offset
+      );
+
+      // Tile iterator writing to destination tensor.
+      typename Epilogue::OutputTileIterator iterator_D(
+        params.params_D,
+        ptr_D,
+        params.problem_size.mn(),
+        thread_idx,
+        threadblock_offset
+      );
+
+      // Tile iterator writing to auxiliary tensor.
+      typename Epilogue::AuxOutputTileIterator iterator_Aux(
+        params.params_Aux,
+        ptr_Aux,
+        params.problem_size.mn(),
+        thread_idx,
+        threadblock_offset
+      );
+
+      // Construct the epilogue
+      Epilogue epilogue(
+        shared_storage.epilogue,
+        thread_idx,
+        warp_idx,
+        lane_idx);
+
+      // Move to appropriate location for this output tile
+      if (ptr_Vector) {
+        ptr_Vector += threadblock_offset.column() + threadblock_tile_offset.m() * params.ldr;
+      }
+
+      // Execute the epilogue operator to update the destination tensor.
+      epilogue(output_op,
+               ptr_Vector,
+               iterator_D,
+               accumulators,
+               iterator_C,
+               iterator_Aux,
+               params.problem_size.mn(),
+               threadblock_offset);
+
+      return;
+    }
+
+    //
+    // Slower path when split-K or batching is needed
+    //
+
+    // Construct the semaphore.
+    Semaphore semaphore(params.semaphore + block_idx, thread_idx);
+
+    if (params.mode == GemmUniversalMode::kGemm) {
+
+      // If performing a reduction via split-K, fetch the initial synchronization
+      if (params.grid_tiled_shape.k() > 1) {
+
+        // Fetch the synchronization lock initially but do not block.
+        semaphore.fetch();
+
+        // Indicate which position in a serial reduction the output operator is currently updating
+        output_op.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k());
+      }
+    }
+    else if (params.mode == GemmUniversalMode::kGemmSplitKParallel) {
+      ptr_D += threadblock_tile_offset.k() * params.batch_stride_D;
+    }
+    else if (params.mode == GemmUniversalMode::kBatched) {
+      ptr_C += threadblock_tile_offset.k() * params.batch_stride_C;
+      ptr_D += threadblock_tile_offset.k() * params.batch_stride_D;
+      if (ptr_Aux) {
+        ptr_Aux += threadblock_tile_offset.k() * params.batch_stride_D;
+      }
+      if (ptr_Vector) {
+        ptr_Vector += threadblock_tile_offset.k() * params.batch_stride_Vector;
+      }
+    }
+    else if (params.mode == GemmUniversalMode::kArray) {
+      ptr_C = static_cast<ElementC * const *>(params.ptr_C)[threadblock_tile_offset.k()];
+      ptr_D = static_cast<ElementC * const *>(params.ptr_D)[threadblock_tile_offset.k()];
+      if (ptr_Aux) {
+        ptr_Aux = static_cast<typename Epilogue::ElementAuxOutput * const *>(params.ptr_Aux)[threadblock_tile_offset.k()];
+      }
+      if (ptr_Vector) {
+        ptr_Vector = static_cast<typename Epilogue::ElementVector * const *>(params.ptr_Vector)[threadblock_tile_offset.k()];
+      }
+    }
+
+    // Tile iterators loading from source tensors.
+    typename Epilogue::OutputTileIterator iterator_C(
+      params.params_C,
+      ptr_C,
+      params.problem_size.mn(),
+      thread_idx,
+      threadblock_offset
+    );
+
+    // Tile iterator writing to destination tensor.
+    typename Epilogue::OutputTileIterator iterator_D(
+      params.params_D,
+      ptr_D,
+      params.problem_size.mn(),
+      thread_idx,
+      threadblock_offset
+    );
+
+    // Tile iterator writing to auxiliary destination tensor.
+    typename Epilogue::AuxOutputTileIterator iterator_Aux(
+      params.params_Aux,
+      // Only the final block writes the auxiliary tensor
+      ((params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) &&
+          (params.grid_tiled_shape.k() != threadblock_tile_offset.k() + 1))
+          ? nullptr
+          : ptr_Aux,
+      params.problem_size.mn(),
+      thread_idx,
+      threadblock_offset
+    );
+
+    // Construct the epilogue
+    Epilogue epilogue(
+      shared_storage.epilogue,
+      thread_idx,
+      warp_idx,
+      lane_idx);
+
+    // Wait on the semaphore - this latency may have been covered by iterator construction
+    if ((params.mode == GemmUniversalMode::kGemm) && params.grid_tiled_shape.k() > 1) {
+
+      // For subsequent threadblocks, the source matrix is held in the 'D' tensor.
+      if (threadblock_tile_offset.k()) {
+        iterator_C = iterator_D;
+      }
+
+      semaphore.wait(threadblock_tile_offset.k());
+
+    }
+
+    // Move to appropriate location for this output tile
+    if (ptr_Vector) {
+      ptr_Vector += threadblock_offset.column() + threadblock_tile_offset.m() * params.ldr;
+    }
+
+    // Execute the epilogue operator to update the destination tensor.
+    epilogue(output_op,
+             // Only the final block uses Vector
+             ((params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() > 1) &&
+              (params.grid_tiled_shape.k() != threadblock_tile_offset.k() + 1))
+                 ? nullptr
+                 : ptr_Vector,
+             iterator_D,
+             accumulators,
+             iterator_C,
+             iterator_Aux,
+             params.problem_size.mn(),
+             threadblock_offset);
+
+    //
+    // Release the semaphore
+    //
+
+    if ((params.mode == GemmUniversalMode::kGemm)  && params.grid_tiled_shape.k() > 1) {
+
+      int lock = 0;
+      if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) {
+
+        // The final threadblock resets the semaphore for subsequent grids.
+        lock = 0;
+      }
+      else {
+        // Otherwise, the semaphore is incremented
+        lock = threadblock_tile_offset.k() + 1;
+      }
+
+      semaphore.release(lock);
+    }
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace kernel
+} // namespace gemm
+} // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

include/cutlass/gemm/kernel/gemm_with_fused_epilogue.h

@@ -54,7 +54,7 @@ namespace kernel {
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 template <
-  typename Mma_,                  ///! Threadblock-scoped matrix multiply-accumulate 
+  typename Mma_,                  ///! Threadblock-scoped matrix multiply-accumulate
   typename Epilogue_,             ///! Epilogue
   typename ThreadblockSwizzle_,   ///! Threadblock swizzling function
   bool IsSingleSource = Epilogue_::kIsSingleSource
@@ -63,7 +63,7 @@ struct GemmWithFusedEpilogue;
 
 // GemmWithFusedEpilogue with two sources
 template <
-  typename Mma_,                  ///! Threadblock-scoped matrix multiply-accumulate 
+  typename Mma_,                  ///! Threadblock-scoped matrix multiply-accumulate
   typename Epilogue_,             ///! Epilogue
   typename ThreadblockSwizzle_    ///! Threadblock swizzling function
 >
@@ -103,7 +103,7 @@ public:
 
   /// Split-K preserves splits that are 128b aligned
   static int const kSplitKAlignment = const_max(
-    128 / sizeof_bits<ElementA>::value, 
+    128 / sizeof_bits<ElementA>::value,
     128 / sizeof_bits<ElementB>::value
   );
 
@@ -147,8 +147,8 @@ public:
     //
     // Methods
     //
-    
-    Arguments(): 
+
+    Arguments():
       ptr_A(nullptr),
       ptr_B(nullptr),
       ptr_C1(nullptr),
@@ -185,14 +185,14 @@ public:
       typename LayoutC::Stride::Index ldt)
     :
       UniversalArgumentsBase(mode, problem_size, batch_count, batch_stride_D),
-      epilogue(epilogue), 
-      ptr_A(ptr_A), ptr_B(ptr_B), ptr_C1(ptr_C1), ptr_C2(ptr_C2), ptr_D(ptr_D), 
-      ptr_Vector(ptr_Vector), 
+      epilogue(epilogue),
+      ptr_A(ptr_A), ptr_B(ptr_B), ptr_C1(ptr_C1), ptr_C2(ptr_C2), ptr_D(ptr_D),
+      ptr_Vector(ptr_Vector),
       ptr_Tensor(ptr_Tensor),
-      batch_stride_A(batch_stride_A), 
-      batch_stride_B(batch_stride_B), 
-      batch_stride_C1(batch_stride_C1), 
-      batch_stride_C2(batch_stride_C2), 
+      batch_stride_A(batch_stride_A),
+      batch_stride_B(batch_stride_B),
+      batch_stride_C1(batch_stride_C1),
+      batch_stride_C2(batch_stride_C2),
       batch_stride_Vector(batch_stride_Vector),
       batch_stride_Tensor(batch_stride_Tensor),
       lda(lda), ldb(ldb), ldc1(ldc1), ldc2(ldc2), ldd(ldd), ldr(ldr), ldt(ldt)
@@ -207,7 +207,7 @@ public:
     /// Returns arguments for the transposed problem
     Arguments transposed_problem() const {
       Arguments args(*this);
-      
+
       std::swap(args.problem_size.m(), args.problem_size.n());
       std::swap(args.ptr_A, args.ptr_B);
       std::swap(args.lda, args.ldb);
@@ -307,7 +307,7 @@ public:
       batch_stride_Vector(args.batch_stride_Vector),
       batch_stride_Tensor(args.batch_stride_Tensor)
     {
-      CUTLASS_TRACE_HOST("GemmWithFusedEpilogue::Params::Params() - problem_size: " << problem_size);
+      CUTLASS_TRACE_HOST("GemmWithFusedEpilogue::Params::Params()");
       CUTLASS_TRACE_HOST("  ptr_Vector: " << (void *)this->ptr_Vector);
       CUTLASS_TRACE_HOST("  ptr_Tensor: " << (void *)this->ptr_Tensor);
       CUTLASS_TRACE_HOST("  ldr: " << this->ldr);
@@ -460,7 +460,7 @@ public:
     int offset_k = 0;
     int problem_size_k = params.problem_size.k();
 
-    ElementA *ptr_A = static_cast<ElementA *>(params.ptr_A); 
+    ElementA *ptr_A = static_cast<ElementA *>(params.ptr_A);
     ElementB *ptr_B = static_cast<ElementB *>(params.ptr_B);
 
 
@@ -468,12 +468,12 @@ public:
     //
     // Fetch pointers based on mode.
     //
-    if (params.mode == GemmUniversalMode::kGemm || 
+    if (params.mode == GemmUniversalMode::kGemm ||
       params.mode == GemmUniversalMode::kGemmSplitKParallel) {
 
       if (threadblock_tile_offset.k() + 1 < params.grid_tiled_shape.k()) {
 
-        problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size; 
+        problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size;
       }
 
       offset_k = threadblock_tile_offset.k() * params.gemm_k_size;
@@ -539,10 +539,10 @@ public:
 
     // Compute threadblock-scoped matrix multiply-add
     mma(
-      gemm_k_iterations, 
-      accumulators, 
-      iterator_A, 
-      iterator_B, 
+      gemm_k_iterations,
+      accumulators,
+      iterator_A,
+      iterator_B,
       accumulators);
 
     //
@@ -571,16 +571,16 @@ public:
     typename Epilogue::ElementTensor *ptr_Tensor = static_cast<typename Epilogue::ElementTensor *>(params.ptr_Tensor);
 
     // Define the reduction output pointer and move to the appropriate place
-    typename Epilogue::ElementVector *ptr_Vector = 
+    typename Epilogue::ElementVector *ptr_Vector =
       static_cast<typename Epilogue::ElementVector *>(params.ptr_Vector);
 
     //
     // Fetch pointers based on mode.
     //
-    
+
     //
     // Special path when split-K not enabled.
-    // 
+    //
 
     if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() == 1) {
 
@@ -621,9 +621,9 @@ public:
 
       // Construct the epilogue
       Epilogue epilogue(
-        shared_storage.epilogue, 
-        thread_idx, 
-        warp_idx, 
+        shared_storage.epilogue,
+        thread_idx,
+        warp_idx,
         lane_idx);
 
       // Move to appropriate location for this output tile
@@ -649,7 +649,7 @@ public:
     // Slower path when split-K or batching is needed
     //
 
-      
+
     #if SPLIT_K_ENABLED
     // Construct the semaphore.
     Semaphore semaphore(params.semaphore + block_idx, thread_idx);
@@ -658,7 +658,7 @@ public:
 
       // If performing a reduction via split-K, fetch the initial synchronization
       if (params.grid_tiled_shape.k() > 1) {
-        
+
         // Fetch the synchronization lock initially but do not block.
         semaphore.fetch();
 
@@ -739,15 +739,15 @@ public:
 
     // Construct the epilogue
     Epilogue epilogue(
-      shared_storage.epilogue, 
-      thread_idx, 
-      warp_idx, 
+      shared_storage.epilogue,
+      thread_idx,
+      warp_idx,
       lane_idx);
 
     #if SPLIT_K_ENABLED
     // Wait on the semaphore - this latency may have been covered by iterator construction
     if ((params.mode == GemmUniversalMode::kGemm) && params.grid_tiled_shape.k() > 1) {
-        
+
       // For subsequent threadblocks, the source matrix is held in the 'D' tensor.
       if (threadblock_tile_offset.k()) {
         iterator_C1 = iterator_D;
@@ -783,7 +783,7 @@ public:
     //
 
     #if SPLIT_K_ENABLED
-    if ((params.mode == GemmUniversalMode::kGemm)  && params.grid_tiled_shape.k() > 1) { 
+    if ((params.mode == GemmUniversalMode::kGemm)  && params.grid_tiled_shape.k() > 1) {
 
       int lock = 0;
       if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) {
@@ -795,7 +795,7 @@ public:
         // Otherwise, the semaphore is incremented
         lock = threadblock_tile_offset.k() + 1;
       }
-      
+
       semaphore.release(lock);
     }
     #endif
@@ -804,7 +804,7 @@ public:
 
 // GemmWithFusedEpilogue with one source
 template <
-  typename Mma_,                  ///! Threadblock-scoped matrix multiply-accumulate 
+  typename Mma_,                  ///! Threadblock-scoped matrix multiply-accumulate
   typename Epilogue_,             ///! Epilogue
   typename ThreadblockSwizzle_    ///! Threadblock swizzling function
 >
@@ -844,7 +844,7 @@ public:
 
   /// Split-K preserves splits that are 128b aligned
   static int const kSplitKAlignment = const_max(
-    128 / sizeof_bits<ElementA>::value, 
+    128 / sizeof_bits<ElementA>::value,
     128 / sizeof_bits<ElementB>::value
   );
 
@@ -885,8 +885,8 @@ public:
     //
     // Methods
     //
-    
-    Arguments(): 
+
+    Arguments():
       ptr_A(nullptr),
       ptr_B(nullptr),
       ptr_C(nullptr),
@@ -919,13 +919,13 @@ public:
       typename LayoutC::Stride::Index ldt)
     :
       UniversalArgumentsBase(mode, problem_size, batch_count, batch_stride_D),
-      epilogue(epilogue), 
-      ptr_A(ptr_A), ptr_B(ptr_B), ptr_C(ptr_C), ptr_D(ptr_D), 
-      ptr_Vector(ptr_Vector), 
+      epilogue(epilogue),
+      ptr_A(ptr_A), ptr_B(ptr_B), ptr_C(ptr_C), ptr_D(ptr_D),
+      ptr_Vector(ptr_Vector),
       ptr_Tensor(ptr_Tensor),
-      batch_stride_A(batch_stride_A), 
-      batch_stride_B(batch_stride_B), 
-      batch_stride_C(batch_stride_C), 
+      batch_stride_A(batch_stride_A),
+      batch_stride_B(batch_stride_B),
+      batch_stride_C(batch_stride_C),
       batch_stride_Vector(batch_stride_Vector),
       batch_stride_Tensor(batch_stride_Tensor),
       lda(lda), ldb(ldb), ldc(ldc), ldd(ldd), ldr(ldr), ldt(ldt)
@@ -940,7 +940,7 @@ public:
     /// Returns arguments for the transposed problem
     Arguments transposed_problem() const {
       Arguments args(*this);
-      
+
       std::swap(args.problem_size.m(), args.problem_size.n());
       std::swap(args.ptr_A, args.ptr_B);
       std::swap(args.lda, args.ldb);
@@ -1035,7 +1035,7 @@ public:
       batch_stride_Vector(args.batch_stride_Vector),
       batch_stride_Tensor(args.batch_stride_Tensor)
     {
-      CUTLASS_TRACE_HOST("GemmWithFusedEpilogue::Params::Params() - problem_size: " << problem_size);
+      CUTLASS_TRACE_HOST("GemmWithFusedEpilogue::Params::Params()");
       CUTLASS_TRACE_HOST("  ptr_Vector: " << (void *)this->ptr_Vector);
       CUTLASS_TRACE_HOST("  ptr_Tensor: " << (void *)this->ptr_Tensor);
       CUTLASS_TRACE_HOST("  ldr: " << this->ldr);
@@ -1186,7 +1186,7 @@ public:
     int offset_k = 0;
     int problem_size_k = params.problem_size.k();
 
-    ElementA *ptr_A = static_cast<ElementA *>(params.ptr_A); 
+    ElementA *ptr_A = static_cast<ElementA *>(params.ptr_A);
     ElementB *ptr_B = static_cast<ElementB *>(params.ptr_B);
 
 
@@ -1194,12 +1194,12 @@ public:
     //
     // Fetch pointers based on mode.
     //
-    if (params.mode == GemmUniversalMode::kGemm || 
+    if (params.mode == GemmUniversalMode::kGemm ||
       params.mode == GemmUniversalMode::kGemmSplitKParallel) {
 
       if (threadblock_tile_offset.k() + 1 < params.grid_tiled_shape.k()) {
 
-        problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size; 
+        problem_size_k = (threadblock_tile_offset.k() + 1) * params.gemm_k_size;
       }
 
       offset_k = threadblock_tile_offset.k() * params.gemm_k_size;
@@ -1265,10 +1265,10 @@ public:
 
     // Compute threadblock-scoped matrix multiply-add
     mma(
-      gemm_k_iterations, 
-      accumulators, 
-      iterator_A, 
-      iterator_B, 
+      gemm_k_iterations,
+      accumulators,
+      iterator_A,
+      iterator_B,
       accumulators);
 
     //
@@ -1296,16 +1296,16 @@ public:
     typename Epilogue::ElementTensor *ptr_Tensor = static_cast<typename Epilogue::ElementTensor *>(params.ptr_Tensor);
 
     // Define the reduction output pointer and move to the appropriate place
-    typename Epilogue::ElementVector *ptr_Vector = 
+    typename Epilogue::ElementVector *ptr_Vector =
       static_cast<typename Epilogue::ElementVector *>(params.ptr_Vector);
 
     //
     // Fetch pointers based on mode.
     //
-    
+
     //
     // Special path when split-K not enabled.
-    // 
+    //
 
     if (params.mode == GemmUniversalMode::kGemm && params.grid_tiled_shape.k() == 1) {
 
@@ -1338,9 +1338,9 @@ public:
 
       // Construct the epilogue
       Epilogue epilogue(
-        shared_storage.epilogue, 
-        thread_idx, 
-        warp_idx, 
+        shared_storage.epilogue,
+        thread_idx,
+        warp_idx,
         lane_idx);
 
       // Move to appropriate location for this output tile
@@ -1365,7 +1365,7 @@ public:
     // Slower path when split-K or batching is needed
     //
 
-      
+
     #if SPLIT_K_ENABLED
     // Construct the semaphore.
     Semaphore semaphore(params.semaphore + block_idx, thread_idx);
@@ -1374,7 +1374,7 @@ public:
 
       // If performing a reduction via split-K, fetch the initial synchronization
       if (params.grid_tiled_shape.k() > 1) {
-        
+
         // Fetch the synchronization lock initially but do not block.
         semaphore.fetch();
 
@@ -1441,15 +1441,15 @@ public:
 
     // Construct the epilogue
     Epilogue epilogue(
-      shared_storage.epilogue, 
-      thread_idx, 
-      warp_idx, 
+      shared_storage.epilogue,
+      thread_idx,
+      warp_idx,
       lane_idx);
 
     #if SPLIT_K_ENABLED
     // Wait on the semaphore - this latency may have been covered by iterator construction
     if ((params.mode == GemmUniversalMode::kGemm) && params.grid_tiled_shape.k() > 1) {
-        
+
       // For subsequent threadblocks, the source matrix is held in the 'D' tensor.
       if (threadblock_tile_offset.k()) {
         iterator_C = iterator_D;
@@ -1484,7 +1484,7 @@ public:
     //
 
     #if SPLIT_K_ENABLED
-    if ((params.mode == GemmUniversalMode::kGemm)  && params.grid_tiled_shape.k() > 1) { 
+    if ((params.mode == GemmUniversalMode::kGemm)  && params.grid_tiled_shape.k() > 1) {
 
       int lock = 0;
       if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) {
@@ -1496,7 +1496,7 @@ public:
         // Otherwise, the semaphore is incremented
         lock = threadblock_tile_offset.k() + 1;
       }
-      
+
       semaphore.release(lock);
     }
     #endif

include/cutlass/gemm/kernel/params_sparse_base.h

@@ -60,26 +60,24 @@ struct SparseParamsBase
   // Data members
   //
 
-  cutlass::gemm::GemmCoord problem_size;
-  cutlass::gemm::GemmCoord grid_tiled_shape;
+  cutlass::gemm::GemmCoord problem_size{};
+  cutlass::gemm::GemmCoord grid_tiled_shape{};
   int swizzle_log_tile;
-  ParamsA params_A;
-  TensorRefA ref_A;
-  ParamsB params_B;
-  TensorRefB ref_B;
-  ParamsE params_E;
-  TensorRefE ref_E;
-  int gemm_k_iterations;
-  int gemm_k_size;
+  ParamsA params_A{};
+  TensorRefA ref_A{};
+  ParamsB params_B{};
+  TensorRefB ref_B{};
+  ParamsE params_E{};
+  TensorRefE ref_E{};
+  int gemm_k_iterations{0};
+  int gemm_k_size{0};
 
   //
   // Host dispatch API
   //
 
   /// Default constructor
-  CUTLASS_HOST_DEVICE
-  SparseParamsBase() : swizzle_log_tile(0), gemm_k_iterations(0), gemm_k_size(0) { }
-
+  SparseParamsBase() = default;
 
   /// Constructor
   CUTLASS_HOST_DEVICE

include/cutlass/gemm/kernel/params_universal_base.h

@@ -54,8 +54,8 @@ template <class LayoutA, class LayoutB>
 CUTLASS_HOST_DEVICE
 static bool 
 is_continous_k_aligned(GemmCoord problem_size, size_t alignmentA, size_t alignmentB) {
-  return (std::is_same<LayoutA, layout::RowMajor>::value && (problem_size.k() % alignmentA) == 0) ||
-         (std::is_same<LayoutB, layout::ColumnMajor>::value && (problem_size.k() % alignmentB) == 0);
+  return (platform::is_same<LayoutA, layout::RowMajor>::value && (problem_size.k() % alignmentA) == 0) ||
+         (platform::is_same<LayoutB, layout::ColumnMajor>::value && (problem_size.k() % alignmentB) == 0);
 }
 
 }  // namespace util
@@ -69,21 +69,16 @@ struct UniversalArgumentsBase
   // Data members
   //
 
-  GemmUniversalMode mode;
-  GemmCoord problem_size;
-  int batch_count;
-
-  int64_t batch_stride_D;
+  GemmUniversalMode mode = cutlass::gemm::GemmUniversalMode::kGemm;
+  GemmCoord problem_size{};
+  int batch_count{1};
+  int64_t batch_stride_D{0};
 
   //
   // Methods
   //
 
-  UniversalArgumentsBase() :
-    mode(GemmUniversalMode::kGemm),
-    batch_count(1),
-    batch_stride_D(0)
-  {}
+  UniversalArgumentsBase() = default;
 
   /// constructs an arguments structure
   UniversalArgumentsBase(
@@ -117,17 +112,14 @@ struct UniversalParamsBase
   // Data members
   //
 
-  GemmCoord problem_size;
-  GemmCoord grid_tiled_shape;
-  int swizzle_log_tile;
-
-  GemmUniversalMode mode;
-  int batch_count;
-  int gemm_k_size;
-
-  int64_t batch_stride_D;
-
-  int *semaphore;
+  GemmCoord problem_size{};
+  GemmCoord grid_tiled_shape{};
+  int swizzle_log_tile{0};
+  GemmUniversalMode mode = cutlass::gemm::GemmUniversalMode::kGemm;
+  int batch_count {0};
+  int gemm_k_size {0};
+  int64_t batch_stride_D {0};
+  int *semaphore = nullptr;
 
 
   //
@@ -137,7 +129,6 @@ struct UniversalParamsBase
   /// Default constructor
   UniversalParamsBase() = default;
 
-
   /// Constructor
   UniversalParamsBase(
     UniversalArgumentsBase const &args, /// GEMM application arguments

include/cutlass/gemm/kernel/rank_2k_grouped.h

@@ -197,51 +197,34 @@ public:
     // Data members
     //
 
-    GemmUniversalMode mode;
-    GemmCoord *problem_sizes;
-    int problem_count;
-    int threadblock_count;
+    GemmUniversalMode mode = GemmUniversalMode::kGemm;
+    GemmCoord *problem_sizes = nullptr;
+    int problem_count{0};
+    int threadblock_count{0};
 
     typename EpilogueOutputOp::Params epilogue;
 
-    ElementA ** ptr_A;
-    ElementB ** ptr_B;
-    ElementC ** ptr_C;
-    ElementC ** ptr_D;
+    ElementA ** ptr_A = nullptr;
+    ElementB ** ptr_B = nullptr;
+    ElementC ** ptr_C = nullptr;
+    ElementC ** ptr_D = nullptr;
 
-    typename LayoutA::Stride::LongIndex *lda;
-    typename LayoutB::Stride::LongIndex *ldb;
-    typename LayoutC::Stride::LongIndex *ldc;
-    typename LayoutC::Stride::LongIndex *ldd;
+    typename LayoutA::Stride::LongIndex *lda = nullptr;
+    typename LayoutB::Stride::LongIndex *ldb = nullptr;
+    typename LayoutC::Stride::LongIndex *ldc = nullptr;
+    typename LayoutC::Stride::LongIndex *ldd = nullptr;
 
     // Only used by device-level operator
-    GemmCoord *host_problem_sizes;
+    GemmCoord *host_problem_sizes = nullptr;
 
-    bool allow_early_exit;
+    bool allow_early_exit = false;
 
     //
     // Methods
     //
 
     /// Default ctor
-    CUTLASS_HOST_DEVICE
-    Arguments():
-      mode(GemmUniversalMode::kGemm),
-      problem_count(0),
-      threadblock_count(0),
-      ptr_A(nullptr),
-      ptr_B(nullptr),
-      ptr_C(nullptr),
-      ptr_D(nullptr),
-      lda(nullptr),
-      ldb(nullptr),
-      ldc(nullptr),
-      ldd(nullptr),
-      host_problem_sizes(nullptr),
-      allow_early_exit(false)
-    {
-
-    }
+    Arguments() = default;
 
     /// Ctor
     CUTLASS_HOST_DEVICE
@@ -290,43 +273,31 @@ public:
   /// Parameters structure
   struct Params {
 
-    typename ProblemVisitor::Params problem_visitor;
-    int threadblock_count;
+    typename ProblemVisitor::Params problem_visitor{};
+    int threadblock_count = 0;
 
-    typename EpilogueOutputOp::Params output_op;
+    typename EpilogueOutputOp::Params output_op{};
 
-    GemmUniversalMode mode;
-    int batch_count;
+    GemmUniversalMode mode = cutlass::gemm::GemmUniversalMode::kGemm;
+    int batch_count = 0;
 
-    ElementA ** ptr_A;
-    ElementB ** ptr_B;
-    ElementC ** ptr_C;
-    ElementC ** ptr_D;
+    ElementA** ptr_A = nullptr;
+    ElementB** ptr_B = nullptr;
+    ElementC** ptr_C = nullptr;
+    ElementC** ptr_D = nullptr;
 
-    typename LayoutA::Stride::LongIndex *lda;
-    typename LayoutB::Stride::LongIndex *ldb;
-    typename LayoutC::Stride::LongIndex *ldc;
-    typename LayoutC::Stride::LongIndex *ldd;
+    typename LayoutA::Stride::LongIndex* lda = nullptr;
+    typename LayoutB::Stride::LongIndex* ldb = nullptr;
+    typename LayoutC::Stride::LongIndex* ldc = nullptr;
+    typename LayoutC::Stride::LongIndex* ldd = nullptr;
 
-    bool allow_early_exit;
+    bool allow_early_exit = false;
 
     //
     // Methods
     //
 
-    CUTLASS_HOST_DEVICE
-    Params():
-      mode(cutlass::gemm::GemmUniversalMode::kGemm),
-      ptr_A(nullptr),
-      ptr_B(nullptr),
-      ptr_C(nullptr),
-      ptr_D(nullptr),
-      lda(nullptr),
-      ldb(nullptr),
-      ldc(nullptr),
-      ldd(nullptr),
-      allow_early_exit(false)
-    { }
+    Params() = default;
 
     CUTLASS_HOST_DEVICE
     Params(Arguments const &args, void *workspace = nullptr, int tile_count = 0):
@@ -380,8 +351,7 @@ public:
   // Methods
   //
 
-  CUTLASS_DEVICE
-  Rank2KGrouped() { }
+  Rank2KGrouped() = default;
 
   /// Determines whether kernel satisfies alignment
   static Status can_implement(cutlass::gemm::GemmCoord const & problem_size) {

include/cutlass/gemm/kernel/rank_2k_universal.h

@@ -119,38 +119,34 @@ public:
     // Data members
     //
 
-    GemmUniversalMode mode;
-    GemmCoord problem_size;
-    int batch_count;
+    GemmUniversalMode mode = cutlass::gemm::GemmUniversalMode::kGemm;
+    GemmCoord problem_size {};
+    int batch_count{1};
 
-    typename EpilogueOutputOp::Params epilogue;
+    typename EpilogueOutputOp::Params epilogue{};
 
-    void const * ptr_A;
-    void const * ptr_B;
-    void const * ptr_C;
-    void * ptr_D;
+    void const * ptr_A = nullptr;
+    void const * ptr_B = nullptr;
+    void const * ptr_C = nullptr;
+    void * ptr_D = nullptr;
 
-    int64_t batch_stride_A;
-    int64_t batch_stride_B;
-    int64_t batch_stride_C;
-    int64_t batch_stride_D;
+    int64_t batch_stride_A {0};
+    int64_t batch_stride_B {0};
+    int64_t batch_stride_C {0};
+    int64_t batch_stride_D {0};
 
-    typename LayoutA::Stride::Index lda;
-    typename LayoutB::Stride::Index ldb;
-    typename LayoutC::Stride::Index ldc;
-    typename LayoutC::Stride::Index ldd;
+    typename LayoutA::Stride::Index lda{0};
+    typename LayoutB::Stride::Index ldb{0};
+    typename LayoutC::Stride::Index ldc{0};
+    typename LayoutC::Stride::Index ldd{0};
 
-    bool allow_early_exit;
+    bool allow_early_exit{false};
 
     //
     // Methods
     //
     
-    Arguments(): 
-      mode(GemmUniversalMode::kGemm), 
-      batch_count(1), 
-      ptr_A(nullptr), ptr_B(nullptr), ptr_C(nullptr), ptr_D(nullptr),
-      allow_early_exit(false) { }
+    Arguments() = default;
 
     /// constructs an arguments structure
     Arguments(
@@ -177,7 +173,8 @@ public:
       batch_count(batch_count),
       epilogue(epilogue), 
       ptr_A(ptr_A), ptr_B(ptr_B), ptr_C(ptr_C), ptr_D(ptr_D), 
-      batch_stride_A(batch_stride_A), batch_stride_C(batch_stride_C), batch_stride_D(batch_stride_D), 
+      batch_stride_A(batch_stride_A), batch_stride_B(0),
+      batch_stride_C(batch_stride_C), batch_stride_D(batch_stride_D), 
       lda(lda), ldb(ldb), ldc(ldc), ldd(ldd),
       allow_early_exit(allow_early_exit) {
 
@@ -203,67 +200,46 @@ public:
   /// Parameters structure
   struct Params {
 
-    cutlass::gemm::GemmCoord problem_size;
-    cutlass::gemm::GemmCoord grid_tiled_shape;
-    int swizzle_log_tile;
+    cutlass::gemm::GemmCoord problem_size{};
+    cutlass::gemm::GemmCoord grid_tiled_shape{};
+    int swizzle_log_tile{0};
     
     // Mma1 Iterator A and B params
-    typename Mma1::IteratorA::Params params_A;
-    typename Mma1::IteratorB::Params params_BT;
+    typename Mma1::IteratorA::Params params_A{};
+    typename Mma1::IteratorB::Params params_BT{};
 
     // Mma2 Iterator A and B params 
-    typename Mma2::IteratorA::Params params_B;
-    typename Mma2::IteratorB::Params params_AT;
+    typename Mma2::IteratorA::Params params_B{};
+    typename Mma2::IteratorB::Params params_AT{};
 
-    typename Epilogue::OutputTileIterator::Params params_C;
-    typename Epilogue::OutputTileIterator::Params params_D;
+    typename Epilogue::OutputTileIterator::Params params_C{};
+    typename Epilogue::OutputTileIterator::Params params_D{};
     
-    typename EpilogueOutputOp::Params output_op;
+    typename EpilogueOutputOp::Params output_op{};
 
-    GemmUniversalMode mode;
-    int batch_count;
-    int gemm_k_size;
+    GemmUniversalMode mode = cutlass::gemm::GemmUniversalMode::kGemm;
+    int batch_count{0};
+    int gemm_k_size{0};
 
-    void * ptr_A;
-    void * ptr_B;
-    void * ptr_C;
-    void * ptr_D;
+    void * ptr_A = nullptr;
+    void * ptr_B = nullptr;
+    void * ptr_C = nullptr;
+    void * ptr_D = nullptr;
 
-    int64_t batch_stride_A;
-    int64_t batch_stride_B;
-    int64_t batch_stride_C;
-    int64_t batch_stride_D;
+    int64_t batch_stride_A{0};
+    int64_t batch_stride_B{0};
+    int64_t batch_stride_C{0};
+    int64_t batch_stride_D{0};
 
-    int *semaphore;
+    int *semaphore = nullptr;
 
-    bool allow_early_exit;
+    bool allow_early_exit {false};
 
     //
     // Methods
     //
 
-    CUTLASS_HOST_DEVICE
-    Params():
-      swizzle_log_tile(0),
-      params_A(0),
-      params_BT(0),
-      params_B(0),
-      params_AT(0),
-      params_C(0),
-      params_D(0),
-      batch_count(0),
-      gemm_k_size(0),
-      mode(cutlass::gemm::GemmUniversalMode::kGemm),
-      ptr_A(nullptr),
-      ptr_B(nullptr),
-      ptr_C(nullptr),
-      ptr_D(nullptr),
-      batch_stride_A(0),
-      batch_stride_B(0),
-      batch_stride_C(0),
-      batch_stride_D(0),
-      semaphore(nullptr),
-      allow_early_exit(false) { }
+    Params() = default;
 
     CUTLASS_HOST_DEVICE
     Params(

include/cutlass/gemm/kernel/rank_k_universal.h

@@ -106,36 +106,32 @@ public:
     // Data members
     //
 
-    GemmUniversalMode mode;
-    GemmCoord problem_size;
-    int batch_count;
+    GemmUniversalMode mode{GemmUniversalMode::kGemm};
+    GemmCoord problem_size{};
+    int batch_count{1};
 
-    typename EpilogueOutputOp::Params epilogue;
+    typename EpilogueOutputOp::Params epilogue{};
 
-    void const * ptr_A;
-    void const * ptr_C;
-    void * ptr_D;
+    void const * ptr_A{nullptr};
+    void const * ptr_C{nullptr};
+    void * ptr_D{nullptr};
 
-    int64_t batch_stride_A;
-    int64_t batch_stride_C;
-    int64_t batch_stride_D;
+    int64_t batch_stride_A{0};
+    int64_t batch_stride_C{0};
+    int64_t batch_stride_D{0};
 
-    typename LayoutA::Stride::Index lda;
-    typename LayoutB::Stride::Index ldb;
-    typename LayoutC::Stride::Index ldc;
-    typename LayoutC::Stride::Index ldd;
+    typename LayoutA::Stride::Index lda{};
+    typename LayoutB::Stride::Index ldb{};
+    typename LayoutC::Stride::Index ldc{};
+    typename LayoutC::Stride::Index ldd{};
 
-    bool allow_early_exit;
+    bool allow_early_exit{false};
 
     //
     // Methods
     //
     
-    Arguments(): 
-      mode(GemmUniversalMode::kGemm), 
-      batch_count(1), 
-      ptr_A(nullptr), ptr_C(nullptr), ptr_D(nullptr),
-      allow_early_exit(false) { }
+    Arguments() = default;
 
     /// constructs an arguments structure
     Arguments(
@@ -160,7 +156,8 @@ public:
       epilogue(epilogue), 
       ptr_A(ptr_A), ptr_C(ptr_C), ptr_D(ptr_D), 
       batch_stride_A(batch_stride_A), batch_stride_C(batch_stride_C), batch_stride_D(batch_stride_D), 
-      lda(lda), ldb(ldb), ldc(ldc), ldd(ldd),
+      lda(lda), ldb(0),
+      ldc(ldc), ldd(ldd),
       allow_early_exit(allow_early_exit) {
 
       }
@@ -174,59 +171,38 @@ public:
   /// Parameters structure
   struct Params {
 
-    cutlass::gemm::GemmCoord problem_size;
-    cutlass::gemm::GemmCoord grid_tiled_shape;
-    int swizzle_log_tile;
+    cutlass::gemm::GemmCoord problem_size{};
+    cutlass::gemm::GemmCoord grid_tiled_shape{};
+    int swizzle_log_tile{0};
    
-    typename Mma::IteratorA::Params params_A;
-    typename Mma::IteratorB::Params params_B;
-    typename Epilogue::OutputTileIterator::Params params_C;
-    typename Epilogue::OutputTileIterator::Params params_D;
-    
-    typename EpilogueOutputOp::Params output_op;
+    typename Mma::IteratorA::Params params_A{};
+    typename Mma::IteratorB::Params params_B{};
+    typename Epilogue::OutputTileIterator::Params params_C{};
+    typename Epilogue::OutputTileIterator::Params params_D{};
+    typename EpilogueOutputOp::Params output_op{};
 
-    GemmUniversalMode mode;
-    int batch_count;
-    int gemm_k_size;
+    GemmUniversalMode mode = cutlass::gemm::GemmUniversalMode::kGemm;
+    int batch_count{0};
+    int gemm_k_size{0};
 
-    void * ptr_A;
-    void * ptr_B;
-    void * ptr_C;
-    void * ptr_D;
+    void * ptr_A{nullptr};
+    void * ptr_B{nullptr};
+    void * ptr_C{nullptr};
+    void * ptr_D{nullptr};
 
-    int64_t batch_stride_A;
-    int64_t batch_stride_B;
-    int64_t batch_stride_C;
-    int64_t batch_stride_D;
+    int64_t batch_stride_A{0};
+    int64_t batch_stride_B{0};
+    int64_t batch_stride_C{0};
+    int64_t batch_stride_D{0};
 
-    int *semaphore;
+    int *semaphore{nullptr};
 
-    bool allow_early_exit;
+    bool allow_early_exit{false};
 
     //
     // Methods
     //
-
-    CUTLASS_HOST_DEVICE
-    Params():
-      swizzle_log_tile(0),
-      params_A(0),
-      params_B(0),
-      params_C(0),
-      params_D(0),
-      batch_count(0),
-      gemm_k_size(0),
-      mode(cutlass::gemm::GemmUniversalMode::kGemm),
-      ptr_A(nullptr),
-      ptr_B(nullptr),
-      ptr_C(nullptr),
-      ptr_D(nullptr),
-      batch_stride_A(0),
-      batch_stride_B(0),
-      batch_stride_C(0),
-      batch_stride_D(0),
-      semaphore(nullptr),
-      allow_early_exit(false) { }
+    Params() = default;
 
     CUTLASS_HOST_DEVICE
     Params(

include/cutlass/gemm/kernel/sm70_gemm.hpp

@@ -116,10 +116,10 @@ static_assert(is_valid_tile_scheduler, "SM70 kernel does not support specializin
 
   // Kernel entry point API
   struct Params {
-    GemmUniversalMode mode;
-    ProblemShape problem_shape;
-    MainloopParams mainloop;
-    EpilogueParams epilogue;
+    GemmUniversalMode mode{};
+    ProblemShape problem_shape{};
+    MainloopParams mainloop{};
+    EpilogueParams epilogue{};
   };
 
   //
@@ -150,15 +150,16 @@ static_assert(is_valid_tile_scheduler, "SM70 kernel does not support specializin
     return mode_implementable && TileScheduler::can_implement(args.scheduler);
   }
 
-  static int
+  static size_t
   get_workspace_size(Arguments const& args) {
-    int workspace_size = 0;
+    size_t workspace_size = 0;
     return workspace_size;
   }
 
   static
   cutlass::Status
-  initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
+  initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr, 
+    CudaHostAdapter* cuda_adapter = nullptr) {
     cutlass::Status status = Status::kSuccess;
 
     return status;
@@ -250,7 +251,6 @@ static_assert(is_valid_tile_scheduler, "SM70 kernel does not support specializin
       thread_idx,
       smem_buf
     );
-
     // Epilogue and write to gD
     CollectiveEpilogue epilogue{params.epilogue};
     epilogue(

include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp

@@ -168,13 +168,13 @@ public:
 
   // Kernel entry point API
   struct Params {
-    GemmUniversalMode mode;
-    ProblemShape problem_shape;
-    MainloopParams mainloop;
-    EpilogueParams epilogue;
-    KernelHardwareInfo hw_info;
-    TileSchedulerParams scheduler;
-    void* workspace;
+    GemmUniversalMode mode{};
+    ProblemShape problem_shape{};
+    MainloopParams mainloop{};
+    EpilogueParams epilogue{};
+    KernelHardwareInfo hw_info{};
+    TileSchedulerParams scheduler{};
+    void* workspace{nullptr};
   };
 
   //
@@ -290,7 +290,8 @@ public:
   }
 
   static cutlass::Status
-  initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
+  initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr,
+    CudaHostAdapter* cuda_adapter = nullptr) {
     Status status = Status::kSuccess;
     uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
     size_t workspace_offset = 0;
@@ -305,7 +306,7 @@ public:
       return status;
     }
 
-    status = CollectiveEpilogue::initialize_workspace(args.problem_shape, args.epilogue, workspace_ptr + workspace_offset, stream);
+    status = CollectiveEpilogue::initialize_workspace(args.problem_shape, args.epilogue, workspace_ptr + workspace_offset, stream, cuda_adapter);
     workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
     workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
 

include/cutlass/gemm/kernel/sm90_gemm_tma.hpp

@@ -135,10 +135,10 @@ public:
 
   // Kernel entry point API
   struct Params {
-    GemmUniversalMode mode;
-    ProblemShape problem_shape;
-    MainloopParams mainloop;
-    EpilogueParams epilogue;
+    GemmUniversalMode mode{};
+    ProblemShape problem_shape{};
+    MainloopParams mainloop{};
+    EpilogueParams epilogue{};
   };
 
   //
@@ -180,13 +180,14 @@ public:
     return implementable;
   }
 
-  static int
+  static size_t
   get_workspace_size(Arguments const& args) {
     return 0;
   }
 
   static cutlass::Status
-  initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
+  initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr,
+    CudaHostAdapter* cuda_adapter = nullptr) {
     return Status::kSuccess;
   }
 

include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp

@@ -141,10 +141,10 @@ public:
 
   // Kernel entry point API
   struct Params {
-    GemmUniversalMode mode;
-    ProblemShape problem_shape;
-    MainloopParams mainloop;
-    EpilogueParams epilogue;
+    GemmUniversalMode mode{};
+    ProblemShape problem_shape{};
+    MainloopParams mainloop{};
+    EpilogueParams epilogue{};
   };
 
   //
@@ -187,14 +187,15 @@ public:
   }
 
   static
-  int
+  size_t
   get_workspace_size(Arguments const& args) {
     return 0;
   }
 
   static
   cutlass::Status
-  initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
+  initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr,
+    CudaHostAdapter* cuda_adapter = nullptr) {
     return Status::kSuccess;
   }
 

include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp

@@ -148,13 +148,13 @@ public:
 
   // Kernel entry point API
   struct Params {
-    GemmUniversalMode mode;
-    ProblemShape problem_shape;
-    MainloopParams mainloop;
-    EpilogueParams epilogue;
-    KernelHardwareInfo hw_info;
-    TileSchedulerParams scheduler;
-    void* workspace;
+    GemmUniversalMode mode{};
+    ProblemShape problem_shape{};
+    MainloopParams mainloop{};
+    EpilogueParams epilogue{};
+    KernelHardwareInfo hw_info{};
+    TileSchedulerParams scheduler{};
+    void* workspace{nullptr};
   };
 
   //
@@ -250,7 +250,8 @@ public:
   }
 
   static cutlass::Status
-  initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
+  initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr, 
+    CudaHostAdapter* cuda_adapter = nullptr) {
     Status status = Status::kSuccess;
     uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
     size_t workspace_offset = 0;
@@ -265,7 +266,7 @@ public:
       return status;
     }
 
-    status = CollectiveEpilogue::initialize_workspace(args.problem_shape, args.epilogue, workspace_ptr + workspace_offset, stream);
+    status = CollectiveEpilogue::initialize_workspace(args.problem_shape, args.epilogue, workspace_ptr + workspace_offset, stream, cuda_adapter);
     workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
     workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
     if (status != Status::kSuccess) {

include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp

@@ -156,12 +156,12 @@ public:
 
   // Kernel entry point API
   struct Params {
-    GemmUniversalMode mode;
-    ProblemShape problem_shape;
-    MainloopParams mainloop;
-    EpilogueParams epilogue;
-    KernelHardwareInfo hw_info;
-    TileSchedulerParams scheduler;
+    GemmUniversalMode mode{};
+    ProblemShape problem_shape{};
+    MainloopParams mainloop{};
+    EpilogueParams epilogue{};
+    KernelHardwareInfo hw_info{};
+    TileSchedulerParams scheduler{};
   };
 
   //
@@ -249,7 +249,8 @@ public:
   }
 
   static cutlass::Status
-  initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
+  initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr,
+    CudaHostAdapter* cuda_adapter = nullptr) {
     Status status = Status::kSuccess;
     uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
     size_t workspace_offset = 0;
@@ -263,7 +264,7 @@ public:
       return status;
     }
 
-    status = CollectiveEpilogue::initialize_workspace(args.problem_shape, args.epilogue, workspace_ptr + workspace_offset, stream);
+    status = CollectiveEpilogue::initialize_workspace(args.problem_shape, args.epilogue, workspace_ptr + workspace_offset, stream, cuda_adapter);
     workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
     workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
     if (status != Status::kSuccess) {

include/cutlass/gemm/kernel/sm90_gemm_warpspecialized.hpp

@@ -145,10 +145,10 @@ public:
 
   // Kernel entry point API
   struct Params {
-    GemmUniversalMode mode;
-    ProblemShape problem_shape;
-    MainloopParams mainloop;
-    EpilogueParams epilogue;
+    GemmUniversalMode mode{};
+    ProblemShape problem_shape{};
+    MainloopParams mainloop{};
+    EpilogueParams epilogue{};
   };
 
   //
@@ -191,14 +191,15 @@ public:
   }
 
   static
-  int
+  size_t
   get_workspace_size(Arguments const& args) {
     return 0;
   }
 
   static
   cutlass::Status
-  initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
+  initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr,
+    CudaHostAdapter* cuda_adapter = nullptr) {
     return Status::kSuccess;
   }
 

include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_cooperative.hpp

@@ -146,12 +146,12 @@ public:
 
   // Kernel entry point API
   struct Params {
-    GemmUniversalMode mode;
-    ProblemShape problem_shape;
-    MainloopParams mainloop;
-    EpilogueParams epilogue;
-    KernelHardwareInfo hw_info;
-    TileSchedulerParams scheduler;
+    GemmUniversalMode mode{};
+    ProblemShape problem_shape{};
+    MainloopParams mainloop{};
+    EpilogueParams epilogue{};
+    KernelHardwareInfo hw_info{};
+    TileSchedulerParams scheduler{};
   };
 
   //
@@ -213,7 +213,7 @@ public:
   }
 
   static
-  int
+  size_t
   get_workspace_size(Arguments const& args) {
     TileScheduler t;
     return t.template get_workspace_size<ProblemShape, ElementAccumulator>(
@@ -222,7 +222,8 @@ public:
 
   static
   cutlass::Status
-  initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
+  initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr,
+    CudaHostAdapter* cuda_adapter = nullptr) {
     TileScheduler t;
     return t.template initialize_workspace<ProblemShape, ElementAccumulator>(
       args.scheduler, workspace, stream, args.problem_shape, args.hw_info, NumMmaWarpGroups);

include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_pingpong.hpp

@@ -156,12 +156,12 @@ public:
 
   // Kernel entry point API
   struct Params {
-    GemmUniversalMode mode;
-    ProblemShape problem_shape;
-    MainloopParams mainloop;
-    EpilogueParams epilogue;
-    KernelHardwareInfo hw_info;
-    TileSchedulerParams scheduler;
+    GemmUniversalMode mode{};
+    ProblemShape problem_shape{};
+    MainloopParams mainloop{};
+    EpilogueParams epilogue{};
+    KernelHardwareInfo hw_info{};
+    TileSchedulerParams scheduler{};
   };
 
   //
@@ -224,14 +224,15 @@ public:
   }
 
   static
-  int
+  size_t
   get_workspace_size(Arguments const& args) {
     return 0;
   }
 
   static
   cutlass::Status
-  initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
+  initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr,
+    CudaHostAdapter* cuda_adapter = nullptr) {
     return Status::kSuccess;
   }
 

include/cutlass/gemm/kernel/sm90_tile_scheduler.hpp

@@ -58,11 +58,34 @@ public:
       FastDivmodU64 const& divmod_cluster_blk_major,
       int32_t log_swizzle_size,
       RasterOrder raster_order) {
+    auto [cta_m_in_cluster, cta_n_in_cluster, _] = cute::block_id_in_cluster();
+    return get_work_idx_m_and_n(
+      blk_per_grid_dim,
+      divmod_cluster_shape_major,
+      divmod_cluster_shape_minor,
+      divmod_cluster_blk_major,
+      log_swizzle_size,
+      raster_order,
+      cta_m_in_cluster,
+      cta_n_in_cluster
+    );
+  }
+
+  static CUTLASS_DEVICE
+  cute::tuple<int32_t, int32_t>
+  get_work_idx_m_and_n(
+      uint64_t blk_per_grid_dim,
+      FastDivmodU64Pow2 const& divmod_cluster_shape_major,
+      FastDivmodU64Pow2 const& divmod_cluster_shape_minor,
+      FastDivmodU64 const& divmod_cluster_blk_major,
+      int32_t log_swizzle_size,
+      RasterOrder raster_order,
+      uint64_t cta_m_in_cluster,
+      uint64_t cta_n_in_cluster) {
 
     uint64_t cluster_id, cluster_major_offset = 0, cluster_minor_offset = 0;
     divmod_cluster_shape_major(cluster_id, cluster_major_offset, blk_per_grid_dim);
 
-    auto [cta_m_in_cluster, cta_n_in_cluster, _] = cute::block_id_in_cluster();
     if (raster_order == RasterOrder::AlongN) {
       cluster_minor_offset = cta_m_in_cluster;
     }
@@ -97,7 +120,7 @@ public:
 
   // The basic tile scheduler does not require any additional workspace
   template <class ProblemShape, class ElementAccumulator>
-  static int
+  static size_t
   get_workspace_size(Arguments const&, ProblemShape, KernelHardwareInfo const&, uint32_t, const uint32_t = 1) {
     return 0;
   }

include/cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp

@@ -400,7 +400,7 @@ public:
 
   // The basic tile scheduler does not require any additional workspace
   template <class ProblemShape, class ElementAccumulator>
-  static int
+  static size_t
   get_workspace_size(Arguments const&, ProblemShape, KernelHardwareInfo const&, uint32_t, const uint32_t = 1) {
     return 0;
   }

include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp

@@ -386,7 +386,8 @@ public:
     WorkTileInfo const& work_tile_info,
     FrgTensorC& accumulators,
     uint32_t num_barriers,
-    uint32_t barrier_idx) {
+    uint32_t barrier_idx,
+    uint32_t num_accumulator_mtxs = 1) {
 
     using ElementAccumulator = typename FrgTensorC::value_type;
 
@@ -412,7 +413,7 @@ public:
     // Reductions use BlockStripedReduce with a width of BarrierManager::ThreadCount under the hood.
     // Thus, the start of the reduction space is the same across all threads in a warp group.
     int reduction_offset =
-      (cute::size<0>(TileShape{}) * cute::size<1>(TileShape{}) * reduction_tile_idx) +
+      (cute::size<0>(TileShape{}) * cute::size<1>(TileShape{}) * reduction_tile_idx * num_accumulator_mtxs) +
       reduction_peer_offset +
       (size(accumulators) * barrier_idx * BarrierManager::ThreadCount);
 
@@ -444,7 +445,7 @@ public:
     }
 
     auto reduction_workspace_size = Params::get_reduction_workspace_size(
-      reduction_tiles, to_gemm_coord(TileShape{}), sizeof_bits<ElementAccumulator>::value);
+      reduction_tiles, to_gemm_coord(TileShape{}), sizeof_bits<ElementAccumulator>::value, num_accumulator_mtxs);
     BarrierType* lock_workspace = reinterpret_cast<BarrierType*>(
       reinterpret_cast<uint8_t*>(params.reduction_workspace_) + reduction_workspace_size);
 
@@ -540,7 +541,7 @@ public:
   }
 
   template <class ProblemShape, class ElementAccumulator>
-  static int
+  static size_t
   get_workspace_size(
     Arguments const& args,
     ProblemShape problem_shape,
@@ -836,7 +837,8 @@ private:
                                           params.divmod_cluster_shape_minor_,
                                           params.divmod_cluster_blk_major_,
                                           params.log_swizzle_size_,
-                                          params.raster_order_);
+                                          params.raster_order_
+                                        );
 
     // Set the M, N, and L block offsets
     work_tile_info.M_idx = work_idx_m;

include/cutlass/gemm/kernel/sparse_gemm_with_absmax.h

@@ -0,0 +1,509 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/*! \file
+    \brief Sparse GEMM kernel with an epilogue that computes the absolute maximum value of the output
+    and a pre-activation-function auxiliary output. The auxiliary output is also (optionally)
+    stored to global memory.
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+
+#include "cutlass/gemm/gemm.h"
+#include "cutlass/gemm/kernel/params_sparse_base.h"
+#include "cutlass/matrix_coord.h"
+#include "cutlass/semaphore.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace kernel {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  typename Mma_,                  ///! Threadblock-scoped matrix multiply-accumulate 
+  typename Epilogue_,             ///! Epilogue
+  typename ThreadblockSwizzle_,   ///! Threadblock swizzling function
+  bool SplitKSerial               ///! If true, code supporting split-K via serial reduction is enabled.
+>
+struct SparseGemmWithAbsmax {
+
+  using Mma = Mma_;
+  using Epilogue = Epilogue_;
+  using OutputOp = typename Epilogue::OutputOp;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
+  static bool const kSplitKSerial = SplitKSerial;
+
+  static int const kSparse = Mma::kSparse;
+  static int const kMetaSizeInBits = Mma::kMetaSizeInBits;
+  static int const kMaxID2 = Mma::kMaxID2;
+  static int const kElementsPerElementE = Mma::kElementsPerElementE;
+
+  using ElementE = typename Mma::ElementE;
+  using LayoutE = typename Mma::LayoutE;
+
+  using LayoutC = typename Epilogue::OutputTileIterator::Layout;
+
+  /// Warp count (concept: GemmShape)
+  using WarpCount = typename Mma::WarpCount;
+  static int const kThreadCount = 32 * WarpCount::kCount;
+
+  using ParamsA = typename Mma::IteratorA::Params;
+  using TensorRefA = typename Mma::IteratorA::TensorRef;
+  using ParamsB = typename Mma::IteratorB::Params;
+  using TensorRefB = typename Mma::IteratorB::TensorRef;
+  using ParamsE = typename Mma::IteratorE::Params;
+  using TensorRefE = typename Mma::IteratorE::TensorRef;
+
+  using ParamsC = typename Epilogue::OutputTileIterator::Params;
+  using TensorRefC = typename Epilogue::OutputTileIterator::TensorRef;
+  using ParamsD = typename Epilogue::OutputTileIterator::Params;
+  using TensorRefD = typename Epilogue::OutputTileIterator::TensorRef;
+  using ParamsAux = typename Epilogue::AuxOutputTileIterator::Params;
+  using TensorRefAux = typename Epilogue::AuxOutputTileIterator::TensorRef;
+
+  /// Argument structure
+  struct Arguments {
+
+    //
+    // Data members
+    //
+
+    GemmCoord problem_size;
+    TensorRefA ref_A;
+    TensorRefB ref_B;
+    TensorRefC ref_C;
+    TensorRefD ref_D;
+    TensorRefE ref_E;
+    TensorRefAux ref_Aux;
+    void* ptr_Vector;
+    typename LayoutC::Stride::Index ldr;
+
+    typename Epilogue::OutputOp::Params epilogue;
+    int split_k_slices;
+
+    //
+    // Methods
+    //
+
+    /// Default ctor
+    CUTLASS_HOST_DEVICE
+    Arguments(): problem_size(0, 0, 0), split_k_slices(1) {
+
+    }
+
+    /// Constructs an Arguments structure 
+    CUTLASS_HOST_DEVICE
+    Arguments(
+      GemmCoord problem_size_,
+      TensorRefA ref_A_,
+      TensorRefB ref_B_,
+      TensorRefC ref_C_,
+      TensorRefD ref_D_,
+      TensorRefE ref_E_,
+      TensorRefAux ref_Aux_,
+      void* ptr_Vector_,
+      typename LayoutC::Stride::Index ldr_,
+      typename OutputOp::Params epilogue_ = 
+        typename OutputOp::Params(),
+      int split_k_slices = 1
+    ):
+      problem_size(problem_size_),
+      ref_A(ref_A_),
+      ref_B(ref_B_),
+      ref_C(ref_C_),
+      ref_D(ref_D_),
+      ref_E(ref_E_),
+      ref_Aux(ref_Aux_),
+      ptr_Vector(ptr_Vector_),
+      ldr(ldr_),
+      epilogue(epilogue_),
+      split_k_slices(split_k_slices) {
+
+    }
+  };
+
+  /// Parameters structure
+  struct Params : public SparseParamsBase<
+      ThreadblockSwizzle, ParamsA, TensorRefA, ParamsB, TensorRefB,
+      ParamsE, TensorRefE> {
+
+    using Base = SparseParamsBase<
+        ThreadblockSwizzle, ParamsA, TensorRefA, ParamsB, TensorRefB,
+        ParamsE, TensorRefE>;
+
+    //
+    // Data members
+    //
+
+    ParamsC params_C;
+    TensorRefC ref_C;
+    ParamsD params_D;
+    TensorRefD ref_D;
+    ParamsAux params_Aux;
+    TensorRefAux ref_Aux;
+
+    void* ptr_Vector;
+    typename LayoutC::Stride::Index ldr;
+
+    typename OutputOp::Params output_op;
+    int *semaphore;
+
+    //
+    // Methods
+    //
+
+    CUTLASS_HOST_DEVICE
+    Params() { }
+
+    CUTLASS_HOST_DEVICE
+    Params(
+      cutlass::gemm::GemmCoord const & problem_size,
+      cutlass::gemm::GemmCoord const & grid_tiled_shape,
+      TensorRefA ref_A,
+      TensorRefB ref_B,
+      TensorRefC ref_C,
+      TensorRefD ref_D,
+      TensorRefE ref_E,
+      TensorRefAux ref_Aux,
+      void* ptr_Vector,
+      typename LayoutC::Stride::Index ldr,
+      typename OutputOp::Params output_op = typename OutputOp::Params(),
+      int *workspace = nullptr
+    ):
+      Base(problem_size, grid_tiled_shape, ref_A, ref_B, ref_E, Mma::Shape::kK),
+      params_C(ref_C.layout()),
+      ref_C(ref_C),
+      params_D(ref_D.layout()),
+      ref_D(ref_D),
+      output_op(output_op),
+      ref_Aux(ref_Aux),
+      params_Aux(ref_Aux.layout()),
+      ptr_Vector(ptr_Vector),
+      ldr(ldr) {
+    semaphore = workspace;
+    }
+  };
+
+  /// Shared memory storage structure
+  union SharedStorage {
+    typename Mma::SharedStorage main_loop;
+    typename Epilogue::SharedStorage epilogue;
+  };
+
+  //
+  // Methods
+  //
+
+  CUTLASS_HOST_DEVICE
+  SparseGemmWithAbsmax() { } 
+
+  /// Determines whether kernel satisfies alignment
+  static Status can_implement(
+      cutlass::gemm::GemmCoord const & problem_size,
+      typename Mma::IteratorA::TensorRef ref_A,
+      typename Mma::IteratorB::TensorRef ref_B,
+      typename Epilogue::OutputTileIterator::TensorRef ref_C,
+      typename Epilogue::OutputTileIterator::TensorRef ref_D,
+      typename Mma::IteratorE::TensorRef ref_E) {
+
+    static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
+    static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
+    static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess;
+    static int const kAlignmentE = Mma::IteratorE::AccessType::kElements;
+
+    if (!TensorRef_aligned(ref_A, kAlignmentA)) {
+      return Status::kErrorMisalignedOperand;
+    }
+
+    if (!TensorRef_aligned(ref_B, kAlignmentB)) {
+      return Status::kErrorMisalignedOperand;
+    }
+
+    if (!TensorRef_aligned(ref_C, kAlignmentC)) {
+      return Status::kErrorMisalignedOperand;
+    }
+
+    if (!TensorRef_aligned(ref_D, kAlignmentC)) {
+      return Status::kErrorMisalignedOperand;
+    }
+
+    if (!TensorRef_aligned(ref_E, kAlignmentE)) {
+      return Status::kErrorMisalignedOperand;
+    }
+
+    if ((problem_size.m() % kAlignmentA) || ((problem_size.k() / kSparse) % kAlignmentA) ||
+      (problem_size.n() % kAlignmentB) || (problem_size.k() % kAlignmentB) ||
+      (problem_size.m() % kAlignmentC) || (problem_size.n() % kAlignmentC) ||
+      (problem_size.m() % kAlignmentE) || ((problem_size.k() / kSparse) % kAlignmentE)) {
+
+      return Status::kErrorMisalignedOperand;
+    }
+
+    // The k dimension has to be the multiple of the Threadblock k because out
+    // of bound meta data would be initialized to 0 by acync.zfill but 0 is not
+    // a valid meta data.
+    if (problem_size.k() % Mma::Shape::kK) {
+      return Status::kErrorMisalignedOperand;
+    }
+
+    // M dimension has to be multiple of 32 (sparse float) or 16 (sparse int) 
+    // because of the row reordering of operand E
+    static int const kAlignmentM = (sizeof(ElementE) == 2) ? 32 : 16;
+
+    if (problem_size.m() % kAlignmentM) {
+      return Status::kErrorMisalignedOperand;
+    }
+
+    return Status::kSuccess;
+  }
+
+  /// Executes one GEMM
+  CUTLASS_DEVICE
+  void operator()(Params const &params, SharedStorage &shared_storage) {
+
+    // Compute threadblock location
+    ThreadblockSwizzle threadblock_swizzle;
+
+    cutlass::gemm::GemmCoord threadblock_tile_offset =
+        threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
+
+    // Early exit if CTA is out of range
+    if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() ||
+      params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) {
+
+      return;
+    }
+
+    // Compute initial location in logical coordinates
+    cutlass::MatrixCoord tb_offset_A{
+      threadblock_tile_offset.m() * Mma::Shape::kM,
+      threadblock_tile_offset.k() * params.gemm_k_size / kSparse,
+    };
+
+    cutlass::MatrixCoord tb_offset_B{
+      threadblock_tile_offset.k() * params.gemm_k_size,
+      threadblock_tile_offset.n() * Mma::Shape::kN
+    };
+
+    cutlass::MatrixCoord tb_offset_E{
+      threadblock_tile_offset.m() * Mma::Shape::kM,
+      threadblock_tile_offset.k() * params.gemm_k_size / kSparse,
+    };
+
+    // Problem size is a function of threadblock index in the K dimension
+    int problem_size_k = min(
+      params.problem_size.k(), 
+      (threadblock_tile_offset.k() + 1) * params.gemm_k_size);
+
+    // Compute threadblock-scoped matrix multiply-add
+    int gemm_k_iterations = (problem_size_k - tb_offset_B.row() + Mma::Shape::kK - 1) / Mma::Shape::kK;
+
+    // Compute position within threadblock
+    int thread_idx = threadIdx.x;
+
+    // Construct iterators to A, B, and E operands
+    typename Mma::IteratorA iterator_A(
+      params.params_A,
+      params.ref_A.data(),
+      {params.problem_size.m(), problem_size_k / kSparse},
+      thread_idx,
+      tb_offset_A);
+
+    typename Mma::IteratorB iterator_B(
+      params.params_B,
+      params.ref_B.data(),
+      {problem_size_k, params.problem_size.n()},
+      thread_idx,
+      tb_offset_B);
+
+    typename Mma::IteratorE iterator_E(
+        params.params_E, params.ref_E.data(),
+        {params.problem_size.m(),
+         problem_size_k / kSparse / kElementsPerElementE},
+        thread_idx, tb_offset_E);
+
+    // Broadcast the warp_id computed by lane 0 to ensure dependent code
+    // is compiled as warp-uniform.
+    int warp_idx = canonical_warp_idx_sync();
+    int lane_idx = threadIdx.x % 32;
+
+    //
+    // Main loop
+    //
+
+    // Construct thread-scoped matrix multiply
+    Mma mma(shared_storage.main_loop, thread_idx, warp_idx, lane_idx);
+
+    typename Mma::FragmentC accumulators;
+
+    accumulators.clear();
+
+    if (!kSplitKSerial || gemm_k_iterations > 0) {
+      // Compute threadblock-scoped matrix multiply-add
+      mma(gemm_k_iterations, accumulators, iterator_A, iterator_B, iterator_E, accumulators);
+    }
+
+    //
+    // Epilogue
+    //
+
+    OutputOp output_op(params.output_op);
+
+    //
+    // Masked tile iterators constructed from members
+    //
+
+    threadblock_tile_offset =
+        threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
+
+    //assume identity swizzle
+    MatrixCoord threadblock_offset(
+      threadblock_tile_offset.m() * Mma::Shape::kM,
+      threadblock_tile_offset.n() * Mma::Shape::kN
+    );
+
+    int block_idx = threadblock_tile_offset.m() + threadblock_tile_offset.n() * params.grid_tiled_shape.m();
+
+    // Construct the semaphore.
+    Semaphore semaphore(params.semaphore + block_idx, thread_idx);
+
+    // If performing a reduction via split-K, fetch the initial synchronization
+    if (kSplitKSerial && params.grid_tiled_shape.k() > 1) {
+      
+      // Fetch the synchronization lock initially but do not block.
+      semaphore.fetch();
+
+      // Indicate which position in a serial reduction the output operator is currently updating
+      output_op.set_k_partition(threadblock_tile_offset.k(), params.grid_tiled_shape.k());
+    }
+
+    typename Epilogue::ElementVector *ptr_Vector = static_cast<typename Epilogue::ElementVector *>(params.ptr_Vector);
+    // Move to appropriate location for this output tile
+    if (ptr_Vector) {
+      ptr_Vector += threadblock_offset.column() + threadblock_tile_offset.m() * params.ldr;
+    }
+
+    // Tile iterator loading from source tensor.
+    typename Epilogue::OutputTileIterator iterator_C(
+      params.params_C,
+      params.ref_C.data(),
+      params.problem_size.mn(),
+      thread_idx,
+      threadblock_offset
+    );
+
+    // Tile iterator writing to destination tensor.
+    typename Epilogue::OutputTileIterator iterator_D(
+      params.params_D,
+      params.ref_D.data(),
+      params.problem_size.mn(),
+      thread_idx,
+      threadblock_offset
+    );
+
+    // Tile iterator writing to auxiliary destination tensor.
+    typename Epilogue::AuxOutputTileIterator iterator_Aux(
+      params.params_Aux,
+      // Only the final block writes the auxiliary tensor
+      ((kSplitKSerial && params.grid_tiled_shape.k() > 1) &&
+          (params.grid_tiled_shape.k() != threadblock_tile_offset.k() + 1))
+          ? nullptr
+          : params.ref_Aux.data(),
+      params.problem_size.mn(),
+      thread_idx,
+      threadblock_offset
+    );
+
+    Epilogue epilogue(
+      shared_storage.epilogue, 
+      thread_idx, 
+      warp_idx, 
+      lane_idx);
+
+    // Wait on the semaphore - this latency may have been covered by iterator construction
+    if (kSplitKSerial && params.grid_tiled_shape.k() > 1) {
+        
+      // For subsequent threadblocks, the source matrix is held in the 'D' tensor.
+      if (threadblock_tile_offset.k()) {
+        iterator_C = iterator_D;
+      }
+
+      semaphore.wait(threadblock_tile_offset.k());
+
+      __threadfence();
+    }
+
+    // Execute the epilogue operator to update the destination tensor.
+    epilogue(output_op,
+             // Only the final block uses Vector
+             ((kSplitKSerial && params.grid_tiled_shape.k() > 1) &&
+              (params.grid_tiled_shape.k() != threadblock_tile_offset.k() + 1))
+                 ? nullptr
+                 : ptr_Vector,
+             iterator_D,
+             accumulators,
+             iterator_C,
+             iterator_Aux,
+             params.problem_size.mn(),
+             threadblock_offset);
+    
+    //
+    // Release the semaphore
+    //
+
+    if (kSplitKSerial && params.grid_tiled_shape.k() > 1) {
+      
+      int lock = 0;
+      if (params.grid_tiled_shape.k() == threadblock_tile_offset.k() + 1) {
+
+        // The final threadblock resets the semaphore for subsequent grids.
+        lock = 0;
+      }
+      else {
+        // Otherwise, the semaphore is incremented
+        lock = threadblock_tile_offset.k() + 1;
+      }
+
+      __threadfence();
+      semaphore.release(lock);
+    }
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace kernel
+} // namespace gemm
+} // namespace cutlass

include/cutlass/gemm/kernel/static_tile_scheduler.hpp

@@ -334,8 +334,8 @@ public:
 
   // Convert CTA-level work tile info to cluster-level tile coord
   CUTLASS_DEVICE
-  cute::Coord<int,int,int,int>
-  tile_info_to_coord_mnkl(WorkTileInfo work_tile_info) const {
+  auto
+  work_tile_to_cluster_coord_mnkl(WorkTileInfo work_tile_info) const {
     // TileScheduler works at CTA-level, kernel works at cluster-level
     int m_coord = idx2crd(work_tile_info.M_idx / scheduler_params.cluster_shape_m_,
                           scheduler_params.problem_tiles_m_);

include/cutlass/gemm/kernel/symm_universal.h

@@ -119,35 +119,32 @@ public:
     // Data members
     //
 
-    GemmUniversalMode mode;
-    GemmCoord problem_size;
-    int batch_count;
+    GemmUniversalMode mode = GemmUniversalMode::kGemm;
+    GemmCoord problem_size{};
+    int batch_count{1};
 
-    typename EpilogueOutputOp::Params epilogue;
+    typename EpilogueOutputOp::Params epilogue{};
 
-    void const * ptr_A;
-    void const * ptr_B;
-    void const * ptr_C;
-    void * ptr_D;
+    void const * ptr_A{nullptr};
+    void const * ptr_B{nullptr};
+    void const * ptr_C{nullptr};
+    void * ptr_D{nullptr};
 
-    int64_t batch_stride_A;
-    int64_t batch_stride_B;
-    int64_t batch_stride_C;
-    int64_t batch_stride_D;
+    int64_t batch_stride_A{0};
+    int64_t batch_stride_B{0};
+    int64_t batch_stride_C{0};
+    int64_t batch_stride_D{0};
 
-    typename LayoutA::Stride::Index lda;
-    typename LayoutB::Stride::Index ldb;
-    typename LayoutC::Stride::Index ldc;
-    typename LayoutC::Stride::Index ldd;
+    typename LayoutA::Stride::Index lda{0};
+    typename LayoutB::Stride::Index ldb{0};
+    typename LayoutC::Stride::Index ldc{0};
+    typename LayoutC::Stride::Index ldd{0};
 
     //
     // Methods
     //
     
-    Arguments(): 
-      mode(GemmUniversalMode::kGemm), 
-      batch_count(1), 
-      ptr_A(nullptr), ptr_B(nullptr), ptr_C(nullptr), ptr_D(nullptr) { }
+    Arguments() = default;
 
     /// constructs an arguments structure
     Arguments(
@@ -173,7 +170,8 @@ public:
       batch_count(batch_count),
       epilogue(epilogue), 
       ptr_A(ptr_A), ptr_B(ptr_B), ptr_C(ptr_C), ptr_D(ptr_D), 
-      batch_stride_A(batch_stride_A), batch_stride_C(batch_stride_C), batch_stride_D(batch_stride_D), 
+      batch_stride_A(batch_stride_A), batch_stride_B(0),
+      batch_stride_C(batch_stride_C), batch_stride_D(batch_stride_D), 
       lda(lda), ldb(ldb), ldc(ldc), ldd(ldd) {
 
       }
@@ -206,64 +204,43 @@ public:
   /// Parameters structure
   struct Params {
 
-    cutlass::gemm::GemmCoord problem_size;
-    cutlass::gemm::GemmCoord grid_tiled_shape;
-    int swizzle_log_tile;
+    cutlass::gemm::GemmCoord problem_size{};
+    cutlass::gemm::GemmCoord grid_tiled_shape{};
+    int swizzle_log_tile{0};
     
     // Mma1 Iterator A and B params
-    typename Mma1::IteratorA::Params params_A_mma1;
-    typename Mma1::IteratorB::Params params_B_mma1;
+    typename Mma1::IteratorA::Params params_A_mma1{};
+    typename Mma1::IteratorB::Params params_B_mma1{};
 
     // Mma2 Iterator A and B params 
-    typename Mma2::IteratorA::Params params_A_mma2;
-    typename Mma2::IteratorB::Params params_B_mma2;
+    typename Mma2::IteratorA::Params params_A_mma2{};
+    typename Mma2::IteratorB::Params params_B_mma2{};
 
-    typename Epilogue::OutputTileIterator::Params params_C;
-    typename Epilogue::OutputTileIterator::Params params_D;
+    typename Epilogue::OutputTileIterator::Params params_C{};
+    typename Epilogue::OutputTileIterator::Params params_D{};
     
-    typename EpilogueOutputOp::Params output_op;
+    typename EpilogueOutputOp::Params output_op{};
 
-    GemmUniversalMode mode;
-    int batch_count;
-    int gemm_k_size;
+    GemmUniversalMode mode = cutlass::gemm::GemmUniversalMode::kGemm;
+    int batch_count {0};
+    int gemm_k_size {0};
 
-    void * ptr_A;
-    void * ptr_B;
-    void * ptr_C;
-    void * ptr_D;
+    void * ptr_A{nullptr};
+    void * ptr_B{nullptr};
+    void * ptr_C{nullptr};
+    void * ptr_D{nullptr};
 
-    int64_t batch_stride_A;
-    int64_t batch_stride_B;
-    int64_t batch_stride_C;
-    int64_t batch_stride_D;
+    int64_t batch_stride_A {0};
+    int64_t batch_stride_B {0};
+    int64_t batch_stride_C {0};
+    int64_t batch_stride_D {0};
 
-    int *semaphore;
+    int *semaphore{nullptr};
 
     //
     // Methods
     //
-
-    CUTLASS_HOST_DEVICE
-    Params():
-      swizzle_log_tile(0),
-      params_A_mma1(0),
-      params_B_mma1(0),
-      params_A_mma2(0),
-      params_B_mma2(0),
-      params_C(0),
-      params_D(0),
-      batch_count(0),
-      gemm_k_size(0),
-      mode(cutlass::gemm::GemmUniversalMode::kGemm),
-      ptr_A(nullptr),
-      ptr_B(nullptr),
-      ptr_C(nullptr),
-      ptr_D(nullptr),
-      batch_stride_A(0),
-      batch_stride_B(0),
-      batch_stride_C(0),
-      batch_stride_D(0),
-      semaphore(nullptr) { }
+    Params() = default;
 
     CUTLASS_HOST_DEVICE
     Params(

include/cutlass/gemm/kernel/tile_scheduler_params.h

@@ -125,9 +125,9 @@ struct PersistentTileSchedulerSm90Params {
     int max_swizzle_size,
     RasterOrderOptions raster_order_option
   ) {
-    
+
     CUTLASS_UNUSED(hw_info);
-    
+
     // Round up to nearest multiple of swizzle_size along each mode
     auto log_swizzle_size = get_log_swizzle_size(problem_blocks.x, problem_blocks.y, max_swizzle_size);
     auto problem_blocks_m = round_up(problem_blocks.x, (1 << log_swizzle_size) * cluster_shape.m());
@@ -619,9 +619,15 @@ struct PersistentTileSchedulerSm90StreamKParams {
     );
 
     uint64_t ctas_per_wave = grid.x * grid.y;
-
+    auto cluster_size = cluster_shape.m() * cluster_shape.n();
     // The number of output tiles to be computed in stream-K and data-parallel fashion, respectively.
-    uint32_t sk_tiles = get_num_sk_tiles(output_tiles, ctas_per_wave, k_tiles_per_output_tile, decomposition_mode);
+    uint32_t sk_tiles = get_num_sk_tiles(
+      output_tiles,
+      ctas_per_wave,
+      cluster_size,
+      k_tiles_per_output_tile,
+      decomposition_mode
+    );
     uint64_t dp_tiles = output_tiles - sk_tiles;
 
     // Calculate the number of work units covering the data-parallel and stream-K tiles.
@@ -634,24 +640,8 @@ struct PersistentTileSchedulerSm90StreamKParams {
     // is needed per data-parallel tile.
     uint64_t dp_units = dp_tiles;
 
-    // Number of k iterations computed by the stream-K units as a whole
-    uint64_t k_tiles_sk_total = k_tiles_per_output_tile * sk_tiles;
-
-    // If there are stream-K tiles to compute and a sufficiently large number of k iterations
-    // across them, they will be covered by a single wave of persistent threadblocks. Thus, there
-    // will be as many work units as there are threadblocks in a single wave.
-    //
-    // When the total k iterations across stream-K tiles is too small to justify distributing
-    // across an entire wave of blocks, we instead distribute the iterations over a smaller
-    // set of blocks.
-
-    // Calculate the number of stream-K units that would be needed if each stream-K unit
-    // computed the minimum allowable k iterations. Truncate this to be in units of clusters.
-    auto cluster_size = cluster_shape.m() * cluster_shape.n();
-    uint64_t min_sized_sk_units = (k_tiles_sk_total / min_iters_per_sk_unit_);
-    min_sized_sk_units = (min_sized_sk_units / cluster_size) * cluster_size;
-
-    uint64_t sk_units = platform::min(ctas_per_wave, min_sized_sk_units);
+    uint64_t ctas_per_sk_wave = ctas_per_wave;
+    uint64_t sk_units = get_num_sk_units(cluster_shape, ctas_per_sk_wave, sk_tiles, k_tiles_per_output_tile);
 
     if (decomposition_mode == DecompositionMode::DataParallel ||
         (decomposition_mode == DecompositionMode::Heuristic && sk_tiles == 0) ||
@@ -869,7 +859,13 @@ struct PersistentTileSchedulerSm90StreamKParams {
   // Returns the number of stream-K tiles that will be computed amongst `output_tiles` total
   // output tiles on a device with `ctas_per_wave` CTAs in each wave.
   static uint32_t
-  get_num_sk_tiles(uint64_t output_tiles, uint64_t ctas_per_wave, uint32_t k_tiles_per_output_tile, DecompositionMode decomposition_mode) {
+  get_num_sk_tiles(
+    uint64_t output_tiles,
+    uint64_t ctas_per_wave,
+    uint64_t cluster_size,
+    uint32_t k_tiles_per_output_tile,
+    DecompositionMode decomposition_mode
+  ) {
     uint32_t full_waves = static_cast<uint32_t>(output_tiles / ctas_per_wave);
     uint32_t total_waves = static_cast<uint32_t>((output_tiles + ctas_per_wave - 1) / ctas_per_wave);
 
@@ -878,6 +874,14 @@ struct PersistentTileSchedulerSm90StreamKParams {
       return 0;
     }
 
+    // If there is wave quantization, assign the first two waves worth of tiles to be
+    // covered by stream-K work and the remainder to be data-parallel. Since we know
+    // that full_waves == total_waves - 1 in this case, the number of data-parallel
+    // waves is simply full_waves-1 (unless full_waves == 0).
+    uint32_t dp_waves = full_waves > 1 ? full_waves - 1 : 0;
+    uint64_t dp_tiles = dp_waves * ctas_per_wave;
+    uint64_t sk_tiles = output_tiles - dp_tiles;
+
     if (decomposition_mode == DecompositionMode::Heuristic) {
       if (full_waves == total_waves || k_tiles_per_output_tile <= min_iters_per_sk_unit_) {
         // All tiles will be data-parallel tiles if there is either no quantization
@@ -892,24 +896,28 @@ struct PersistentTileSchedulerSm90StreamKParams {
       // Rudimentary heuristic: prefer data-parallel decomposition if we have more than
       // one wave and the tail wave is more than half full. This is subject to change.
       uint64_t tail_tiles = output_tiles - (full_waves * ctas_per_wave);
-      if (tail_tiles >= (ctas_per_wave / 2)) {
+      if (2 * tail_tiles >= ctas_per_wave) {
         return 0;
       }
     }
 
-    // If there is wave quantization, assign the first two waves worth of tiles to be
-    // covered by stream-K work and the remainder to be data-parallel. Since we know
-    // that full_waves == total_waves - 1 in this case, the number of data-parallel
-    // waves is simply full_waves-1 (unless full_waves == 0).
-    uint32_t dp_waves = full_waves > 0 ? full_waves - 1 : 0;
-
-    uint64_t dp_tiles = dp_waves * ctas_per_wave;
-    return static_cast<uint32_t>(output_tiles - dp_tiles);
+    return static_cast<uint32_t>(sk_tiles);
   }
 
   CUTLASS_HOST_DEVICE
   static uint64_t
-  get_num_sk_units(GemmCoord cluster_shape, uint64_t ctas_per_wave, uint32_t sk_tiles, uint32_t k_tiles_per_output_tile) {
+  get_num_sk_units(GemmCoord cluster_shape, uint64_t ctas_per_sk_wave, uint32_t sk_tiles, uint32_t k_tiles_per_output_tile) {
+    // If there are stream-K tiles to compute and a sufficiently large number of k iterations
+    // across them, they will be covered by a single wave of persistent threadblocks. Thus, there
+    // will be as many work units as there are threadblocks in a single wave.
+    //
+    // When the total k iterations across stream-K tiles is too small to justify distributing
+    // across an entire wave of blocks, we instead distribute the iterations over a smaller
+    // set of blocks.
+
+    // Calculate the number of stream-K units that would be needed if each stream-K unit
+    // computed the minimum allowable k iterations. Truncate this to be in units of clusters.
+
     // Number of k iterations computed by the stream-K units as a whole
     uint64_t k_tiles_sk_total = k_tiles_per_output_tile * sk_tiles;
 
@@ -919,7 +927,7 @@ struct PersistentTileSchedulerSm90StreamKParams {
     uint64_t min_sized_sk_units = (k_tiles_sk_total / min_iters_per_sk_unit_);
     min_sized_sk_units = (min_sized_sk_units / cluster_size) * cluster_size;
 
-    uint64_t sk_units = platform::min(ctas_per_wave, min_sized_sk_units);
+    uint64_t sk_units = platform::min(ctas_per_sk_wave, min_sized_sk_units);
     return sk_units;
   }
 
@@ -934,9 +942,9 @@ struct PersistentTileSchedulerSm90StreamKParams {
   // Calculates the size of the workspace needed for holding partial outputs from splits
   CUTLASS_HOST_DEVICE
   static int
-  get_reduction_workspace_size(uint64_t num_tiles, GemmCoord tile_shape, uint32_t accumulator_bits) {
+  get_reduction_workspace_size(uint64_t num_tiles, GemmCoord tile_shape, uint32_t accumulator_bits, uint32_t num_accumulator_mtxs = 1) {
     auto output_tile_size = tile_shape.m() * tile_shape.n();
-    auto workspace_bits = accumulator_bits * output_tile_size * num_tiles;
+    auto workspace_bits = accumulator_bits * output_tile_size * num_tiles * num_accumulator_mtxs;
     return round_up_to_l2_alignment(bits_to_bytes(static_cast<int>(workspace_bits)));
   }
 
@@ -957,7 +965,8 @@ struct PersistentTileSchedulerSm90StreamKParams {
     uint32_t mma_warp_groups,
     uint32_t barrier_bits,
     uint32_t accumulator_bits,
-    uint32_t epilogue_subtile = 1) {
+    uint32_t epilogue_subtile = 1,
+    uint32_t num_accumulator_mtxs = 1) {
 
     auto log_swizzle_size = UnderlyingParams::get_log_swizzle_size(problem_blocks.x, problem_blocks.y, max_swizzle);
     problem_blocks.x = round_up(problem_blocks.x, (1 << log_swizzle_size) * cluster_shape.m());
@@ -975,7 +984,7 @@ struct PersistentTileSchedulerSm90StreamKParams {
         (decomposition_mode == DecompositionMode::Heuristic && splits > 1)) {
       // Basic split-K variant requires workspace for all output tiles
       barrier_workspace_size = get_barrier_workspace_size(output_tiles, mma_warp_groups, barrier_bits);
-      reduction_workspace_size = get_reduction_workspace_size(output_tiles, tile_shape, accumulator_bits);
+      reduction_workspace_size = get_reduction_workspace_size(output_tiles, tile_shape, accumulator_bits, num_accumulator_mtxs);
     }
     else {
       KernelHardwareInfo new_hw_info;
@@ -995,8 +1004,16 @@ struct PersistentTileSchedulerSm90StreamKParams {
         raster_order_option
       );
       uint64_t ctas_per_wave = grid.x * grid.y;
-      uint32_t sk_tiles = get_num_sk_tiles(output_tiles, ctas_per_wave, static_cast<uint32_t>(k_tiles_per_output_tile), decomposition_mode);
-      uint64_t sk_units = get_num_sk_units(cluster_shape, ctas_per_wave, sk_tiles, k_tiles_per_output_tile);
+      uint64_t cluster_size = cluster_shape.m() * cluster_shape.n();
+      uint32_t sk_tiles = get_num_sk_tiles(
+        output_tiles,
+        ctas_per_wave,
+        cluster_size,
+        static_cast<uint32_t>(k_tiles_per_output_tile),
+        decomposition_mode
+      );
+      uint64_t ctas_per_sk_wave = ctas_per_wave;
+      uint64_t sk_units = get_num_sk_units(cluster_shape, ctas_per_sk_wave, sk_tiles, k_tiles_per_output_tile);
       uint64_t dp_tiles = output_tiles - sk_tiles;
 
       uint64_t reduction_tiles = sk_tiles;
@@ -1012,7 +1029,7 @@ struct PersistentTileSchedulerSm90StreamKParams {
       // written its accumulator to scratch space. The separate reduction unit will only begin
       // performing the reduction when the barrier has reached the number of peers for the output tile.
       barrier_workspace_size = get_barrier_workspace_size(sk_tiles, mma_warp_groups, barrier_bits);
-      reduction_workspace_size = get_reduction_workspace_size(reduction_tiles, tile_shape, accumulator_bits);
+      reduction_workspace_size = get_reduction_workspace_size(reduction_tiles, tile_shape, accumulator_bits, num_accumulator_mtxs);
     }
   }
   #endif // !defined(__CUDACC_RTC__)
@@ -1030,7 +1047,7 @@ struct PersistentTileSchedulerSm90StreamKParams {
 
   // Get the amount of scratch workspace needed for the kernel. This variant of the method should only be used when
   // problem_shape and tile_shape contain modes of only rank 1.
-  static int
+  static size_t
   get_workspace_size(
     BatchedGemmCoord problem_shape,
     GemmCoord tile_shape,
@@ -1043,7 +1060,8 @@ struct PersistentTileSchedulerSm90StreamKParams {
     uint32_t mma_warp_groups,
     uint32_t barrier_bits,
     uint32_t element_accumulator_bits,
-    uint32_t epilogue_subtile) {
+    uint32_t epilogue_subtile,
+    uint32_t num_accumulator_mtxs) {
 
     dim3 problem_blocks = UnderlyingParams::get_tiled_cta_shape_mnl(problem_shape, tile_shape, cluster_shape);
     uint32_t k_tiles_per_output_tile = (problem_shape.k() + tile_shape.k() - 1) / tile_shape.k();
@@ -1061,14 +1079,15 @@ struct PersistentTileSchedulerSm90StreamKParams {
       mma_warp_groups,
       barrier_bits,
       element_accumulator_bits,
-      epilogue_subtile
+      epilogue_subtile,
+      num_accumulator_mtxs
     );
   }
 
   // Version of get_workspace_size that takes in as input the number of CTAs in the M and N dimensions.
   // This is useful for calculating the tiled shape when a mode of problem and/or CTA shape has rank > 1,
   // for which using CuTe algebra for calculating tile shapes is easiest.
-  static int
+  static size_t
   get_workspace_size(
     dim3 problem_blocks,
     uint32_t k_tiles_per_output_tile,
@@ -1082,7 +1101,8 @@ struct PersistentTileSchedulerSm90StreamKParams {
     uint32_t mma_warp_groups,
     uint32_t barrier_bits,
     uint32_t element_accumulator_bits,
-    uint32_t epilogue_subtile = 1) {
+    uint32_t epilogue_subtile = 1,
+    uint32_t num_accumulator_mtxs = 1) {
 
     int barrier_workspace_size = 0;
     int reduction_workspace_size = 0;
@@ -1103,7 +1123,8 @@ struct PersistentTileSchedulerSm90StreamKParams {
         mma_warp_groups,
         barrier_bits,
         element_accumulator_bits,
-        epilogue_subtile
+        epilogue_subtile,
+        num_accumulator_mtxs
       );
     #endif
 
@@ -1170,7 +1191,8 @@ struct PersistentTileSchedulerSm90StreamKParams {
     uint32_t mma_warp_groups,
     uint32_t barrier_bits,
     uint32_t element_accumulator_bits,
-    uint32_t epilogue_subtile = 1) {
+    uint32_t epilogue_subtile = 1,
+    uint32_t num_accumulator_mtxs = 1) {
 
     #if !defined(__CUDACC_RTC__)
       int barrier_workspace_size = 0;
@@ -1191,7 +1213,8 @@ struct PersistentTileSchedulerSm90StreamKParams {
         mma_warp_groups,
         barrier_bits,
         element_accumulator_bits,
-        epilogue_subtile
+        epilogue_subtile,
+        num_accumulator_mtxs
       );
 
       if (barrier_workspace_size > 0) {
@@ -1301,9 +1324,9 @@ struct PersistentTileSchedulerSm90GroupParams {
     int max_swizzle_size,
     RasterOrderOptions raster_order_option
   ) {
-    
+
     CUTLASS_UNUSED(hw_info);
-    
+
     // Round up to nearest multiple of swizzle_size along each mode
     auto log_swizzle_size = get_log_swizzle_size(problem_blocks.x, problem_blocks.y, max_swizzle_size);
     auto problem_blocks_m = round_up(problem_blocks.x, (1 << log_swizzle_size) * cluster_shape.m());

include/cutlass/gemm/kernel/trmm_universal.h

@@ -110,32 +110,29 @@ public:
     // Data members
     //
 
-    GemmUniversalMode mode;
-    GemmCoord problem_size;
-    int batch_count;
+    GemmUniversalMode mode{GemmUniversalMode::kGemm};
+    GemmCoord problem_size{};
+    int batch_count{1};
 
-    typename EpilogueOutputOp::Params epilogue;
+    typename EpilogueOutputOp::Params epilogue{};
 
-    void const * ptr_A;
-    void const * ptr_B;
-    void * ptr_D;
+    void const * ptr_A{nullptr};
+    void const * ptr_B{nullptr};
+    void * ptr_D{nullptr};
 
-    int64_t batch_stride_A;
-    int64_t batch_stride_B;
-    int64_t batch_stride_D;
+    int64_t batch_stride_A{0};
+    int64_t batch_stride_B{0};
+    int64_t batch_stride_D{0};
 
-    typename LayoutA::Stride::Index lda;
-    typename LayoutB::Stride::Index ldb;
-    typename LayoutC::Stride::Index ldd;
+    typename LayoutA::Stride::Index lda{0};
+    typename LayoutB::Stride::Index ldb{0};
+    typename LayoutC::Stride::Index ldd{0};
 
     //
     // Methods
     //
-    
-    Arguments(): 
-      mode(GemmUniversalMode::kGemm), 
-      batch_count(1), 
-      ptr_A(nullptr), ptr_B(nullptr), ptr_D(nullptr) { }
+
+    Arguments() = default;
 
     /// constructs an arguments structure
     Arguments(
@@ -161,7 +158,7 @@ public:
       batch_stride_A(batch_stride_A), batch_stride_B(batch_stride_B), batch_stride_D(batch_stride_D), 
       lda(lda), ldb(ldb), ldd(ldd) {
       }
-
+    
     /// Returns arguments for the transposed problem sizes
     Arguments transposed_problem_size() const {
       Arguments args(*this);
@@ -190,50 +187,34 @@ public:
   /// Parameters structure
   struct Params {
 
-    cutlass::gemm::GemmCoord problem_size;
-    cutlass::gemm::GemmCoord grid_tiled_shape;
-    int swizzle_log_tile;
+    cutlass::gemm::GemmCoord problem_size{};
+    cutlass::gemm::GemmCoord grid_tiled_shape{};
+    int swizzle_log_tile{0};
    
-    typename Mma::IteratorA::Params params_A;
-    typename Mma::IteratorB::Params params_B;
-    typename Epilogue::OutputTileIterator::Params params_D;
+    typename Mma::IteratorA::Params params_A{};
+    typename Mma::IteratorB::Params params_B{};
+    typename Epilogue::OutputTileIterator::Params params_D{};
     
-    typename EpilogueOutputOp::Params output_op;
+    typename EpilogueOutputOp::Params output_op{};
 
-    GemmUniversalMode mode;
-    int batch_count;
-    int gemm_k_size;
+    GemmUniversalMode mode = cutlass::gemm::GemmUniversalMode::kGemm;
+    int batch_count {0};
+    int gemm_k_size {0};
 
-    void * ptr_A;
-    void * ptr_B;
-    void * ptr_D;
+    void * ptr_A{nullptr};
+    void * ptr_B{nullptr};
+    void * ptr_D{nullptr};
 
-    int64_t batch_stride_A;
-    int64_t batch_stride_B;
-    int64_t batch_stride_D;
+    int64_t batch_stride_A {0};
+    int64_t batch_stride_B {0};
+    int64_t batch_stride_D {0};
 
-    int *semaphore;
+    int *semaphore{nullptr};
 
     //
     // Methods
     //
-
-    CUTLASS_HOST_DEVICE
-    Params():
-      swizzle_log_tile(0),
-      params_A(0),
-      params_B(0),
-      params_D(0),
-      batch_count(0),
-      gemm_k_size(0),
-      mode(cutlass::gemm::GemmUniversalMode::kGemm),
-      ptr_A(nullptr),
-      ptr_B(nullptr),
-      ptr_D(nullptr),
-      batch_stride_A(0),
-      batch_stride_B(0),
-      batch_stride_D(0),
-      semaphore(nullptr) { }
+    Params() = default;
 
     CUTLASS_HOST_DEVICE
     Params(