expose StoreT parameter for potential speed (#838)

* expose StoreT parameter for potential speed * add storeT to more elementwise --------- Co-authored-by: Haicheng Wu <[email protected]>
NVIDIA · Mar 10, 2023 · 86cae03 · 86cae03
1 parent 29801e3
commit 86cae03
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 6 deletions.
diff --git a/include/cutlass/epilogue/thread/linear_combination_bias_elementwise.h b/include/cutlass/epilogue/thread/linear_combination_bias_elementwise.h
@@ -62,6 +62,7 @@ template <
   int ElementsPerAccess,
   typename ElementwiseOp_ = Identity<ElementCompute_>,
   typename BinaryOp_ = plus<ElementCompute_>,
+  bool StoreT_ = true,
   typename ElementVector_ = ElementC_
 >
 class LinearCombinationBiasElementwise {
@@ -97,7 +98,7 @@ class LinearCombinationBiasElementwise {
   static bool const kStoreZ = true;
 
   /// If true, the 'T' tensor is stored
-  static bool const kStoreT = true;
+  static bool const kStoreT = StoreT_;
 
   /// Host-constructable parameters structure
   struct Params {

diff --git a/include/cutlass/epilogue/thread/linear_combination_bias_relu.h b/include/cutlass/epilogue/thread/linear_combination_bias_relu.h
@@ -204,7 +204,7 @@ template <
   typename ElementCompute_,
   typename ElementZ_,
   int ElementsPerAccess,
-  bool StoreT = true,
+  bool StoreT_ = true,
   typename ElementVector_ = ElementC_
 >
 class LinearCombinationBiasRelu {
@@ -238,7 +238,7 @@ class LinearCombinationBiasRelu {
   static bool const kStoreZ = true;
 
   /// If true, the 'T' tensor is stored
-  static bool const kStoreT = StoreT;
+  static bool const kStoreT = StoreT_;
 
   /// Host-constructable parameters structure
   struct Params {

diff --git a/include/cutlass/epilogue/thread/linear_combination_residual_block.h b/include/cutlass/epilogue/thread/linear_combination_residual_block.h
@@ -60,6 +60,7 @@ template <typename ElementOutput_, typename ElementAccumulator_,
           template <typename T> class BinaryOp1_,
           template <typename T> class UnaryOp_,
           template <typename T> class BinaryOp2_ = detail::NoOp,
+          bool StoreT_ = false,
           typename ElementVector_ = ElementC_>
 class LinearCombinationResidualBlock {
 public:
@@ -90,7 +91,7 @@ class LinearCombinationResidualBlock {
 
   static bool const kIsHeavy = true;
   static bool const kStoreZ = true;
-  static bool const kStoreT = false;
+  static bool const kStoreT = StoreT_;
 
   /// Host-constructable parameters structure
   struct Params {
@@ -182,11 +183,12 @@ template <typename ElementOutput_, typename ElementAccumulator_,
           template <typename T> class ActivationOp_,
           template <typename T> class BinaryOp1_,
           template <typename T> class UnaryOp_,
+          bool StoreT_,
           typename ElementVector_>
 class LinearCombinationResidualBlock<ElementOutput_, ElementAccumulator_,
           ElementCompute_, ElementC_, ElementsPerAccess,
           ActivationOp_, BinaryOp1_, UnaryOp_,
-          detail::NoOp, ElementVector_> {
+          detail::NoOp, StoreT_, ElementVector_> {
 public:
   static bool const kIsSingleSource = true;
 
@@ -214,7 +216,7 @@ class LinearCombinationResidualBlock<ElementOutput_, ElementAccumulator_,
 
   static bool const kIsHeavy = true;
   static bool const kStoreZ = true;
-  static bool const kStoreT = false;
+  static bool const kStoreT = StoreT_;
 
   /// Host-constructable parameters structure
   struct Params {