CUTLASS 2.2 (#96)

Adds support for NVIDIA Ampere Architecture features. CUDA 11 Toolkit recommended.

tools/library/src/gemm_operation.h

@@ -1,5 +1,5 @@
 /***************************************************************************************************
- * Copyright (c) 2017-2019, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2017-2020, NVIDIA CORPORATION.  All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without modification, are permitted
  * provided that the following conditions are met:
@@ -29,13 +29,14 @@
 #pragma once
 
 #include "cutlass/cutlass.h"
-#include "cutlass/gemm/kernel/default_gemm_planar_complex_universal.h"
 
 #include "cutlass/gemm/device/gemm.h"
 #include "cutlass/gemm/device/gemm_complex.h"
 #include "cutlass/gemm/device/gemm_batched.h"
 #include "cutlass/gemm/device/gemm_array.h"
 #include "cutlass/gemm/device/gemm_universal_adapter.h"
+#include "cutlass/gemm/kernel/default_gemm_universal.h"
+#include "cutlass/gemm/kernel/default_gemm_planar_complex_universal.h"
 
 #include "cutlass/library/library.h"
 #include "library_internal.h"
@@ -104,10 +105,10 @@ public:
       MathOperationMap<typename Operator::Operator>::kId;
 
     description_.tile_description.minimum_compute_capability = 
-      ArchMap<typename Operator::ArchTag>::kMin;
+      ArchMap<typename Operator::ArchTag, typename Operator::OperatorClass>::kMin;
 
     description_.tile_description.maximum_compute_capability = 
-      ArchMap<typename Operator::ArchTag>::kMax;
+      ArchMap<typename Operator::ArchTag, typename Operator::OperatorClass>::kMax;
     
     description_.A = make_TensorDescription<ElementA, LayoutA>(Operator::kAlignmentA);
     description_.B = make_TensorDescription<ElementB, LayoutB>(Operator::kAlignmentB);
@@ -698,6 +699,201 @@ public:
   }
 };
 
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename Operator_>
+class GemmUniversalOperation : public GemmOperationBase<Operator_> {
+public:
+
+  using Operator = Operator_;
+  using ElementA = typename Operator::ElementA;
+  using LayoutA = typename Operator::LayoutA;
+  using ElementB = typename Operator::ElementB;
+  using LayoutB = typename Operator::LayoutB;
+  using ElementC = typename Operator::ElementC;
+  using LayoutC = typename Operator::LayoutC;
+  using ElementAccumulator = typename Operator::ElementAccumulator;
+  using ElementCompute = typename Operator::EpilogueOutputOp::ElementCompute;
+
+  using OperatorArguments = typename Operator::Arguments;
+
+public:
+
+  /// Constructor
+  GemmUniversalOperation(char const *name = "unknown_gemm"): 
+    GemmOperationBase<Operator_>(name) {
+
+    this->description_.gemm_kind = GemmKind::kUniversal;
+  }
+
+protected:
+
+  /// Constructs the arguments structure given the configuration and arguments
+  static Status construct_arguments_(
+    OperatorArguments &operator_args,
+    GemmUniversalConfiguration const *configuration) {
+
+    operator_args.mode = configuration->mode;
+
+    operator_args.problem_size = configuration->problem_size;
+    operator_args.batch_count = configuration->batch_count;
+
+    operator_args.lda = int(configuration->lda);
+    operator_args.ldb = int(configuration->ldb);
+    operator_args.ldc = int(configuration->ldc);
+    operator_args.ldd = int(configuration->ldd);
+
+    return Status::kSuccess;
+  }
+
+  /// Constructs the arguments structure given the configuration and arguments
+  static Status update_arguments_(
+    OperatorArguments &operator_args,
+    GemmUniversalArguments const *arguments) {
+    
+    if (arguments->pointer_mode == ScalarPointerMode::kHost) {
+      typename Operator::EpilogueOutputOp::Params params(
+        *static_cast<ElementCompute const *>(arguments->alpha),
+        *static_cast<ElementCompute const *>(arguments->beta)
+      );
+      operator_args.epilogue = params;
+    }
+    else if (arguments->pointer_mode == ScalarPointerMode::kDevice){
+      typename Operator::EpilogueOutputOp::Params params(
+        static_cast<ElementCompute const *>(arguments->alpha),
+        static_cast<ElementCompute const *>(arguments->beta)
+      );
+      operator_args.epilogue = params; 
+    }
+    else {
+      return Status::kErrorInvalidProblem;
+    }
+
+    // update arguments
+    operator_args.ptr_A = arguments->A;
+    operator_args.ptr_B = arguments->B;
+    operator_args.ptr_C = arguments->C;
+    operator_args.ptr_D = arguments->D;
+
+    operator_args.batch_stride_A = arguments->batch_stride_A;
+    operator_args.batch_stride_B = arguments->batch_stride_B;
+    operator_args.batch_stride_C = arguments->batch_stride_C;
+    operator_args.batch_stride_D = arguments->batch_stride_D;
+    
+    return Status::kSuccess;
+  }
+
+public:
+
+  /// Returns success if the operation can proceed
+  virtual Status can_implement(
+    void const *configuration_ptr, 
+    void const *arguments_ptr) const {
+    
+    GemmUniversalConfiguration const *configuration = 
+      static_cast<GemmUniversalConfiguration const *>(configuration_ptr);
+
+    GemmUniversalArguments const *arguments = 
+      static_cast<GemmUniversalArguments const *>(arguments_ptr);
+
+    OperatorArguments args;
+
+    Status status = construct_arguments_(args, configuration);
+
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    status = update_arguments_(args, arguments);
+
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    return Operator::can_implement(args);
+  }
+  
+  /// Gets the host-side workspace
+  virtual uint64_t get_host_workspace_size(
+    void const *configuration) const {
+
+    return sizeof(Operator);
+  }
+  
+  /// Gets the device-side workspace
+  virtual uint64_t get_device_workspace_size(
+    void const *configuration_ptr) const {
+
+    OperatorArguments args;
+
+    Status status = construct_arguments_(
+      args, 
+      static_cast<GemmUniversalConfiguration const *>(configuration_ptr));
+
+    if (status != Status::kSuccess) {
+      return 0;
+    }
+
+    uint64_t size = Operator::get_workspace_size(args);
+
+    return size;
+  }
+  
+  /// Initializes the workspace
+  virtual Status initialize(
+    void const *configuration_ptr, 
+    void *host_workspace, 
+    void *device_workspace, 
+    cudaStream_t stream = nullptr) const {
+
+    OperatorArguments args;
+
+    Status status = construct_arguments_(
+      args, 
+      static_cast<GemmUniversalConfiguration const *>(configuration_ptr));
+
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    Operator *op = new (host_workspace) Operator;
+
+    status = op->initialize(args, device_workspace, stream);
+    
+    return status;
+  }
+
+  /// Runs the kernel
+  virtual Status run(
+    void const *arguments_ptr,
+    void *host_workspace, 
+    void *device_workspace = nullptr, 
+    cudaStream_t stream = nullptr) const {
+
+    OperatorArguments args;
+    
+    Status status = update_arguments_(
+      args, 
+      static_cast<GemmUniversalArguments const *>(arguments_ptr));
+
+    if (status != Status::kSuccess) {
+      return status;
+    }
+    
+    Operator *op = static_cast<Operator *>(host_workspace);
+    
+    status = op->update(args, device_workspace);
+
+    if (status != Status::kSuccess) {
+      return status;
+    }
+    
+    status = op->run(stream);
+    
+    return status;
+  }
+};
+
 ///////////////////////////////////////////////////////////////////////////////////////////////////
 
 template <typename Operator_>

tools/library/src/handle.cu

@@ -26,7 +26,7 @@
 /*! \file
     \brief CUTLASS Library handle.
 */
-
+#include <iostream> 
 #include <stdexcept>
 #include <cstdint>
 
@@ -43,7 +43,8 @@ namespace library {
 Handle::Handle(
   cudaStream_t stream, 
   size_t workspace_size
-): 
+):
+  provider_(Provider::kCUTLASS), 
   stream_(stream), 
   workspace_(nullptr), 
   workspace_size_(0), 
@@ -95,6 +96,7 @@ Handle::Handle(Handle && handle) {
 /// Move assignment operator
 Handle & Handle::operator=(Handle && handle) {
 
+  provider_ = handle.provider_;
   device_ = handle.device_;
   workspace_size_ = handle.workspace_size_;
   workspace_ = handle.workspace_;
@@ -121,6 +123,16 @@ cudaStream_t Handle::get_stream() const {
   return stream_;
 }
 
+/// Gets the current provider
+Provider Handle::get_provider() const {
+  return provider_;
+}
+
+/// Sets the provider of operations
+void Handle::set_provider(Provider provider) {
+  provider_ = provider;
+}
+
 /// Gets the device workspace size
 size_t Handle::get_workspace_size() const {
   return workspace_size_;
@@ -351,6 +363,8 @@ Status Handle::gemm(
   //
 
   GemmFunctionalKey key(
+    provider_,
+    GemmKind::kGemm,
     element_compute,
     element_scalar,
     element_A,
@@ -457,6 +471,188 @@ Status Handle::gemm(
 
 ///////////////////////////////////////////////////////////////////////////////////////////////////
 
+/// Executes a GEMM computation: D <= alpha * A*B + beta * C.
+//
+// Supports batched-strided, batched array or split-K serial or split-K parallel.
+//
+Status Handle::gemm_universal(
+
+  GemmUniversalMode mode,                   /// indicates the mode in which the kUniversal GEMM is launched
+
+  int M,                                    /// GEMM M dimension
+  int N,                                    /// GEMM N dimension
+  int K,                                    /// GEMM K dimension
+
+  NumericTypeID element_compute,            /// Data type of internal accumulation
+
+  NumericTypeID element_scalar,             /// Data type of alpha/beta scalars
+
+  void const *alpha,                        /// Pointer to alpha scalar
+
+  NumericTypeID element_A,                  /// Data type of A matrix elements
+  LayoutTypeID layout_A,                    /// Layout of A matrix
+  ComplexTransform transform_A,             /// Complex transformation applied to A matrix - ignored for real-valued matrices
+
+  void const * ptr_A,                       /// Pointer to A matrix in Global Memory
+  int lda,                                  /// Leading dimension of A matrix
+
+  NumericTypeID element_B,                  /// Data type of B matrix elements
+  LayoutTypeID layout_B,                    /// Layout of B matrix
+  ComplexTransform transform_B,             /// Complex transformation applied to B matrix - ignored for real-valued matrices
+
+  void const * ptr_B,                       /// Pointer to B matrix in Global Memory
+  int ldb,                                  /// Leading dimension of B matrix
+
+  void const * beta,                        /// Pointer to beta scalar
+
+  NumericTypeID element_C,                  /// Data type of C and D matrices
+
+  void const * ptr_C,                       /// Pointer to C matrix
+  int ldc,                                  /// Leading dimension of C matrix
+
+  void * ptr_D,                             /// Pointer to D matrix
+  int ldd,                                  /// Leading dimension of D matrix
+
+  int batch_count,                          /// Batch count or number of split-K slices
+
+  int64_t batch_stride_A,                   /// Batch stride of A operand
+  int64_t batch_stride_B,                   /// Batch stride of B operand
+  int64_t batch_stride_C,                   /// Batch stride of C operand
+  int64_t batch_stride_D                    /// Batch stride of D operand
+) {
+  
+  //
+  // Find the operation
+  //
+
+  GemmFunctionalKey key(
+    provider_,
+    GemmKind::kUniversal,
+    element_compute,
+    element_scalar,
+    element_A,
+    layout_A,
+    transform_A,
+    element_B,
+    layout_B,
+    transform_B,
+    element_C
+  );
+
+  auto operators_it = Singleton::get().operation_table.gemm_operations.find(key);
+
+  if (operators_it == Singleton::get().operation_table.gemm_operations.end()) {
+    return cutlass::Status::kErrorNotSupported;
+  }
+  
+  if (operators_it->second.empty()) {
+    return cutlass::Status::kErrorNotSupported;
+  }
+
+  //
+  // Compute the largest alignment restriction the kernel can satisfy.
+  //
+
+  // Maximum alignment expectation among all kernels (in units of bytes)
+  int const kMaximumAlignmentSize = 16;
+
+  void const *ptr_A_check = ptr_A;
+  void const *ptr_B_check = ptr_B;
+  void const *ptr_C_check = ptr_C;
+  void *      ptr_D_check = ptr_D;
+
+  // Ignore alignment of pointers to pointers. We can't check this from the host,
+  // as each batch index has its own pointer in device memory.
+  if (mode == GemmUniversalMode::kArray) {
+    ptr_A_check = nullptr; 
+    ptr_B_check = nullptr; 
+    ptr_C_check = nullptr; 
+    ptr_D_check = nullptr; 
+  }
+
+  int alignment = gemm_problem_alignment(
+    M, N, K, 
+    element_A, ptr_A_check, lda, 0,
+    element_B, ptr_B_check, ldb, 0,
+    element_C, ptr_C_check, ldc, 0,
+    ptr_D_check, ldd, 0, kMaximumAlignmentSize
+  );
+
+  //
+  // Find the best kernel in descending order of preference.
+  //
+
+  GemmPreferenceKey preference_key(compute_capability(), alignment);
+
+  Operation const *operation = find_gemm_operation(operators_it, preference_key);
+
+  if (!operation) {
+    return cutlass::Status::kErrorNotSupported;
+  }
+
+  last_operation_ = operation;
+
+  //
+  // Configure operation
+  //
+
+  GemmUniversalConfiguration configuration{
+    mode,
+    {M, N, K},
+    batch_count,
+    lda,
+    ldb,
+    ldc,
+    ldd
+  };
+
+  // Query host work space size
+  uint64_t host_workspace_size_needed = operation->get_host_workspace_size(&configuration);
+
+  if (uint64_t(kHostWorkspaceSize) < host_workspace_size_needed) {
+    return cutlass::Status::kErrorNotSupported;
+  }
+
+  char host_workspace[kHostWorkspaceSize];
+
+  // Query device workspace size
+  uint64_t device_workspace_size_needed = operation->get_device_workspace_size(&configuration);
+
+  if (uint64_t(workspace_size_) < device_workspace_size_needed) {
+    return cutlass::Status::kErrorNotSupported;
+  }
+
+  // Initialize host and device workspaces
+  Status status = operation->initialize(
+    &configuration,
+    host_workspace,
+    workspace_,
+    stream_);
+
+  if (status != cutlass::Status::kSuccess) {
+    return status;
+  }
+
+  // Run the operator
+  GemmUniversalArguments arguments{
+    ptr_A,
+    ptr_B,
+    ptr_C,
+    ptr_D,
+    alpha,
+    beta,
+    scalar_pointer_mode_,
+    batch_stride_A,
+    batch_stride_B,
+    batch_stride_C,
+    batch_stride_D
+  };
+
+  return operation->run(&arguments, host_workspace, workspace_, stream_);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
 /// Planar complex GEMM
 Status Handle::gemm_planar_complex(
 
@@ -522,6 +718,8 @@ Status Handle::gemm_planar_complex(
   //
 
   GemmFunctionalKey key(
+    provider_,
+    GemmKind::kPlanarComplex,
     element_compute,
     element_scalar,
     element_A,
@@ -533,9 +731,9 @@ Status Handle::gemm_planar_complex(
     element_C
   );
 
-  auto operators_it = Singleton::get().operation_table.gemm_planar_complex_operations.find(key);
+  auto operators_it = Singleton::get().operation_table.gemm_operations.find(key);
 
-  if (operators_it == Singleton::get().operation_table.gemm_planar_complex_operations.end()) {
+  if (operators_it == Singleton::get().operation_table.gemm_operations.end()) {
     return cutlass::Status::kErrorNotSupported;
   }
   
@@ -714,6 +912,8 @@ Status Handle::gemm_planar_complex_array(
   //
 
   GemmFunctionalKey key(
+    provider_,
+    GemmKind::kPlanarComplexArray,
     element_compute,
     element_scalar,
     element_A,
@@ -725,9 +925,9 @@ Status Handle::gemm_planar_complex_array(
     element_C
   );
 
-  auto operators_it = Singleton::get().operation_table.gemm_planar_complex_array_operations.find(key);
+  auto operators_it = Singleton::get().operation_table.gemm_operations.find(key);
 
-  if (operators_it == Singleton::get().operation_table.gemm_planar_complex_array_operations.end()) {
+  if (operators_it == Singleton::get().operation_table.gemm_operations.end()) {
     return cutlass::Status::kErrorNotSupported;
   }
   
@@ -837,7 +1037,6 @@ Status Handle::gemm_planar_complex_array(
 }
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
-
 } // namespace library
 } // namespace cutlass
 

tools/library/src/library_internal.h

@@ -1,5 +1,5 @@
 /***************************************************************************************************
- * Copyright (c) 2019, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2017-2020, NVIDIA CORPORATION.  All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without modification, are permitted
  * provided that the following conditions are met:
@@ -125,6 +125,14 @@ template <> struct NumericTypeMap<cutlass::complex<double> > {
   static NumericTypeID const kId = NumericTypeID::kCF64;
 };
 
+template <> struct NumericTypeMap<cutlass::bfloat16_t> {
+  static NumericTypeID const kId = NumericTypeID::kBF16;
+};
+
+template <> struct NumericTypeMap<cutlass::tfloat32_t> {
+  static NumericTypeID const kId = NumericTypeID::kTF32;
+};
+
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 template <typename T> struct MathOperationMap {
@@ -143,6 +151,10 @@ template <> struct MathOperationMap<cutlass::arch::OpMultiplyAddComplex> {
   static MathOperationID const kId = MathOperationID::kMultiplyAddComplex;
 };
 
+template <> struct MathOperationMap<cutlass::arch::OpMultiplyAddGaussianComplex> {
+  static MathOperationID const kId = MathOperationID::kMultiplyAddGaussianComplex;
+};
+
 template <> struct MathOperationMap<cutlass::arch::OpXorPopc> {
   static MathOperationID const kId = MathOperationID::kXorPopc;
 };
@@ -217,33 +229,43 @@ template <> struct ComplexTransformMap<cutlass::ComplexTransform::kConjugate> {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-template <typename T> struct ArchMap;
+template <typename ArchTag, typename OperatorClass> struct ArchMap;
 
-template <> struct ArchMap<arch::Sm50> {
+template <> struct ArchMap<arch::Sm50, arch::OpClassSimt> {
   static int const kMin = 50;
   static int const kMax = 1024;
 };
 
-template <> struct ArchMap<arch::Sm60> {
+template <> struct ArchMap<arch::Sm60, arch::OpClassSimt> {
   static int const kMin = 60;
   static int const kMax = 1024;
 };
 
-template <> struct ArchMap<arch::Sm61> {
+template <> struct ArchMap<arch::Sm61, arch::OpClassSimt> {
   static int const kMin = 61;
   static int const kMax = 1024;
 };
 
-template <> struct ArchMap<arch::Sm70> {
+template <> struct ArchMap<arch::Sm70, arch::OpClassWmmaTensorOp> {
+  static int const kMin = 70;
+  static int const kMax = 1024;
+};
+
+template <> struct ArchMap<arch::Sm70, arch::OpClassTensorOp> {
   static int const kMin = 70;
   static int const kMax = 75;
 };
 
-template <> struct ArchMap<arch::Sm75> {
+template <typename OperatorClass> struct ArchMap<arch::Sm75, OperatorClass> {
   static int const kMin = 75;
   static int const kMax = 1024;
 };
 
+template <typename OperatorClass> struct ArchMap<arch::Sm80, OperatorClass> {
+  static int const kMin = 80;
+  static int const kMax = 1024;
+};
+
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 template <typename Element, typename Layout>

tools/library/src/manifest.cpp

@@ -1,5 +1,5 @@
 /***************************************************************************************************
- * Copyright (c) 2017-2019, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2017-2020, NVIDIA CORPORATION.  All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without modification, are permitted
  * provided that the following conditions are met:
@@ -37,11 +37,6 @@ namespace library {
 
 //////////////////////////////////////////////////////////////////////////////////////////////////////////
 
-// init and insert all cutlass op in manifest object (procedurally generated using generator.py)
-void initialize_all(Manifest &manifest);         
-
-/////////////////////////////////////////////////////////////////////////////////////////////////////////
-
 /// Top-level initialization
 Status Manifest::initialize() {
 
@@ -49,13 +44,8 @@ Status Manifest::initialize() {
     operations_.clear();
   }
 
-  switch(provider_) {
-    case Provider::kCUTLASS: 
-      initialize_all(*this); break;  
-
-    default: 
-      break;
-  }
+  // initialize procedurally generated cutlass op in manifest object
+  initialize_all(*this);
 
   return Status::kSuccess;
 }

tools/library/src/operation_table.cu

@@ -28,30 +28,7 @@
         instances may be queried.
 */
 
-#include <fstream>
-
-#include "cutlass/library/library.h"
 #include "cutlass/library/operation_table.h"
-#include "cutlass/library/util.h"
-
-/////////////////////////////////////////////////////////////////////////////////////////////////
-
-std::ostream & operator<<(std::ostream &out, cutlass::library::GemmFunctionalKey const &k) {
-
-  out << "{\n"
-    << "  element_compute: " << to_string(k.element_compute) << "\n"
-    << "   element_scalar: " << to_string(k.element_scalar) << "\n"
-    << "        element_A: " << to_string(k.element_A) << "\n"
-    << "         layout_A: " << to_string(k.layout_A) << "\n"
-    << "      transform_A: " << to_string(k.transform_A) << "\n"
-    << "        element_B: " << to_string(k.element_B) << "\n"
-    << "         layout_B: " << to_string(k.layout_B) << "\n"
-    << "      transform_B: " << to_string(k.transform_B) << "\n"
-    << "        element_C: " << to_string(k.element_C) << "\n"
-    << "}";
-
-  return out;
-}
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
@@ -67,85 +44,38 @@ void OperationTable::append(Manifest const &manifest) {
 
     OperationDescription const &desc = operation->description();
 
+    // insert all gemm operation into operation table
     if (desc.kind == OperationKind::kGemm) {
       GemmDescription const &gemm_desc = static_cast<GemmDescription const &>(desc);
     
-      if (gemm_desc.gemm_kind == GemmKind::kGemm) {
 
-        GemmFunctionalKey functional_key(
-          gemm_desc.tile_description.math_instruction.element_accumulator,
-          gemm_desc.element_epilogue,
-          gemm_desc.A.element,
-          gemm_desc.A.layout,
-          gemm_desc.transform_A,
-          gemm_desc.B.element,
-          gemm_desc.B.layout,
-          gemm_desc.transform_B,
-          gemm_desc.C.element
-        );
+      GemmFunctionalKey functional_key(
+        gemm_desc.provider,
+        gemm_desc.gemm_kind,
+        gemm_desc.tile_description.math_instruction.element_accumulator,
+        gemm_desc.element_epilogue,
+        gemm_desc.A.element,
+        gemm_desc.A.layout,
+        gemm_desc.transform_A,
+        gemm_desc.B.element,
+        gemm_desc.B.layout,
+        gemm_desc.transform_B,
+        gemm_desc.C.element
+      );
 
-        Operation const *op = operation.get();
+      Operation const *op = operation.get();
 
-        int cc = gemm_desc.tile_description.minimum_compute_capability;
+      int cc = gemm_desc.tile_description.minimum_compute_capability;
         
-        int alignment = std::max(std::max(
-          gemm_desc.A.alignment, gemm_desc.B.alignment), gemm_desc.C.alignment);
+      int alignment = std::max(std::max(
+        gemm_desc.A.alignment, gemm_desc.B.alignment), gemm_desc.C.alignment);
 
-        GemmPreferenceKey preference_key(cc, alignment);
+      GemmPreferenceKey preference_key(cc, alignment);
 
-        gemm_operations[functional_key][preference_key].push_back(op);
-      }
-      else if (gemm_desc.gemm_kind == GemmKind::kPlanarComplex) {
-
-        GemmFunctionalKey functional_key(
-          gemm_desc.tile_description.math_instruction.element_accumulator,
-          gemm_desc.element_epilogue,
-          gemm_desc.A.element,
-          gemm_desc.A.layout,
-          gemm_desc.transform_A,
-          gemm_desc.B.element,
-          gemm_desc.B.layout,
-          gemm_desc.transform_B,
-          gemm_desc.C.element
-        );
-
-        Operation const *op = operation.get();
-
-        int cc = gemm_desc.tile_description.minimum_compute_capability;
-        
-        int alignment = std::max(std::max(
-          gemm_desc.A.alignment, gemm_desc.B.alignment), gemm_desc.C.alignment);
-
-        GemmPreferenceKey preference_key(cc, alignment);
-
-        gemm_planar_complex_operations[functional_key][preference_key].push_back(op);
-      }
-      else if (gemm_desc.gemm_kind == GemmKind::kPlanarComplexArray) {
-
-        GemmFunctionalKey functional_key(
-          gemm_desc.tile_description.math_instruction.element_accumulator,
-          gemm_desc.element_epilogue,
-          gemm_desc.A.element,
-          gemm_desc.A.layout,
-          gemm_desc.transform_A,
-          gemm_desc.B.element,
-          gemm_desc.B.layout,
-          gemm_desc.transform_B,
-          gemm_desc.C.element
-        );
-
-        Operation const *op = operation.get();
-
-        int cc = gemm_desc.tile_description.minimum_compute_capability;
-        
-        int alignment = std::max(std::max(
-          gemm_desc.A.alignment, gemm_desc.B.alignment), gemm_desc.C.alignment);
-
-        GemmPreferenceKey preference_key(cc, alignment);
-
-        gemm_planar_complex_array_operations[functional_key][preference_key].push_back(op);
-      }
+      gemm_operations[functional_key][preference_key].push_back(op);
     }
+
+
   }
 
 }

tools/library/src/util.cu

@@ -45,6 +45,7 @@ static struct {
   Provider enumerant;
 }
 Provider_enumerants[] = {
+  {"none", "None", Provider::kNone},
   {"cutlass", "CUTLASS", Provider::kCUTLASS},
   {"host", "reference_host", Provider::kReferenceHost},
   {"device", "reference_device", Provider::kReferenceDevice},
@@ -83,6 +84,38 @@ Provider from_string<Provider>(std::string const &str) {
 }
 
 
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+static struct {
+  char const *text;
+  char const *pretty;
+  GemmKind enumerant;
+}
+GemmKind_enumerants[] = {
+  {"gemm", "<Gemm>", GemmKind::kGemm},
+  {"batched", "<Batched>", GemmKind::kBatched},
+  {"array", "<Array>", GemmKind::kArray},
+  {"universal", "<Universal>", GemmKind::kUniversal},
+  {"planar_complex", "<PlanarComplex>", GemmKind::kPlanarComplex},
+  {"planar_complex_array", "<PlanarComplexArray>", GemmKind::kPlanarComplexArray},
+};
+
+/// Converts a ConvKind enumerant to a string
+char const *to_string(GemmKind type, bool pretty) {
+
+  for (auto const & possible : GemmKind_enumerants) {
+    if (type == possible.enumerant) {
+      if (pretty) {
+        return possible.pretty;
+      }
+      else {
+        return possible.text;
+      }
+    }
+  }
+
+  return pretty ? "Invalid" : "invalid";
+}
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
@@ -92,6 +125,7 @@ static struct {
   OperationKind enumerant;
 }
 OperationKind_enumerants[] = {
+  {"eq_gemm", "EqGemm", OperationKind::kEqGemm}, 
   {"gemm", "Gemm", OperationKind::kGemm},               
 };
 
@@ -194,10 +228,14 @@ NumericTypeID_enumerants[] = {
   {"s32", "S32", NumericTypeID::kS32},
   {"s64", "S64", NumericTypeID::kS64},
   {"f16", "F16", NumericTypeID::kF16},
+  {"bf16", "BF16", NumericTypeID::kBF16},
   {"f32", "F32", NumericTypeID::kF32},
+  {"tf32", "TF32", NumericTypeID::kTF32},
   {"f64", "F64", NumericTypeID::kF64},
   {"cf16", "CF16", NumericTypeID::kCF16},
+  {"cbf16", "CBF16", NumericTypeID::kCBF16},
   {"cf32", "CF32", NumericTypeID::kCF32},
+  {"ctf32", "CTF32", NumericTypeID::kCTF32},
   {"cf64", "CF64", NumericTypeID::kCF64},
   {"cu4", "CU4", NumericTypeID::kCU4},
   {"cu8", "CU8", NumericTypeID::kCU8},
@@ -249,10 +287,14 @@ NumericTypeID from_string<NumericTypeID>(std::string const &str) {
 int sizeof_bits(NumericTypeID type) {
   switch (type) {
     case NumericTypeID::kF16: return 16;
+    case NumericTypeID::kBF16: return 16;
+    case NumericTypeID::kTF32: return 32;
     case NumericTypeID::kF32: return 32;
     case NumericTypeID::kF64: return 64;
     case NumericTypeID::kCF16: return 32;
+    case NumericTypeID::kCBF16: return 32;
     case NumericTypeID::kCF32: return 64;
+    case NumericTypeID::kCTF32: return 64;
     case NumericTypeID::kCF64: return 128;
     case NumericTypeID::kS4: return 4;
     case NumericTypeID::kS8: return 8;
@@ -276,6 +318,8 @@ bool is_complex_type(NumericTypeID type) {
     case NumericTypeID::kCF16: return true;
     case NumericTypeID::kCF32: return true;
     case NumericTypeID::kCF64: return true;
+    case NumericTypeID::kCBF16: return true;
+    case NumericTypeID::kCTF32: return true;
     default: break;
   }
   return false;
@@ -287,6 +331,8 @@ NumericTypeID get_real_type(NumericTypeID type) {
     case NumericTypeID::kCF16: return NumericTypeID::kF16;
     case NumericTypeID::kCF32: return NumericTypeID::kF32;
     case NumericTypeID::kCF64: return NumericTypeID::kF64;
+    case NumericTypeID::kCBF16: return NumericTypeID::kBF16;
+    case NumericTypeID::kCTF32: return NumericTypeID::kTF32;
     default: break;
   }
   return type;
@@ -314,6 +360,8 @@ bool is_integer_type(NumericTypeID type) {
 bool is_signed_type(NumericTypeID type) {
   switch (type) {
     case NumericTypeID::kF16: return true;
+    case NumericTypeID::kBF16: return true;
+    case NumericTypeID::kTF32: return true;
     case NumericTypeID::kF32: return true;
     case NumericTypeID::kF64: return true;
     case NumericTypeID::kS4: return true;
@@ -340,9 +388,13 @@ bool is_unsigned_integer(NumericTypeID type) {
 bool is_float_type(NumericTypeID type) {
   switch (type) {
   case NumericTypeID::kF16: return true;
+  case NumericTypeID::kBF16: return true;
+  case NumericTypeID::kTF32: return true;
   case NumericTypeID::kF32: return true;
   case NumericTypeID::kF64: return true;
   case NumericTypeID::kCF16: return true;
+  case NumericTypeID::kCBF16: return true;
+  case NumericTypeID::kCTF32: return true;
   case NumericTypeID::kCF32: return true;
   case NumericTypeID::kCF64: return true;
   default: break;
@@ -431,7 +483,7 @@ OpcodeClassID_enumerants[] = {
   {"simt", "<simt>", OpcodeClassID::kSimt},
   {"tensorop", "<tensorop>", OpcodeClassID::kTensorOp},
   {"wmmatensorop", "<wmmatensorop>", OpcodeClassID::kWmmaTensorOp},
-  {"wmma", "<wmma>", OpcodeClassID::kWmmaTensorOp}
+  {"wmma", "<wmma>", OpcodeClassID::kWmmaTensorOp},
 };
 
 /// Converts a OpcodeClassID enumerant to a string
@@ -509,6 +561,47 @@ ComplexTransform from_string<ComplexTransform>(std::string const &str) {
 }
 
 
+static struct {
+  char const *text;
+  char const *pretty;
+  SplitKMode enumerant;
+}
+SplitKMode_enumerants[] = {
+  {"serial", "<serial>", SplitKMode::kSerial},
+  {"parallel", "<parallel>", SplitKMode::kParallel},
+};
+
+/// Converts a SplitKMode enumerant to a string
+char const *to_string(SplitKMode type, bool pretty) {
+
+  for (auto const & possible : SplitKMode_enumerants) {
+    if (type == possible.enumerant) {
+      if (pretty) {
+        return possible.pretty;
+      }
+      else {
+        return possible.text;
+      }
+    }
+  }
+
+  return pretty ? "Invalid" : "invalid";
+}
+
+/// Converts a SplitKMode enumerant from a string
+template <>
+SplitKMode from_string<SplitKMode>(std::string const &str) {
+
+  for (auto const & possible : SplitKMode_enumerants) {
+    if ((str.compare(possible.text) == 0) ||
+        (str.compare(possible.pretty) == 0)) {
+      return possible.enumerant;
+    }
+  }
+
+  return SplitKMode::kInvalid;
+}
+
 /////////////////////////////////////////////////////////////////////////////////////////////////
 /// Lexical cast a string to a byte array. Returns true if cast is successful or false if invalid.
 bool lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type, std::string const &str) {
@@ -570,6 +663,20 @@ bool lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type, std::string c
     *reinterpret_cast<half_t *>(bytes.data()) = static_cast<half_t>(tmp);
   }
     break;
+  case NumericTypeID::kBF16:
+  {
+    float tmp;
+    ss >> tmp;
+    *reinterpret_cast<bfloat16_t *>(bytes.data()) = static_cast<bfloat16_t>(tmp);
+  }
+    break;
+  case NumericTypeID::kTF32:
+  {
+    float tmp;
+    ss >> tmp;
+    *reinterpret_cast<tfloat32_t *>(bytes.data()) = static_cast<tfloat32_t>(tmp);
+  }
+    break;
   case NumericTypeID::kF32:
   {
     ss >> *reinterpret_cast<float *>(bytes.data());
@@ -589,11 +696,29 @@ bool lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type, std::string c
     x->imag() = static_cast<half_t>(std::imag(tmp));
   }
     break;
+  case NumericTypeID::kCBF16:
+  {
+    std::complex<float> tmp;
+    ss >> tmp;
+    cutlass::complex<cutlass::bfloat16_t> *x = reinterpret_cast<cutlass::complex<bfloat16_t> *>(bytes.data());
+    x->real() = static_cast<bfloat16_t>(std::real(tmp));
+    x->imag() = static_cast<bfloat16_t>(std::imag(tmp));
+  }
+    break;
   case NumericTypeID::kCF32:
   {
     ss >> *reinterpret_cast<std::complex<float>*>(bytes.data());
   }
     break;
+  case NumericTypeID::kCTF32:
+  {
+    std::complex<float> tmp;
+    ss >> tmp;
+    cutlass::complex<cutlass::tfloat32_t> *x = reinterpret_cast<cutlass::complex<tfloat32_t> *>(bytes.data());
+    x->real() = static_cast<tfloat32_t>(std::real(tmp));
+    x->imag() = static_cast<tfloat32_t>(std::imag(tmp));
+  }
+    break;
   case NumericTypeID::kCF64:
   {
     ss >> *reinterpret_cast<std::complex<double>*>(bytes.data());
@@ -674,6 +799,18 @@ std::string lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type) {
     ss << tmp;
   }
     break;
+  case NumericTypeID::kBF16:
+  {
+    float tmp = *reinterpret_cast<bfloat16_t *>(bytes.data());;
+    ss << tmp;
+  }
+    break;
+  case NumericTypeID::kTF32:
+  {
+    float tmp = *reinterpret_cast<tfloat32_t *>(bytes.data());;
+    ss << tmp;
+  }
+    break;
   case NumericTypeID::kF32:
   {
     ss << *reinterpret_cast<float *>(bytes.data());
@@ -696,6 +833,18 @@ std::string lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type) {
     }
   }
     break;
+  case NumericTypeID::kCBF16:
+  {
+    cutlass::complex<bfloat16_t> const *x = 
+      reinterpret_cast<cutlass::complex<bfloat16_t> const *>(bytes.data());
+
+    ss << float(x->real());
+
+    if (x->imag() != cutlass::bfloat16_t()) {
+      ss << "+i" << float(x->imag());
+    }
+  }
+    break;
   case NumericTypeID::kCF32:
   {
     cutlass::complex<float> const * x = reinterpret_cast<cutlass::complex<float> const *>(bytes.data());
@@ -707,6 +856,17 @@ std::string lexical_cast(std::vector<uint8_t> &bytes, NumericTypeID type) {
     }
   }
     break;
+  case NumericTypeID::kCTF32:
+  {
+    cutlass::complex<tfloat32_t> const * x = reinterpret_cast<cutlass::complex<tfloat32_t> const *>(bytes.data());
+
+    ss << float(x->real());
+
+    if (x->imag() != tfloat32_t()) {
+      ss << "+i" << float(x->imag());
+    }
+  }
+    break;
   case NumericTypeID::kCF64:
   {
     cutlass::complex<double> const * x = reinterpret_cast<cutlass::complex<double> const *>(bytes.data());
@@ -780,6 +940,16 @@ bool cast_from_int64(std::vector<uint8_t> &bytes, NumericTypeID type, int64_t sr
     *reinterpret_cast<half_t *>(bytes.data()) = static_cast<half_t>(float(src));
   }
     break;
+  case NumericTypeID::kBF16:
+  {
+    *reinterpret_cast<bfloat16_t *>(bytes.data()) = static_cast<bfloat16_t>(float(src));
+  }
+    break;
+  case NumericTypeID::kTF32:
+  {
+    *reinterpret_cast<tfloat32_t *>(bytes.data()) = static_cast<tfloat32_t>(float(src));
+  }
+    break;
   case NumericTypeID::kF32:
   {
     *reinterpret_cast<float *>(bytes.data()) = static_cast<float>(src);
@@ -870,6 +1040,16 @@ bool cast_from_uint64(std::vector<uint8_t> &bytes, NumericTypeID type, uint64_t
     *reinterpret_cast<half_t *>(bytes.data()) = static_cast<half_t>(float(src));
   }
     break;
+  case NumericTypeID::kBF16:
+  {
+    *reinterpret_cast<bfloat16_t *>(bytes.data()) = static_cast<bfloat16_t>(float(src));
+  }
+    break;
+  case NumericTypeID::kTF32:
+  {
+    *reinterpret_cast<tfloat32_t *>(bytes.data()) = static_cast<tfloat32_t>(float(src));
+  }
+    break;
   case NumericTypeID::kF32:
   {
     *reinterpret_cast<float *>(bytes.data()) = static_cast<float>(src);
@@ -961,6 +1141,16 @@ bool cast_from_double(std::vector<uint8_t> &bytes, NumericTypeID type, double sr
     *reinterpret_cast<half_t *>(bytes.data()) = static_cast<half_t>(float(src));
   }
     break;
+  case NumericTypeID::kBF16:
+  {
+    *reinterpret_cast<bfloat16_t *>(bytes.data()) = static_cast<bfloat16_t>(float(src));
+  }
+    break;
+  case NumericTypeID::kTF32:
+  {
+    *reinterpret_cast<tfloat32_t *>(bytes.data()) = static_cast<tfloat32_t>(float(src));
+  }
+    break;
   case NumericTypeID::kF32:
   {
     *reinterpret_cast<float *>(bytes.data()) = static_cast<float>(src);
@@ -978,11 +1168,23 @@ bool cast_from_double(std::vector<uint8_t> &bytes, NumericTypeID type, double sr
     x->imag() = static_cast<half_t>(float(0));
   }
     break;
+  case NumericTypeID::kCBF16:
+  {
+    cutlass::complex<cutlass::bfloat16_t> *x = reinterpret_cast<cutlass::complex<bfloat16_t> *>(bytes.data());
+    x->real() = static_cast<bfloat16_t>(bfloat16_t(src));
+    x->imag() = static_cast<bfloat16_t>(bfloat16_t(0));
+  }
+    break;
   case NumericTypeID::kCF32:
   {
     *reinterpret_cast<std::complex<float>*>(bytes.data()) = std::complex<float>(float(src), float(0));
   }
     break;
+  case NumericTypeID::kCTF32:
+  {
+    *reinterpret_cast<std::complex<tfloat32_t>*>(bytes.data()) = std::complex<tfloat32_t>(tfloat32_t(src), tfloat32_t(0));
+  }
+    break;
   case NumericTypeID::kCF64:
   {
     *reinterpret_cast<std::complex<double>*>(bytes.data()) = std::complex<double>(src, double(0));