CUTLASS 3.6.0 (#1850)

* v3.6

* update changelog

* update readme

* fix typo

* fixing typos

* hopper gemm with weight prefetch

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>
Co-authored-by: Haicheng Wu <haichengw@nvidia.com>

tools/profiler/include/cutlass/profiler/cublas_helpers.h

@@ -36,6 +36,7 @@
 
 #if CUTLASS_ENABLE_CUBLAS
 #include <cublas_v2.h>
+#include <cublasLt.h>
 
 #include "cutlass/cutlass.h"
 #include "cutlass/library/library.h"
@@ -90,25 +91,48 @@ Status cublas_satisfies(library::SymmDescription const &desc);
 /// Additionally, it provides implicit cast from CublasCreate's object to cublasHandle_t's object
 class CublasCreate {
 private:
-	cublasHandle_t handle;
-	cublasStatus_t status;
+  cublasHandle_t handle;
+  cublasStatus_t status;
 
 public:
-	CublasCreate() {
-		status = cublasCreate(&handle);
-	}
+  CublasCreate() {
+    status = cublasCreate(&handle);
+  }
 
-	~CublasCreate() {
-		cublasDestroy(handle);
-	}
+  ~CublasCreate() {
+    cublasDestroy(handle);
+  }
 
-    /// Implicit cast CublasCreate object to cublasHandle_t
-    operator cublasHandle_t() const { return handle; }
+  /// Implicit cast CublasCreate object to cublasHandle_t
+  operator cublasHandle_t() const { return handle; }
 
-    /// returns cublasStatus_t for handle creation
-    cublasStatus_t get_cublas_create_status() { return status; }
+  /// returns cublasStatus_t for handle creation
+  cublasStatus_t get_cublas_create_status() { return status; }
 };
 
+/// This is a helper class to create cublasLtHandle_t automatically on CublasLtCreate object creation and 
+/// to destroy cublasLtHandle_t on CublasLtCreate object destruction. 
+/// Additionally, it provides implicit cast from CublasLtCreate's object to cublasLtHandle_t's object
+class CublasLtCreate {
+private:
+  cublasLtHandle_t handle;
+  cublasStatus_t status;
+
+public:
+  CublasLtCreate() {
+    status = cublasLtCreate(&handle);
+  }
+
+  ~CublasLtCreate() {
+    cublasLtDestroy(handle);
+  }
+
+  /// Implicit cast CublasLtCreate object to cublasLtHandle_t
+  operator cublasLtHandle_t() const { return handle; }
+
+  /// returns cublasLtStatus_t for handle creation
+  cublasStatus_t get_cublaslt_create_status() { return status; }
+};
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 namespace detail {
@@ -226,6 +250,80 @@ struct cublasGemmExDispatcher {
   cublasStatus_t operator()(cublasHandle_t handle);
 };
 
+/// Dispatcher to cublaslt kernels 
+//
+struct cublasLtGemmExDispatcher {
+
+  //
+  // Data members
+  //
+  library::GemmDescription const &op_desc;
+  library::GemmUniversalConfiguration configuration;
+  library::GemmUniversalArguments arguments;
+
+  // cublas-specific data structures to fill cublas API call arguments
+  cublasOperation_t trans_A;
+  cublasOperation_t trans_B;
+  cudaDataType_t data_type_A;
+  cudaDataType_t data_type_B;
+  cudaDataType_t data_type_C;
+  cudaDataType_t compute_data_type = CUDA_R_32F;
+
+  //cublasLt-specific data structures
+  cublasLtMatmulDesc_t operationDesc = NULL;
+  cublasLtMatrixLayout_t Adesc = NULL, Bdesc = NULL, Cdesc = NULL, Ddesc = NULL;
+  cublasLtMatmulPreference_t preference = NULL;
+  
+  //is set by call to get_cublaslt_algo()
+  cublasLtMatmulHeuristicResult_t heuristicResult_;
+  void *workspace = nullptr;
+
+  Status status;
+
+#if (__CUDACC_VER_MAJOR__ >= 11)
+  cublasComputeType_t compute_type;
+#endif
+
+  //
+  // Methods
+  //
+
+  cublasLtGemmExDispatcher( 
+    library::GemmDescription const &op_desc,
+    library::GemmUniversalConfiguration configuration_,
+    library::GemmUniversalArguments arguments_
+  );
+
+  /// Initialize the cublasLt variables
+  void initialize_cublaslt();
+  
+
+  /// Runs auto-tuning for the cublas heuristics
+  bool get_cublaslt_algo(cublasLtHandle_t handle,
+    AlgorithmMode algorithm_mode 
+    ); 
+
+  /// Executes GEMM using these arguments
+  cublasStatus_t operator()(cublasLtHandle_t handle);
+
+  ~cublasLtGemmExDispatcher(){
+
+    // descriptors are no longer needed as all GPU work was already enqueued
+    if (preference) cublasLtMatmulPreferenceDestroy(preference);
+    if (Ddesc) cublasLtMatrixLayoutDestroy(Ddesc);
+    if (Cdesc) cublasLtMatrixLayoutDestroy(Cdesc);
+    if (Bdesc) cublasLtMatrixLayoutDestroy(Bdesc);
+    if (Adesc) cublasLtMatrixLayoutDestroy(Adesc);
+    if (operationDesc) cublasLtMatmulDescDestroy(operationDesc);
+
+    if (workspace) {
+      cudaFree(workspace);
+    }
+
+  } 
+
+};
+
 ///////////////////////////////////////////////////////////////////////////////////////////////////
 
 /// Dispatcher to cublas rank k update kernels 

tools/profiler/include/cutlass/profiler/cutlass_profiler.h

@@ -48,7 +48,7 @@ namespace profiler {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-/// CUTLASS Profiler application 
+/// CUTLASS Profiler application
 class CutlassProfiler {
 private:
 
@@ -66,13 +66,10 @@ private:
 
   /// Prints usage
   void print_usage_(std::ostream &);
-  
+
   /// Prints usage
   void print_options_(std::ostream &);
 
-  /// Initializes the device
-  void initialize_device_();
-
   /// Enumerates all operations
   void enumerate_();
 

tools/profiler/include/cutlass/profiler/device_allocation.h

@@ -81,6 +81,9 @@ private:
   /// Buffer holding TensorRef instance to recently allocated memory
   std::vector<uint8_t> tensor_ref_buffer_;
 
+  /// The device ID where the allocation is made
+  int device_;
+
 public:
   //
   // Static member functions
@@ -91,7 +94,7 @@ public:
 
   /// Returns the stride of a packed layout
   static std::vector<int64_t> get_packed_layout(
-    library::LayoutTypeID layout_id, 
+    library::LayoutTypeID layout_id,
     std::vector<int> const &extent);
 
   /// returns the capacity needed
@@ -103,16 +106,16 @@ public:
 
   /// Returns true if two blocks have exactly the same value
   static bool block_compare_equal(
-    library::NumericTypeID numeric_type, 
-    void const *ptr_A, 
-    void const *ptr_B, 
+    library::NumericTypeID numeric_type,
+    void const *ptr_A,
+    void const *ptr_B,
     size_t capacity);
 
   /// Returns true if two blocks have approximately the same value
   static bool block_compare_relatively_equal(
-    library::NumericTypeID numeric_type, 
-    void const *ptr_A, 
-    void const *ptr_B, 
+    library::NumericTypeID numeric_type,
+    void const *ptr_A,
+    void const *ptr_B,
     size_t capacity,
     double epsilon,
     double nonzero_floor);
@@ -123,15 +126,19 @@ public:
   //
 
   DeviceAllocation();
-  
-  DeviceAllocation(library::NumericTypeID type, size_t capacity);
-  
+
   DeviceAllocation(
-    library::NumericTypeID type, 
-    library::LayoutTypeID layout_id, 
-    std::vector<int> const &extent, 
+    library::NumericTypeID type,
+    size_t capacity,
+    int device = -1);
+
+  DeviceAllocation(
+    library::NumericTypeID type,
+    library::LayoutTypeID layout_id,
+    std::vector<int> const &extent,
     std::vector<int64_t> const &stride = std::vector<int64_t>(),
-    int batch_count = 1);
+    int batch_count = 1,
+    int device = -1);
 
   ~DeviceAllocation();
 
@@ -142,9 +149,9 @@ public:
 
   /// Allocates memory for a given layout and tensor
   DeviceAllocation &reset(
-    library::NumericTypeID type, 
-    library::LayoutTypeID layout_id, 
-    std::vector<int> const &extent, 
+    library::NumericTypeID type,
+    library::LayoutTypeID layout_id,
+    std::vector<int> const &extent,
     std::vector<int64_t> const &stride = std::vector<int64_t>(),
     int batch_count = 1);
 
@@ -157,7 +164,7 @@ public:
 
   /// Data type of contained elements
   library::NumericTypeID type() const;
-  
+
   /// Pointer to start of device memory allocation
   void *data() const;
 
@@ -184,7 +191,7 @@ public:
 
   /// Capacity of allocation in number of elements
   size_t capacity() const;
-  
+
   /// Capacity of allocation in bytes
   size_t bytes() const;
 
@@ -205,7 +212,7 @@ public:
 
   /// Initializes a host allocation to a random distribution using std::cout
   void initialize_random_sparsemeta_host(int seed, int MetaSizeInBits);
-  
+
   /// Uniformly fills a tensor with a value when provided o.w. zero
   void fill_device(double value);
 
@@ -221,8 +228,12 @@ public:
   /// Copies from an equivalent-sized tensor in device memory
   void copy_to_host(void *ptr);
 
-  /// Writes a tensor to csv 
+  /// Writes a tensor to csv
   void write_tensor_csv(std::ostream &out);
+
+private:
+  /// A wrapper that sets the device, performs malloc, and sets back
+  cudaError_t malloc(void** ptr, size_t size);
 };
 
 using DeviceAllocationList = std::list<DeviceAllocation>;

tools/profiler/include/cutlass/profiler/device_context.h

@@ -29,7 +29,7 @@
  *
  **************************************************************************************************/
 /* \file
-   \brief 
+   \brief
 */
 
 #pragma once
@@ -68,46 +68,52 @@ private:
 
   /// Non-owning set of named allocations
   AllocationMap allocations_;
-  
+
 public:
 
   /// Allocates memory of a given type, capacity (elements), and name
   DeviceAllocation *allocate_block(
+    Options const &options,
     std::string const &name,
-    library::NumericTypeID type, 
-    size_t capacity);
-
-  /// Allocates memory of a given type, capacity (elements), and name
-  DeviceAllocation *allocate_tensor(
-    std::string const &name,
-    library::NumericTypeID type, 
-    library::LayoutTypeID layout_id, 
-    std::vector<int> const &extent, 
-    std::vector<int64_t> const &stride = std::vector<int64_t>(),
-    int batch_count = 1);
+    library::NumericTypeID type,
+    size_t capacity,
+    size_t device_index);
 
   /// Allocates memory of a given type, capacity (elements), and name
   DeviceAllocation *allocate_tensor(
     Options const &options,
     std::string const &name,
-    library::NumericTypeID type, 
-    library::LayoutTypeID layout_id, 
-    std::vector<int> const &extent, 
+    library::NumericTypeID type,
+    library::LayoutTypeID layout_id,
+    std::vector<int> const &extent,
     std::vector<int64_t> const &stride,
     int batch_count,
-    int seed_shift = 0);
+    size_t device_index);
 
-  /// Allocates memory for sparse meta data 
-  DeviceAllocation *allocate_sparsemeta_tensor(
+  /// Allocates memory of a given type, capacity (elements), and name
+  DeviceAllocation *allocate_and_initialize_tensor(
     Options const &options,
     std::string const &name,
-    library::NumericTypeID type, 
-    library::LayoutTypeID layout_id, 
+    library::NumericTypeID type,
+    library::LayoutTypeID layout_id,
+    std::vector<int> const &extent,
+    std::vector<int64_t> const &stride,
+    int batch_count,
+    int seed_shift,
+    size_t device_index);
+
+  /// Allocates memory for sparse meta data
+  DeviceAllocation *allocate_and_initialize_sparsemeta_tensor(
+    Options const &options,
+    std::string const &name,
+    library::NumericTypeID type,
+    library::LayoutTypeID layout_id,
     library::NumericTypeID type_a,
-    std::vector<int> const &extent, 
+    std::vector<int> const &extent,
     std::vector<int64_t> const &stride,
     int batch_count,
-    int seed_shift = 0);
+    int seed_shift,
+    size_t device_index);
 
   /// Clears named allocations (but does not necessarily free memory)
   void clear();

tools/profiler/include/cutlass/profiler/options.h

@@ -82,12 +82,16 @@ public:
   struct Device {
 
     /// Device ID
-    int device;
+    std::vector<int> devices;
+
+    /// Number of total devices
+    /// This is not set by the user, it is set by automatically
+    int num_devices;
 
     /// CUDA Device properties
-    cudaDeviceProp properties;
+    std::vector<cudaDeviceProp> properties;
 
-    /// Total memory allocation on device
+    /// Total memory allocation on each device
     size_t maximum_capacity;
 
     //
@@ -100,8 +104,11 @@ public:
     void print_options(std::ostream &out, int indent = 0) const;
     void print_device_info(std::ostream &out) const;
 
-    /// Returns the compute capability of the listed device (e.g. 61, 60, 70, 75)
-    int compute_capability() const;
+    /// Returns the device ID from a device index
+    int device_id(size_t device_index) const;
+
+    /// Returns the compute capability of the listed devices (e.g. 61, 60, 70, 75)
+    int compute_capability(int device_index) const;
   };
 
   /// Options related to initializing input tensors
@@ -129,7 +136,7 @@ public:
     //
 
     explicit Initialization(CommandLine const &cmdline);
-    
+
     void print_usage(std::ostream &out) const;
     void print_options(std::ostream &out, int indent = 0) const;
 
@@ -171,13 +178,13 @@ public:
     //
 
     explicit Verification(CommandLine const &cmdline);
-  
+
     void print_usage(std::ostream &out) const;
     void print_options(std::ostream &out, int indent = 0) const;
 
     /// Returns true if a provider is enabled
     bool provider_enabled(library::Provider provider) const;
-    
+
     /// Returns the index of a provider if its enabled
     size_t index(library::Provider provider) const;
   };
@@ -225,7 +232,7 @@ public:
     /// Returns the index of a provider if its enabled
     size_t index(library::Provider provider) const;
   };
-  
+
   /// Options related to reporting
   struct Report {
 
@@ -260,7 +267,7 @@ public:
     //
 
     explicit Report(CommandLine const &cmdline);
-    
+
     void print_usage(std::ostream &out) const;
     void print_options(std::ostream &out, int indent = 0) const;
   };
@@ -282,7 +289,7 @@ public:
     //
 
     explicit About(CommandLine const &cmdline);
-    
+
     void print_usage(std::ostream &out) const;
     void print_options(std::ostream &out, int indent = 0) const;
 
@@ -303,7 +310,7 @@ public:
 
   /// Vector of operation name substrings
   std::vector<std::string> operation_names;
-  
+
   /// Vector of operation name substrings
   std::vector<std::string> excluded_operation_names;