Support parallel split K mode for porfiling

tools/library/include/cutlass/library/handle.h

@@ -1,5 +1,5 @@
 /***************************************************************************************************
- * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2017-2022, 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:
@@ -338,6 +338,9 @@ using HandlePtr = std::unique_ptr<Handle>;
 /// Finds conv2d operation instances with Conv2d::ElementC = Reduction::ElementWorkspace
 Operation const* find_conv_operation_for_parallel_reduction(Operation const *operation);
 /////////////////////////////////////////////////////////////////////////////////////////////////
+/// Finds gemm operation instances with ElementC = Reduction::ElementWorkspace
+Operation const* find_gemm_operation_for_parallel_reduction(Operation const *operation);
+/////////////////////////////////////////////////////////////////////////////////////////////////
 
 } // namespace library
 } // namespace cutlass

tools/library/include/cutlass/library/library.h

@@ -590,7 +590,8 @@ class Operation {
     void const *configuration) const = 0;
 
   virtual uint64_t get_device_workspace_size(
-    void const *configuration) const = 0;
+    void const *configuration,
+    void const *arguments = nullptr) const = 0;
 
   virtual Status initialize(
     void const *configuration, 

tools/library/src/conv2d_operation.h

@@ -1,5 +1,5 @@
 /***************************************************************************************************
- * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2017-2022, 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:
@@ -272,7 +272,8 @@ class Conv2dOperation : public Conv2dOperationBase<Operator_> {
 
   /// Gets the device-side workspace
   virtual uint64_t get_device_workspace_size(
-    void const *configuration_ptr) const {
+    void const *configuration_ptr,
+    void const *arguments_ptr = nullptr) const {
 
     OperatorArguments args;
 

tools/library/src/conv3d_operation.h

@@ -1,5 +1,5 @@
 /***************************************************************************************************
- * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2017-2022, 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:
@@ -266,7 +266,8 @@ class Conv3dOperation : public Conv3dOperationBase<Operator_> {
 
   /// Gets the device-side workspace
   virtual uint64_t get_device_workspace_size(
-    void const *configuration_ptr) const {
+    void const *configuration_ptr,
+    void const *arguments_ptr = nullptr) const {
 
     OperatorArguments args;
 

tools/library/src/gemm_operation.h

@@ -1,5 +1,5 @@
 /***************************************************************************************************
- * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2017-2022, 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:
@@ -242,7 +242,8 @@ class GemmOperation : public GemmOperationBase<Operator_> {
 
   /// Gets the device-side workspace
   virtual uint64_t get_device_workspace_size(
-    void const *configuration_ptr) const {
+    void const *configuration_ptr,
+    void const *arguments_ptr = nullptr) const {
 
     OperatorArguments args;
 
@@ -443,7 +444,8 @@ class GemmSparseOperation : public GemmOperationBase<Operator_> {
 
   /// Gets the device-side workspace
   virtual uint64_t get_device_workspace_size(
-    void const *configuration_ptr) const {
+    void const *configuration_ptr,
+    void const *arguments_ptr = nullptr) const {
 
     OperatorArguments args;
 
@@ -569,7 +571,7 @@ class GemmUniversalOperation : public GemmOperationBase<Operator_> {
     operator_args.ldb = (configuration->ldb);
     operator_args.ldc = (configuration->ldc);
     operator_args.ldd = (configuration->ldd);
-    
+
     return Status::kSuccess;
   }
 
@@ -649,7 +651,8 @@ class GemmUniversalOperation : public GemmOperationBase<Operator_> {
 
   /// Gets the device-side workspace
   virtual uint64_t get_device_workspace_size(
-    void const *configuration_ptr) const {
+    void const *configuration_ptr,
+    void const *arguments_ptr) const {
 
     OperatorArguments args;
 
@@ -661,6 +664,14 @@ class GemmUniversalOperation : public GemmOperationBase<Operator_> {
       return 0;
     }
 
+    status = update_arguments_(
+      args,
+      static_cast<GemmUniversalArguments const *>(arguments_ptr));
+
+    if (status != Status::kSuccess) {
+      return 0;
+    }
+
     uint64_t size = Operator::get_workspace_size(args);
 
     return size;
@@ -855,7 +866,8 @@ class GemmPlanarComplexOperation : public GemmOperationBase<Operator_> {
 
   /// Gets the device-side workspace
   virtual uint64_t get_device_workspace_size(
-    void const *configuration_ptr) const {
+    void const *configuration_ptr,
+    void const *arguments_ptr = nullptr) const {
 
     OperatorArguments args;
 
@@ -1055,7 +1067,8 @@ class GemmPlanarComplexArrayOperation : public GemmOperationBase<Operator_> {
 
   /// Gets the device-side workspace
   virtual uint64_t get_device_workspace_size(
-    void const *configuration_ptr) const {
+    void const *configuration_ptr,
+    void const *arguments_ptr = nullptr) const {
 
     OperatorArguments args;
 

tools/library/src/handle.cu

@@ -1,5 +1,5 @@
 /***************************************************************************************************
- * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2017-2022, 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:
@@ -1098,6 +1098,59 @@ Operation const* find_conv_operation_for_parallel_reduction(Operation const *ope
 
   return nullptr;
 }
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Finds gemm operation instances with Gemm::ElementC = Reduction::ElementWorkspace
+Operation const* find_gemm_operation_for_parallel_reduction(Operation const *operation) {
+
+  GemmDescription const &gemm_desc = 
+    static_cast<GemmDescription const &>(operation->description());
+
+  // if the curren gemm operation accumulator and output data type match return operation
+  if(gemm_desc.tile_description.math_instruction.element_accumulator == gemm_desc.C.element) {
+    return operation;
+  }
+
+  // find gemm operation to match gemm output and reduction workspace data type
+  GemmFunctionalKey key(
+    library::Provider::kCUTLASS,
+    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.tile_description.math_instruction.element_accumulator);
+
+  // gemm operation table
+  auto gemm_operations = Singleton::get().operation_table.gemm_operations;
+
+  // find ConvFunctionalKey in gemm operation table
+  auto operators_it = gemm_operations.find(key);
+
+  if (operators_it == gemm_operations.end()) {
+    return nullptr;
+  }
+
+  if (operators_it->second.empty()) {
+    return nullptr;
+  }
+
+  // A and B uses the same alignment in the generator.py
+  int alignment = gemm_desc.A.alignment;
+
+  // gemm operation for same compute capability and iterator algorithm
+  GemmPreferenceKey preference_key(
+    gemm_desc.tile_description.minimum_compute_capability, 
+    alignment);
+
+  return find_gemm_operation(operators_it, preference_key);
+}
+
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 } // namespace library

tools/library/src/reduction/init_reduction_operations.cu

@@ -1,5 +1,5 @@
 /***************************************************************************************************
- * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2017-2022, 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:

tools/library/src/reduction/reduction_device.cu

@@ -1,5 +1,5 @@
 /***************************************************************************************************
- * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2017-2022, 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:
@@ -46,7 +46,7 @@ void initialize_reduce_add_linear_combination_f32_f32_f16(Manifest &manifest) {
 
   using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombination<
     ElementOutput,
-    128 / cutlass::sizeof_bits<ElementOutput>::value,
+    128 / cutlass::sizeof_bits<ElementWorkspace>::value,
     ElementAccumulator,
     ElementCompute
   >;
@@ -81,7 +81,7 @@ void initialize_reduce_add_linear_combination_f32_f32_f32(Manifest &manifest) {
 
   using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombination<
     ElementOutput,
-    128 / cutlass::sizeof_bits<ElementOutput>::value,
+    128 / cutlass::sizeof_bits<ElementWorkspace>::value,
     ElementAccumulator,
     ElementCompute
   >;
@@ -115,7 +115,7 @@ void initialize_reduce_add_linear_combination_cf32_cf32_cf32(Manifest &manifest)
 
   using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombination<
     ElementOutput,
-    128 / cutlass::sizeof_bits<ElementOutput>::value,
+    128 / cutlass::sizeof_bits<ElementWorkspace>::value,
     ElementAccumulator,
     ElementCompute
   >;
@@ -140,6 +140,5 @@ void initialize_reduce_add_linear_combination_cf32_cf32_cf32(Manifest &manifest)
   ));
 }
 
-
 } 
 }

tools/library/src/reduction/reduction_operation.h

@@ -30,6 +30,7 @@
 #include <iostream>
 #include "cutlass/cutlass.h"
 #include "cutlass/epilogue/thread/linear_combination.h"
+#include "cutlass/epilogue/thread/linear_combination_clamp.h"
 #include "cutlass/reduction/thread/reduction_operators.h"
 #include "cutlass/reduction/device/reduce_split_k.h"
 
@@ -180,7 +181,8 @@ class ReductionOperation : public Operation {
 
   /// Gets the device-side workspace
   virtual uint64_t get_device_workspace_size(
-    void const *configuration_ptr) const {
+    void const *configuration_ptr,
+    void const *arguments_ptr = nullptr) const {
 
     OperatorArguments args;
 

tools/library/src/reference/conv_reference_operation.h

@@ -1,5 +1,5 @@
 /***************************************************************************************************
- * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2017-2022, 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:
@@ -403,7 +403,8 @@ class ConvReferenceOperation : public Operation {
   }
 
   virtual uint64_t get_device_workspace_size(
-    void const *configuration) const {
+    void const *configuration,
+    void const *arguments = nullptr) const {
 
     return 0;
   }

tools/library/src/reference/gemm_reference_operation.h

@@ -1,5 +1,5 @@
 /***************************************************************************************************
- * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2017-2022, 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:
@@ -161,7 +161,8 @@ class GemmReferenceOperation : public Operation {
   }
 
   virtual uint64_t get_device_workspace_size(
-    void const *configuration) const {
+    void const *configuration,
+    void const *arguments = nullptr) const {
 
     return 0;
   }