GEMM API

template <
  /// Element type for A matrix operand
  typename ElementA_,
  /// Layout type for A matrix operand
  typename LayoutA_,
  /// Element type for B matrix operand
  typename ElementB_,
  /// Layout type for B matrix operand
  typename LayoutB_,
  /// 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_ = ElementC_,
  /// Operator class tag
  typename OperatorClass_ = arch::OpClassSimt,
  /// Tag indicating architecture to tune for
  typename ArchTag_ = arch::Sm70,
  /// Threadblock-level tile size (concept: GemmShape)
  typename ThreadblockShape_ = /* default based on config */,
  /// Warp-level tile size (concept: GemmShape)
  typename WarpShape_ = /* default based on config */,
  /// Instruction-level tile size (concept: GemmShape)
  typename InstructionShape_ = /* default based on config */,
  /// Epilogue output operator
  typename EpilogueOutputOp_ = /* default based on config */,
  /// Threadblock-level swizzling operator
  typename ThreadblockSwizzle_ = threadblock::GemmIdentityThreadblockSwizzle<>,
  /// Number of stages used in the pipelined mainloop
  int Stages = /* default based on config */,
  /// Access granularity of A matrix in units of elements
  int AlignmentA = /* default based on config */,
  /// Access granularity of B matrix in units of elements
  int AlignmentB = /* default based on config */,
  /// If true, kernel supports split-K with serial reduction
  bool SplitKSerial = false,
  /// Operation performed by GEMM
  typename Operator_ = /* default based on config */,
  /// Gather operand A by using an index array
  bool GatherA = false,
  /// Gather operand B by using an index array
  bool GatherB = false,
  /// Scatter result D by using an index array
  bool ScatterD = false,
  /// Permute result D
  typename PermuteDLayout = layout::NoPermute
>
class Gemm;

struct Arguments {
  GemmCoord problem_size;
  TensorRef<ElementA const, LayoutA> ref_A;
  TensorRef<ElementB const, LayoutB> ref_B;
  TensorRef<ElementC const, LayoutC> ref_C;
  TensorRef<ElementC, LayoutC> ref_D;
  typename EpilogueOutputOp::Params epilogue;
  int split_k_slices;
  // For gather+scatter operations
  int const *gather_A_indices;
  int const *gather_B_indices;
  int const *scatter_D_indices;
};

int split_k_slices = 4;  // Partition K dimension

CutlassGemm::Arguments args(
  {M, N, K},
  {A, lda},
  {B, ldb},
  {C, ldc},
  {D, ldd},
  {alpha, beta},
  split_k_slices  // Enable split-K
);

// Get workspace size
size_t workspace_size = CutlassGemm::get_workspace_size(args);
void* workspace;
cudaMalloc(&workspace, workspace_size);

// Launch with workspace
CutlassGemm gemm_op;
gemm_op(args, workspace, stream);

#include "cutlass/gemm/device/gemm_batched.h"

using GemmBatched = cutlass::gemm::device::GemmBatched<
  float, cutlass::layout::RowMajor,
  float, cutlass::layout::RowMajor,
  float, cutlass::layout::RowMajor
>;

GemmBatched::Arguments args(
  {M, N, K},
  {A, lda},
  lda * M,          // batch_stride_A
  {B, ldb},
  ldb * K,          // batch_stride_B
  {C, ldc},
  ldc * M,          // batch_stride_C
  {D, ldd},
  ldd * M,          // batch_stride_D
  {alpha, beta},
  batch_count
);

CutlassGemm gemm_op;

// Check if problem is supported
Status status = CutlassGemm::can_implement(args);
if (status != Status::kSuccess) {
  std::cerr << "Problem not supported: " 
            << cutlass::cutlassGetStatusString(status) 
            << std::endl;
  return;
}

// Initialize and run
status = gemm_op.initialize(args, workspace, stream);
if (status != Status::kSuccess) {
  std::cerr << "Initialization failed" << std::endl;
  return;
}

status = gemm_op.run(stream);
if (status != Status::kSuccess) {
  std::cerr << "Kernel execution failed" << std::endl;
  return;
}