CUTLASS 3.1 (#915)

Co-authored-by: Aniket Shivam <ashivam@nvidia.com>

include/cutlass/epilogue/thread/linear_combination.h

@@ -52,7 +52,7 @@ namespace thread {
 
 /// Applies a linear combination operator to an array of elements.
 ///
-/// D = alpha * accumulator + beta * source + uniform
+/// D = alpha * accumulator + beta * source
 ///
 template <
   typename ElementOutput_,                             ///< Data type used to load and store tensors
@@ -69,6 +69,7 @@ class LinearCombination {
 public:
 
   using ElementOutput = ElementOutput_;
+  using ElementSource = ElementSource_;
   using ElementAccumulator = ElementAccumulator_;
   using ElementCompute = ElementCompute_;
   using ElementC = ElementSource_;
@@ -77,14 +78,15 @@ public:
   static int const kCount = Count;
   static const ScaleType::Kind kScale = Scale;
   using FragmentOutput = Array<ElementOutput, kCount>;
+  using FragmentSource = Array<ElementSource, kCount>;
   using FragmentAccumulator = Array<ElementAccumulator, kCount>;
-  using ComputeFragment = Array<ElementCompute, kCount>;
+  using FragmentCompute = Array<ElementCompute, kCount>;
 
-  using ParamsBase = LinearCombinationParams;
   static FloatRoundStyle const kRound = Round;
 
   /// Host-constructable parameters structure
-  struct Params : ParamsBase{
+  struct Params 
+  {
     ElementCompute alpha;                  ///< scales accumulators
     ElementCompute beta;                   ///< scales source tensor
     ElementCompute const *alpha_ptr;       ///< pointer to accumulator scalar - if not null, loads it from memory
@@ -92,10 +94,6 @@ public:
 
     CUTLASS_HOST_DEVICE
     Params():
-      ParamsBase(
-        ElementCompute(1),
-        ElementCompute(0)
-      ),
       alpha(ElementCompute(1)),
       beta(ElementCompute(0)),
       alpha_ptr(nullptr),
@@ -106,14 +104,12 @@ public:
       ElementCompute alpha,
       ElementCompute beta
     ):
-      ParamsBase(alpha, beta),
       alpha(alpha), beta(beta), alpha_ptr(nullptr), beta_ptr(nullptr) { }
 
     CUTLASS_HOST_DEVICE
     Params(
       ElementCompute alpha
     ):
-      ParamsBase(alpha, ElementCompute(0)),
       alpha(alpha), beta(0), alpha_ptr(nullptr), beta_ptr(nullptr) { }
 
     CUTLASS_HOST_DEVICE
@@ -121,28 +117,13 @@ public:
       ElementCompute const *alpha_ptr,
       ElementCompute const *beta_ptr
     ):
-      ParamsBase(*alpha_ptr, *beta_ptr),
       alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) { }
 
     CUTLASS_HOST_DEVICE
     Params(
       ElementCompute const *alpha_ptr
     ):
-      ParamsBase(*alpha_ptr, ElementCompute(0)),
       alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(nullptr) { }
-
-    CUTLASS_HOST_DEVICE
-    Params(
-      ParamsBase const& base
-    ): ParamsBase(base), alpha_ptr(nullptr), beta_ptr(nullptr) {
-      #if defined(__CUDA_ARCH__)
-      alpha = reinterpret_cast<ElementCompute const&>(base.alpha_data);
-      beta = reinterpret_cast<ElementCompute const&>(base.beta_data);
-      #else
-      memcpy( alpha, base.alpha_data, sizeof(ElementCompute) );
-      memcpy( beta, base.alpha_data, sizeof(ElementCompute) );
-      #endif
-    }
   };
 
 private:
@@ -183,30 +164,73 @@ public:
     }
   }
 
-  /// Computes linear scaling: D = alpha * accumulator + beta * source
+  /// Computes intermediate: X = beta * source
+  CUTLASS_HOST_DEVICE
+  FragmentCompute compute_intermediate(
+      FragmentSource const &source) const {
+
+    // Convert source to internal compute numeric type
+    NumericArrayConverter<ElementCompute, ElementSource, kCount, Round> source_converter;
+    FragmentCompute converted_source = source_converter(source);
+
+    if (Scale == ScaleType::NoBetaScaling) {
+      return converted_source;
+    }
+    else {
+      multiplies<FragmentCompute> mul_source;
+      return mul_source(beta_, converted_source);
+    }
+  }
+
+  /// Computes linear scaling with intermediate: D = alpha * accumulator + X
+  CUTLASS_HOST_DEVICE
+  FragmentOutput with_intermediate(
+      FragmentAccumulator const& accumulator,
+      FragmentCompute const& intermediate) const {
+
+    // Convert accumulator to internal compute numeric type
+    NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
+
+    // Convert to destination numeric type
+    NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
+    
+    FragmentCompute converted_accumulator = accumulator_converter(accumulator);
+
+    if (Scale == ScaleType::Nothing) {
+      return destination_converter(converted_accumulator);
+    } else {
+      // Perform binary operations
+      multiply_add<FragmentCompute> mul_add_accumulator;
+      FragmentCompute computed_output = mul_add_accumulator(alpha_, converted_accumulator, intermediate);
+
+      return destination_converter(computed_output);
+    }
+  }
+
+  /// Computes linear scaling with source: D = alpha * accumulator + beta * source
   CUTLASS_HOST_DEVICE
   FragmentOutput operator()(
-    FragmentAccumulator const &accumulator,
-    FragmentOutput const &source) const {
+      FragmentAccumulator const &accumulator,
+      FragmentSource const &source) const {
 
-    // Convert source to interal compute numeric type
-    NumericArrayConverter<ElementCompute, ElementOutput, kCount, Round> source_converter;
+    // Convert source to internal compute numeric type
+    NumericArrayConverter<ElementCompute, ElementSource, kCount, Round> source_converter;
     NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
 
     // Convert to destination numeric type
     NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
 
-    ComputeFragment converted_source = source_converter(source);
-    ComputeFragment converted_accumulator = accumulator_converter(accumulator);
+    FragmentCompute converted_source = source_converter(source);
+    FragmentCompute converted_accumulator = accumulator_converter(accumulator);
 
     if (Scale == ScaleType::Nothing)
       return destination_converter(converted_accumulator);
 
     // Perform binary operations
-    ComputeFragment intermediate;
+    FragmentCompute intermediate;
 
-    multiplies<ComputeFragment> mul_add_source;
-    multiply_add<ComputeFragment> mul_add_accumulator;
+    multiplies<FragmentCompute> mul_add_source;
+    multiply_add<FragmentCompute> mul_add_accumulator;
 
     if (Scale == ScaleType::NoBetaScaling)
       intermediate = converted_source;
@@ -221,7 +245,7 @@ public:
   /// Computes linear scaling: D = alpha * accumulator
   CUTLASS_HOST_DEVICE
   FragmentOutput operator()(
-    FragmentAccumulator const &accumulator) const {
+      FragmentAccumulator const &accumulator) const {
 
     // Convert source to interal compute numeric type
     NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
@@ -229,14 +253,14 @@ public:
     // Convert to destination numeric type
     NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
 
-    ComputeFragment converted_accumulator = accumulator_converter(accumulator);
+    FragmentCompute converted_accumulator = accumulator_converter(accumulator);
 
     if (Scale == ScaleType::Nothing)
       return destination_converter(converted_accumulator);
 
     // Perform binary operations
-    ComputeFragment intermediate;
-    multiplies<ComputeFragment> mul_accumulator;
+    FragmentCompute intermediate;
+    multiplies<FragmentCompute> mul_accumulator;
 
     intermediate = mul_accumulator(alpha_, converted_accumulator);    // D = alpha * Accum