3.6.0 update (#2005)

* 3.6.0 update

* doc and swap stuff

---------

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

tools/profiler/src/conv2d_operation_profiler.cu

@@ -396,6 +396,29 @@ Status Conv2dOperationProfiler::initialize_configuration(
       problem_, operation_desc.conv_kind, operation_desc.A.layout,
       operation_desc.B.layout, operation_desc.C.layout);
 
+#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
+  {
+    auto print_vector = [] (const auto& vec) {
+      printf("[");
+      for (size_t k = 0; k < vec.size(); ++k) {
+        cute::print(vec[k]);
+        if (k + 1 < vec.size()) {
+          printf(",");
+        }
+      }
+      printf("]");
+    };
+  
+    printf("\n    conv_workspace_.configuration.stride_a: ");
+    print_vector(conv_workspace_.configuration.stride_a);
+    printf("\n    conv_workspace_.configuration.stride_b: ");
+    print_vector(conv_workspace_.configuration.stride_b);
+    printf("\n    conv_workspace_.configuration.stride_c: ");
+    print_vector(conv_workspace_.configuration.stride_c);
+    printf("\n");
+  }
+#endif
+
   // initialize library::ConvArguments
   conv_workspace_.arguments.A            = nullptr;
   conv_workspace_.arguments.B            = nullptr;
@@ -1237,7 +1260,7 @@ bool Conv2dOperationProfiler::profile(
     }
 
     results_.back().status = profile_cutlass_(
-      results_.back().runtime,
+      results_.back(),
       options,
       operation,
       &conv_workspace_.arguments,
@@ -1251,7 +1274,7 @@ bool Conv2dOperationProfiler::profile(
 
 /// Method to profile a CUTLASS Operation
 Status Conv2dOperationProfiler::profile_cutlass_(
-  double &runtime,
+  PerformanceResult &result,
   Options const &options,
   library::Operation const *operation,
   void *arguments,
@@ -1387,7 +1410,7 @@ Status Conv2dOperationProfiler::profile_cutlass_(
   // Update performance result
   //
 
-  runtime = timer.duration(iteration);
+  result.runtime = timer.duration(iteration);
 
   return status;
 }

tools/profiler/src/conv3d_operation_profiler.cu

@@ -1099,7 +1099,7 @@ bool Conv3dOperationProfiler::profile(
     set_cutlass_operator_arguments_();
 
     results_.back().status = profile_cutlass_(
-      results_.back().runtime,
+      results_.back(),
       options,
       operation,
       &conv_workspace_.arguments,
@@ -1141,7 +1141,7 @@ void Conv3dOperationProfiler::set_cutlass_operator_arguments_(int problem_idx) {
 
 /// Method to profile a CUTLASS Operation
 Status Conv3dOperationProfiler::profile_cutlass_(
-  double &runtime,
+  PerformanceResult &result,
   Options const &options,
   library::Operation const *operation,
   void *arguments,
@@ -1248,7 +1248,7 @@ Status Conv3dOperationProfiler::profile_cutlass_(
   // Update performance result
   //
 
-  runtime = timer.duration(iteration);
+  result.runtime = timer.duration(iteration);
 
   return status;
 }

tools/profiler/src/cublas_helpers.cu

@@ -656,7 +656,7 @@ bool cublasLtGemmExDispatcher::get_cublaslt_algo(cublasLtHandle_t handle,
   return true;
 }
 
-cublasStatus_t cublasLtGemmExDispatcher::operator()(cublasLtHandle_t handle) 
+cublasStatus_t cublasLtGemmExDispatcher::operator()(cublasLtHandle_t handle, cudaStream_t stream)
 {
   return cublasLtMatmul(handle,
     operationDesc,
@@ -673,7 +673,7 @@ cublasStatus_t cublasLtGemmExDispatcher::operator()(cublasLtHandle_t handle)
     &heuristicResult_.algo,
     workspace,
     heuristicResult_.workspaceSize,
-    0); //number of streams is set to 0
+    stream); //number of streams is set to 0
   
 }
 

tools/profiler/src/device_allocation.cu

@@ -290,9 +290,8 @@ DeviceAllocation::DeviceAllocation():
   capacity_(0),
   pointer_(nullptr),
   layout_(library::LayoutTypeID::kUnknown),
-  batch_count_(1),
-  device_(-1) {
-
+  batch_count_(1) {
+  cudaGetDevice(&device_);
 }
 
 DeviceAllocation::DeviceAllocation(
@@ -329,13 +328,33 @@ DeviceAllocation::DeviceAllocation(
 
 DeviceAllocation::~DeviceAllocation() {
   if (pointer_) {
+    int current_device;
+    cudaGetDevice(&current_device);
+
+    if (current_device != device_) {
+      cudaSetDevice(device_);
+    }
     cudaFree(pointer_);
+
+    if (current_device != device_) {
+      cudaSetDevice(current_device);
+    }
   }
 }
 
 DeviceAllocation &DeviceAllocation::reset() {
   if (pointer_) {
+    int current_device;
+    cudaGetDevice(&current_device);
+
+    if (current_device != device_) {
+      cudaSetDevice(device_);
+    }
     cudaFree(pointer_);
+
+    if (current_device != device_) {
+      cudaSetDevice(current_device);
+    }
   }
 
   type_ = library::NumericTypeID::kInvalid;
@@ -2438,25 +2457,11 @@ void DeviceAllocation::fill_host(double val = 0.0) {
 
 cudaError_t DeviceAllocation::malloc(void** ptr, size_t size) {
   cudaError_t result;
-  int set_device_back_to = -1;
+  int current_device;
+  cudaGetDevice(&current_device);
 
-  /// When needed this sets the device to the allocation's device remembering
-  /// the current device so that it can be set back after the cudaMalloc is
-  /// performed.
-  if (device_ >= 0) {
-    int current_device;
-    result = cudaGetDevice(&current_device);
-    if (result != cudaSuccess) {
-      return result;
-    }
-
-    if (current_device != device_) {
-      set_device_back_to = current_device;
-      result = cudaSetDevice(device_);
-      if (result != cudaSuccess) {
-        return result;
-      }
-    }
+  if (current_device != device_) {
+    cudaSetDevice(device_);
   }
 
   // This performs the cudaMalloc
@@ -2465,13 +2470,8 @@ cudaError_t DeviceAllocation::malloc(void** ptr, size_t size) {
     return result;
   }
 
-  /// When needed this sets the device back to what it was when the function was
-  /// called.
-  if (set_device_back_to != -1) {
-    result = cudaSetDevice(set_device_back_to);
-    if (result != cudaSuccess) {
-      return result;
-    }
+  if (current_device != device_) {
+    cudaSetDevice(current_device);
   }
 
   return cudaSuccess;

tools/profiler/src/gpu_timer.cpp

@@ -33,9 +33,11 @@
 */
 
 #include <stdexcept>
+#include <cstring>
 
 #include "cutlass/profiler/gpu_timer.h"
 
+
 namespace cutlass {
 namespace profiler {
 
@@ -52,32 +54,39 @@ GpuTimer::GpuTimer() {
   }
 }
 
+GpuTimer::GpuTimer(GpuTimer&& gpu_timer) noexcept {
+  memcpy(events, gpu_timer.events, sizeof(events));
+  memset(gpu_timer.events, 0, sizeof(gpu_timer.events));
+}
+
 GpuTimer::~GpuTimer() {
-  for (auto & event : events) {
-    cudaEventDestroy(event);
+  for (const auto & event : events) {
+    if (event != nullptr) {
+      cudaEventDestroy(event);
+    }
   }
 }
 
-/// Records a start event in the stream
-void GpuTimer::start(cudaStream_t stream) {
-  cudaError_t result = cudaEventRecord(events[0], stream);
+/// Records a start event in the stream, the flag is for cudaEventRecordWithFlags
+void GpuTimer::start(cudaStream_t stream, const unsigned int flag) {
+  cudaError_t result = cudaEventRecordWithFlags(events[0], stream, flag);
   if (result != cudaSuccess) {
     throw std::runtime_error("Failed to record start event.");
   }
 }
 
-/// Records a stop event in the stream
-void GpuTimer::stop(cudaStream_t stream) {
-cudaError_t result = cudaEventRecord(events[1], stream);
+/// Records a stop event in the stream, the flag is for cudaEventRecordWithFlags
+void GpuTimer::stop(cudaStream_t stream, const unsigned int flag) {
+cudaError_t result = cudaEventRecordWithFlags(events[1], stream, flag);
   if (result != cudaSuccess) {
     throw std::runtime_error("Failed to record stop event.");
   }
 }
 
-/// Records a stop event in the stream and synchronizes on the stream
-void GpuTimer::stop_and_wait(cudaStream_t stream) {
+/// Records a stop event in the stream and synchronizes on the stream, the flag is for cudaEventRecordWithFlags
+void GpuTimer::stop_and_wait(cudaStream_t stream, const unsigned int flag) {
 
-  stop(stream);
+  stop(stream, flag);
 
   cudaError_t result;
   if (stream) {

tools/profiler/src/operation_profiler.cu

@@ -658,7 +658,7 @@ void OperationProfiler::save_workspace(
 
 /// Method to profile a CUTLASS Operation
 Status OperationProfiler::profile_cutlass_(
-  double &runtime,
+  PerformanceResult &result,
   Options const &options,
   library::Operation const *operation,
   void *arguments,
@@ -726,7 +726,7 @@ Status OperationProfiler::profile_cutlass_(
   // Update performance result
   //
 
-  runtime = timer.duration(iteration);
+  result.runtime = timer.duration(iteration);
 
   return status;
 }

tools/profiler/src/options.cu

@@ -307,12 +307,6 @@ void Options::Initialization::get_distribution(
     {0, 0}
   };
 
-  // Initalize pnz values to a default value of 100%
-  dist.gaussian.pnz = 1.0;
-  dist.gaussian.pnzA = 1.0;
-  dist.gaussian.pnzB = 1.0;
-  dist.gaussian.pnzC = 1.0;
-
   using KeyValueVector = std::vector<std::pair<std::string, std::string> >;
 
   KeyValueVector values;
@@ -330,6 +324,25 @@ void Options::Initialization::get_distribution(
     ++it;
   }
 
+  // Default initialization
+  switch (dist.kind) {
+    case cutlass::Distribution::Uniform:
+      dist.set_uniform(-4/*min*/, 4/*max*/);
+      break;
+    case cutlass::Distribution::Gaussian:
+      dist.set_gaussian(0/*mean*/, 4/*stddev*/);
+      break;
+    case cutlass::Distribution::Identity:
+      dist.set_identity();
+      break;
+    case cutlass::Distribution::Sequential:
+      dist.set_sequential(0/*start*/, 4/*delta*/);
+      break;
+    default:
+      dist.set_uniform(-4/*min*/, 4/*max*/);
+      return;
+  }
+
   // Subsequent key-value pairs update the named field of the distribution struct.
   for (; it != values.end(); ++it) {
     // Integer scaling factor - if < 0, no integer rounding is performed.

tools/profiler/src/performance_report.cpp

@@ -337,7 +337,15 @@ std::ostream & PerformanceReport::print_csv_header_(
     << ",Bytes"
     << ",Flops"
     << ",Flops/Byte"
-    << ",Runtime"
+    << ",Runtime";
+
+  if (options_.device.devices.size() > 1) {
+    for (size_t i = 0; i < options_.device.devices.size(); i++) {
+      out << ",Runtime_" << i;
+    }
+  }
+
+  out
     << ",GB/s"
     << ",GFLOPs"
     ;
@@ -376,6 +384,16 @@ std::ostream & PerformanceReport::print_result_csv_(
     << "," << result.flops / result.bytes
     << "," << result.runtime;
 
+  if (options_.device.devices.size() > 1) {
+    if (result.runtime_vector.size() != options_.device.devices.size()) {
+      throw std::runtime_error("Runtime vector size mismatch");
+    }
+
+    for (const auto runtime : result.runtime_vector) {
+      out << "," << runtime;
+    }
+  }
+
   if (result.good()) {
 
     out

tools/profiler/src/rank_2k_operation_profiler.cu

@@ -733,7 +733,7 @@ bool Rank2KOperationProfiler::profile(
     rank_k_workspace_.arguments.pointer_mode = library::ScalarPointerMode::kHost;
 
     results_.back().status = profile_cutlass_(
-      results_.back().runtime,
+      results_.back(),
       options,
       operation,
       &rank_k_workspace_.arguments,

tools/profiler/src/rank_k_operation_profiler.cu

@@ -718,7 +718,7 @@ bool RankKOperationProfiler::profile(
     rank_k_workspace_.arguments.pointer_mode = library::ScalarPointerMode::kHost;
 
     results_.back().status = profile_cutlass_(
-      results_.back().runtime,
+      results_.back(),
       options,
       operation,
       &rank_k_workspace_.arguments,

tools/profiler/src/sparse_gemm_operation_profiler.cu

@@ -578,7 +578,7 @@ bool SparseGemmOperationProfiler::profile(
     gemm_workspace_.arguments.pointer_mode = library::ScalarPointerMode::kHost;
 
     results_.back().status = profile_cutlass_(
-      results_.back().runtime,
+      results_.back(),
       options,
       operation,
       &gemm_workspace_.arguments,

tools/profiler/src/symm_operation_profiler.cu

@@ -771,7 +771,7 @@ bool SymmOperationProfiler::profile(
     symm_workspace_.arguments.pointer_mode = library::ScalarPointerMode::kHost;
 
     results_.back().status = profile_cutlass_(
-      results_.back().runtime,
+      results_.back(),
       options,
       operation,
       &symm_workspace_.arguments,

tools/profiler/src/trmm_operation_profiler.cu

@@ -709,7 +709,7 @@ bool TrmmOperationProfiler::profile(
     trmm_workspace_.arguments.pointer_mode = library::ScalarPointerMode::kHost;
 
     results_.back().status = profile_cutlass_(
-      results_.back().runtime,
+      results_.back(),
       options,
       operation,
       &trmm_workspace_.arguments,