support arbitrary input/ouput and non-3-dim mm_out

tests/ap/__main__.py

@@ -1,3 +1 @@
-# import test_trivial_reduce
-# import test_binary_trivial_reduce
-import test_matmul_binary
+import test_matmul_epilogue

tests/ap/kernel_arg_id_util.py

@@ -15,6 +15,7 @@ def get_or_create_kernel_arg_id_manul_var_name(self, kernel_arg_id, cpp_var_name
     return self.all_kernel_arg_id2unique_name.get_or_create(kernel_arg_id, create)
 
   def get_in_tensor_data_ptr_var_name(self, in_ir_value_name):
+    print('in_ir_value_name: ', in_ir_value_name)
     ir_value = getattr(self.tensor_match_ctx, in_ir_value_name)
     kernel_arg_id = self.code_gen_ctx.in_tensor_data_ptr_kernel_arg_id(ir_value)
     create = self._get_creator(kernel_arg_id, self._create_in_tensor_data_ptr_var_name)
@@ -29,6 +30,7 @@ def _create_in_tensor_data_ptr_var_name(self):
     return name
 
   def get_out_tensor_data_ptr_var_name(self, out_ir_value_name):
+    out_ir_value_name = out_ir_value_name.replace("out", "output")
     ir_value = getattr(self.tensor_match_ctx, out_ir_value_name)
     kernel_arg_id = self.code_gen_ctx.out_tensor_data_ptr_kernel_arg_id(ir_value)
     create = self._get_creator(kernel_arg_id, self._create_out_tensor_data_ptr_var_name)

tests/ap/make_axpr.sh

@@ -14,11 +14,13 @@ FILENAMES_ARRAY=(
     "__main__"
     "topo_drr_pass"
     "op_convertion_drr_pass"
+    "umprime"
     "access_topo_drr"
     "abstract_drr"
+    "matmul_epilogue_pass_of_remove_functions"
     "ap_tpl_codegen"
     "matmul_binary_tpl"
-    "test_matmul_binary"
+    "test_matmul_epilogue"
 )
 for filename in "${FILENAMES_ARRAY[@]}"
 do

tests/ap/matmul/matmul.h

@@ -89,7 +89,9 @@ struct GemmEpilogueParams {
   std::vector<int64_t> input0_shape;
   std::vector<int64_t> input1_shape;
   std::vector<const void *> epilogue_in_ptrs;
+  std::vector<void *> epilogue_out_ptrs;
   std::vector<std::vector<int64_t>> epilogue_in_shapes;
+  std::vector<std::vector<int64_t>> epilogue_out_shapes;
 
   GemmEpilogueParams() {}
   GemmEpilogueParams(cudaStream_t stream, const void *input, const void *weight,
@@ -156,16 +158,23 @@ struct GemmEpilogueParams {
     shape_args.ldc_bias = (!bias || is_C_bias) ? 0 : n;
   }
 
-  void SetEpilogues(const std::vector<const void *> &in_ptrs) {
+  void SetEpilogues(const std::vector<const void *> &in_ptrs, 
+                    const std::vector< void *> &out_ptrs) {
     epilogue_in_ptrs = in_ptrs;
+    epilogue_out_ptrs = out_ptrs;
   }
 
   void
   SetEpilogueAndShapes(const std::vector<const void *> &in_ptrs,
-                       const std::vector<std::vector<int64_t>> &in_shapes) {
+                       const std::vector<std::vector<int64_t>> &in_shapes,
+                       const std::vector<void *> &out_ptrs,
+                       const std::vector<std::vector<int64_t>> &out_shapes) {
     ASSERT_CHECK(in_ptrs.size() == in_shapes.size());
     epilogue_in_ptrs = in_ptrs;
     epilogue_in_shapes = in_shapes;
+    ASSERT_CHECK(out_ptrs.size() == out_shapes.size());
+    epilogue_out_ptrs = out_ptrs;
+    epilogue_out_shapes = out_shapes;
   }
 };
 

tests/ap/matmul/tests/matmul_binary_kernel.cu

@@ -47,13 +47,15 @@ void MatmulAddBinaryKernel(
     cudaStream_t *stream, const void *input, const void *weight,
     const void *bias, void *output,
     const std::vector<const void *> &epilogue_ins,
+    const std::vector<void *> &epilogue_outs,
     const std::vector<int64_t> &input_shape,
     const std::vector<int64_t> &weight_shape,
     const std::vector<int64_t> &bias_shape,
-    const std::vector<std::vector<int64_t>> &epilogue_shapes) {
+    const std::vector<std::vector<int64_t>> &epilogue_in_shapes,
+    const std::vector<std::vector<int64_t>> &epilogue_out_shapes) {
   GemmEpilogueParams params(*stream, input, weight, bias, output, input_shape,
                             weight_shape, bias_shape);
-  params.SetEpilogueAndShapes(epilogue_ins, epilogue_shapes);
+  params.SetEpilogueAndShapes(epilogue_ins, epilogue_in_shapes, epilogue_outs, epilogue_out_shapes);
 
 #if AP_ENABLE_AUTOTUNE
 #if AP_USE_FLOAT16

tests/ap/matmul_binary_tpl.py

@@ -16,6 +16,11 @@ def is_in_tensor_karg(kernel_arg_id):
     )
     return kernel_arg_id_type_name == "InTensorDataPtrKernelArgId"
 
+def is_out_tensor_karg(kernel_arg_id):
+    kernel_arg_id_type_name = f"{type(kernel_arg_id)}".replace("<class '", "").replace(
+        "'>", ""
+    )
+    return kernel_arg_id_type_name == "OutTensorDataPtrKernelArgId"
 
 class MatmulBinaryTemplate:
     def __init__(
@@ -39,6 +44,7 @@ def __init__(
         )
         self.input_dim_karg_to_shape_access = MutableOrderedDict()
         self.input_tensor_karg_to_shape_access = MutableOrderedDict()
+        self.output_tensor_karg_to_shape_access = MutableOrderedDict()
         self.kernel_name = "MatmulBinaryKernel"
         self.library_name = "matmul_binary_kernel"
 
@@ -105,6 +111,11 @@ def get_kernel_arg_runtime_getters(self):
             lambda pair: pair[0].runtime_getter, all_kernel_arg_id_and_unique_names
         )
 
+    def init_outputs(self):
+        out_tensor_data_nums = self.mut_kernel_arg_id_registry.out_tensor_data_ptr_seq_no
+        stmt = map(lambda i: f"out{i}", range(out_tensor_data_nums + 1))
+        return "T " + f", ".join(stmt) + ";"
+
     def get_kernel_arg_types(self):
         all_kernel_arg_id_and_unique_names = (
             self.mut_kernel_arg_id_registry.all_kernel_arg_id2unique_name.items()
@@ -159,6 +170,7 @@ def get_epilogue_arguments_init_str(
         def declare_epilogue_arguments_assign(pair):
             kernel_arg_id = pair[0]
             is_in_tensor_type = is_in_tensor_karg(kernel_arg_id)
+            is_out_tensor_type = is_out_tensor_karg(kernel_arg_id)
 
             var_name = pair[1]
             field_name = self.kernel_arg_translator.get_param_struct_field_name(
@@ -169,30 +181,36 @@ def get_in_tensor_statement():
                 param_name_for_var = self.input_tensor_karg_to_shape_access[var_name]
                 return f"reinterpret_cast<const {output_dtype} *>({params_name}.{param_name_for_var})"
 
+            def get_out_tensor_statement():
+                param_name_for_var = self.output_tensor_karg_to_shape_access[var_name]
+                return f"reinterpret_cast<{output_dtype} *>({params_name}.{param_name_for_var})"
+
             def get_dim_expr_statement():
                 param_name_for_var = self.input_dim_karg_to_shape_access[var_name]
                 return f"{params_name}.{param_name_for_var}"
 
             statement = (
                 get_in_tensor_statement()
                 if is_in_tensor_type
-                else get_dim_expr_statement()
+                else get_out_tensor_statement()
+                     if is_out_tensor_type
+                     else get_dim_expr_statement()
             )
             return f"{obj_name}.{field_name} = {statement};"
 
         generated_kernel_arg_id_and_names = (
             self.mut_kernel_arg_id_registry.generated_kernel_arg_id2unique_name.items()
         )
+
         return f"\n{indent}".join(
             map(declare_epilogue_arguments_assign, generated_kernel_arg_id_and_names)
         )
 
-    def get_params_epilogue_ptrs_init_str(self, obj_name, indent):
+    def get_params_epilogue_ptrs_init_str(self, in_obj_name, out_obj_name, indent):
         in_tensor_id = 0
-
-        def declare_params_epilogue_arguments_assign(pair):
+        def declare_in_params_epilogue_arguments_assign(pair):
             def get_creator():
-                return f"{obj_name}[{in_tensor_id}]"
+                return f"{in_obj_name}[{in_tensor_id}]"
 
             kernel_arg_id = pair[0]
             is_in_tensor_type = is_in_tensor_karg(kernel_arg_id)
@@ -201,7 +219,7 @@ def generate_statement():
                 self.input_tensor_karg_to_shape_access.get_or_create(
                     pair[1], get_creator
                 )
-                statement = f"{obj_name}.push_back({pair[1]});"
+                statement = f"{in_obj_name}.push_back({pair[1]});"
                 in_tensor_id = in_tensor_id + 1
                 return statement
 
@@ -210,13 +228,39 @@ def generate_statement():
         generated_kernel_arg_id_and_names = (
             self.mut_kernel_arg_id_registry.generated_kernel_arg_id2unique_name.items()
         )
-        return f"\n{indent}".join(
-            map(
-                declare_params_epilogue_arguments_assign,
+        in_str_list = map(
+                declare_in_params_epilogue_arguments_assign,
                 generated_kernel_arg_id_and_names,
-            )
         )
 
+        out_tensor_id = 0
+        def declare_out_params_epilogue_arguments_assign(pair):
+            def get_creator():
+                return f"{out_obj_name}[{out_tensor_id}]"
+
+            kernel_arg_id = pair[0]
+            is_out_tensor_type = is_out_tensor_karg(kernel_arg_id)
+
+            def generate_statement():
+                self.output_tensor_karg_to_shape_access.get_or_create(
+                    pair[1], get_creator
+                )
+                statement = f"{out_obj_name}.push_back({pair[1]});"
+                out_tensor_id = out_tensor_id + 1
+                return statement
+
+            return generate_statement() if is_out_tensor_type else ""
+
+        out_str_list = map(
+            declare_out_params_epilogue_arguments_assign,
+            generated_kernel_arg_id_and_names,
+        )
+        str_list = filter(
+            lambda ss: ss != "",
+            [*in_str_list, *out_str_list]
+        )
+        return f"\n{indent}".join(str_list)
+
     def get_params_input_shape_init_str(self, input_name, input_shape_kargs, indent):
         def init_input_shape_with_args(i):
             def get_creator():
@@ -264,9 +308,9 @@ def make_project(
   // Note: need to support vectorized operation
   __forceinline__ __host__ __device__
   T operator()(T x, const Arguments& args, const MatrixCoord& coord) const {
-    T out;
+    AP_OUTPUTS_INIT
     AP_GENERATED_BINARY_EPILOGUE_STRING
-    return out;
+    return out0;
   }
 };
 
@@ -303,9 +347,10 @@ def make_project(
       *cuda_stream_ptr, ${input0}, ${input1}, nullptr, ${output}, ${input0}_shape, ${input1}_shape, std::vector<int64_t>{});
 
   std::vector<const void *> epilogue_in_ptrs;
+  std::vector<void *> epilogue_out_ptrs;
   AP_PARAMS_EPILOGUE_PTRS_INIT
 
-  params.SetEpilogues(epilogue_in_ptrs);
+  params.SetEpilogues(epilogue_in_ptrs, epilogue_out_ptrs);
 
 #if AP_ENABLE_AUTOTUNE
   AP_AUTOTUNE_${output_dtype}(ap::RunMatmulWithVariadicKernel);
@@ -321,6 +366,7 @@ def make_project(
             code_template.replace(
                 "AP_GENERATED_BINARY_EPILOGUE_STRING", trivial_code_str
             )
+            .replace("AP_OUTPUTS_INIT", self.init_outputs())
             .replace("AP_KERNEL_ARGS_DECLARE", self.get_kernel_arg_list_str())
             .replace(
                 "AP_PARAMS_INPUT0_SHAPE_INIT",
@@ -336,7 +382,7 @@ def make_project(
             )
             .replace(
                 "AP_PARAMS_EPILOGUE_PTRS_INIT",
-                self.get_params_epilogue_ptrs_init_str("epilogue_in_ptrs", indent="  "),
+                self.get_params_epilogue_ptrs_init_str("epilogue_in_ptrs", "epilogue_out_ptrs", indent="  "),
             )
             .replace(
                 "AP_EPILOGUE_ARGUMENTS_FIELDS",
@@ -356,7 +402,7 @@ def make_project(
             .replace("${k_value}", f"{input0_shape_kargs[-1].value}")
             .replace("${n_value}", f"{input1_shape_kargs[-1].value}")
         )
-
+        print('cuda code is: ', code)
         source_dir = "/work/abstract_pass/Athena/tests/ap/matmul"
         cutlass_dir = "/work/abstract_pass/Athena/tests/ap/matmul/cutlass"
         compile_cmd = (