Adding more Threadblock Tiles for Mixed-input TensorOp (BF16 * S8) in cutlass_library (#1132)

* Adding more tiles in the cutlass_library for mixed-input support.

* fix rebase issue

* more tiles to upcast a

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

@@ -215,6 +215,18 @@ struct GemmPreferenceKey {
     return compute_capability == rhs.compute_capability;
   }
 };
+/////////////////////////////////////////////////////////////////////////////////////////////////
+inline
+std::ostream& operator<< (std::ostream& out, const cutlass::library::GemmPreferenceKey& key) {
+    out << "{\n"
+      << "compute_capability : " << key.compute_capability << std::endl
+      << "alignment          : " << key.alignment << std::endl
+      << "}";
+  
+  return out;
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 

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

@@ -172,6 +172,7 @@ enum class MathOperationID {
   kAdd,
   kMultiplyAdd,
   kMultiplyAddSaturate,
+  kMultiplyAddMixedInputUpcast,
   kMultiplyAddFastBF16,
   kMultiplyAddFastF16,
   kMultiplyAddFastF32,

tools/library/src/library_internal.h

@@ -174,6 +174,10 @@ template <> struct MathOperationMap<cutlass::arch::OpMultiplyAddSaturate> {
   static MathOperationID const kId = MathOperationID::kMultiplyAddSaturate;
 };
 
+template <> struct MathOperationMap<cutlass::arch::OpMultiplyAddMixedInputUpcast> {
+  static MathOperationID const kId = MathOperationID::kMultiplyAddMixedInputUpcast;
+};
+
 template <> struct MathOperationMap<cutlass::arch::OpMultiplyAddComplex> {
   static MathOperationID const kId = MathOperationID::kMultiplyAddComplex;
 };

tools/library/src/reduction/init_reduction_operations.cu

@@ -44,6 +44,7 @@ namespace library {
 ///////////////////////////////////////////////////////////////////////////////////////////////
 void initialize_reduce_add_linear_combination_f16_f16_f16(Manifest &manifest);
 void initialize_reduce_add_linear_combination_f32_f32_f16(Manifest &manifest);
+void initialize_reduce_add_linear_combination_f32_f32_bf16(Manifest &manifest);
 void initialize_reduce_add_linear_combination_f32_f32_f32(Manifest &manifest);
 void initialize_reduce_add_linear_combination_f64_f64_f64(Manifest &manifest);
 void initialize_reduce_add_linear_combination_cf32_cf32_cf32(Manifest &manifest);
@@ -55,6 +56,7 @@ void initialize_all_reduction_op(Manifest &manifest) {
 
   initialize_reduce_add_linear_combination_f16_f16_f16(manifest);
   initialize_reduce_add_linear_combination_f32_f32_f16(manifest);
+  initialize_reduce_add_linear_combination_f32_f32_bf16(manifest);
   initialize_reduce_add_linear_combination_f32_f32_f32(manifest);
   initialize_reduce_add_linear_combination_f64_f64_f64(manifest);
   initialize_reduce_add_linear_combination_cf32_cf32_cf32(manifest);

tools/library/src/reduction/reduction_device.cu

@@ -112,6 +112,40 @@ void initialize_reduce_add_linear_combination_f32_f32_f16(Manifest &manifest) {
   ));
 }
 
+void initialize_reduce_add_linear_combination_f32_f32_bf16(Manifest &manifest) {
+
+  using ElementWorkspace = float;
+  using ElementAccumulator = float;
+  using ElementOutput = cutlass::bfloat16_t;
+  using ElementCompute = float;
+
+  using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombination<
+    ElementOutput,
+    128 / cutlass::sizeof_bits<ElementWorkspace>::value,
+    ElementAccumulator,
+    ElementCompute
+  >;
+
+  using ReductionOp = cutlass::reduction::thread::ReduceAdd<
+    ElementAccumulator,
+    typename EpilogueOutputOp::ElementAccumulator,
+    EpilogueOutputOp::kCount
+  >;
+
+  using Operation_reduce_add_linear_combination_f32_f32_bf16 = cutlass::reduction::device::ReduceSplitK<
+    cutlass::reduction::kernel::ReduceSplitK<
+      cutlass::MatrixShape<4, 32 * EpilogueOutputOp::kCount>,
+      EpilogueOutputOp,
+      ReductionOp
+    >
+  >;
+
+  manifest.append(new ReductionOperation<
+    Operation_reduce_add_linear_combination_f32_f32_bf16>(
+      "reduce_add_linear_combination_f32_f32_bf16"
+  ));
+}
+
 
 void initialize_reduce_add_linear_combination_f32_f32_f32(Manifest &manifest) {
 

tools/library/src/reference/gemm_fp_mixed_input.cu

@@ -107,7 +107,23 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) {
   make_gemm_real_canonical_layouts<
     int8_t,
     bfloat16_t,
+    float,
+    float,
+    float
+  >(manifest);
+
+  make_gemm_real_canonical_layouts<
+    int8_t,
     bfloat16_t,
+    bfloat16_t,
+    float,
+    float
+  >(manifest);
+
+  make_gemm_real_canonical_layouts<
+    bfloat16_t,
+    uint8_t,
+    float,
     float,
     float
   >(manifest);