chore: import upstream snapshot with attribution

python/paddle/apy/device/cuda/compile_command_util.py

@@ -0,0 +1,48 @@
+# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import ap
+
+
+class CompileCommandGenerator:
+    def __init__(self):
+        self.file_ext = "cu"
+        self.op_type2generate_func = ap.OrderedDict(
+            [
+                ['matmul', self.generate_matmul_compile_command],
+            ]
+        )
+
+    def __call__(self, op_type, tpl_dirname, library_name):
+        return self.op_type2generate_func[op_type](tpl_dirname, library_name)
+
+    def generate_matmul_compile_command(self, tpl_dirname, library_name):
+        matmul_source_dir = f"{tpl_dirname}/matmul"
+
+        compile_cmd = "nvcc -std=c++20 -O3 -Xcompiler=-fPIC -arch=sm_80 --expt-relaxed-constexpr"
+        compile_cmd = compile_cmd + " -I ${AP_CUTLASS_DIR}/include"
+        compile_cmd = compile_cmd + " -I ${AP_CUTLASS_DIR}/tools/util/include"
+        compile_cmd = compile_cmd + " -I " + matmul_source_dir
+        compile_cmd = (
+            compile_cmd
+            + " -DCUTLASS_ENABLE_TENSOR_CORE_MMA=1 -DCUTLASS_DEBUG_TRACE_LEVEL=0"
+        )
+        compile_cmd = (
+            compile_cmd + " -DAP_ENABLE_AUTOTUNE=1 -DAP_ENABLE_DEBUG=0"
+        )
+        compile_cmd = (
+            compile_cmd
+            + f" --shared {library_name}.{self.file_ext} -o lib{library_name}.so"
+        )
+        return compile_cmd

python/paddle/apy/device/dcu/compile_command_util.py

@@ -0,0 +1,46 @@
+# Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import ap
+
+
+class CompileCommandGenerator:
+    def __init__(self):
+        self.file_ext = "cu"
+        self.op_type2generate_func = ap.OrderedDict(
+            [
+                ['matmul', self.generate_matmul_compile_command],
+            ]
+        )
+
+    def __call__(self, op_type, tpl_dirname, library_name):
+        return self.op_type2generate_func[op_type](tpl_dirname, library_name)
+
+    def generate_matmul_compile_command(self, tpl_dirname, library_name):
+        matmul_source_dir = f"{tpl_dirname}/matmul"
+
+        compile_cmd = (
+            "hipcc -std=c++17 -O3 -fPIC --offload-arch=gfx928 -Wno-return-type"
+        )
+        compile_cmd = compile_cmd + " -I ${AP_CUTLASS_DIR}/include"
+        compile_cmd = compile_cmd + " -I ${AP_CUTLASS_DIR}/tools/util/include"
+        compile_cmd = compile_cmd + " -I " + matmul_source_dir
+        compile_cmd = (
+            compile_cmd + " -DAP_ENABLE_AUTOTUNE=0 -DAP_ENABLE_DEBUG=0"
+        )
+        compile_cmd = (
+            compile_cmd
+            + f" --shared {library_name}.{self.file_ext} -o lib{library_name}.so"
+        )
+        return compile_cmd

python/paddle/apy/matmul_pass/__main__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import matmul_variadic_ptn  # noqa: F401

python/paddle/apy/matmul_pass/abstract_drr.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+class DrrPass:
+    def make_drr_ctx(self):
+        drr_ctx = DrrCtx()  # noqa: F821
+        drr_ctx.set_drr_pass_type(self.drr_pass_type())
+        drr_ctx.init_source_pattern(self.source_pattern)
+        drr_ctx.init_constraint_func(self.constraint)
+        drr_ctx.init_result_pattern(self.result_pattern)
+        return drr_ctx
+
+    def constraint(self, o, t):
+        return True
+
+    def drr_pass_type(self):
+        return "abstract_drr_pass_type"
+
+
+class register_drr_pass:
+    def __init__(self, pass_name, nice):
+        self.pass_name = pass_name
+        self.nice = nice
+
+    def __call__(self, drr_pass_cls):
+        Registry.abstract_drr_pass(  # noqa: F821
+            self.pass_name, self.nice, drr_pass_cls
+        )
+        return drr_pass_cls

python/paddle/apy/matmul_pass/access_topo_drr.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+class DrrPass:
+    def make_drr_ctx(self):
+        drr_ctx = DrrCtx()  # noqa: F821
+        drr_ctx.set_drr_pass_type(self.drr_pass_type())
+        drr_ctx.init_source_pattern(self.source_pattern)
+        drr_ctx.init_constraint_func(self.constraint)
+        drr_ctx.init_result_pattern(self.result_pattern)
+        return drr_ctx
+
+    def constraint(self, o, t):
+        return True
+
+    def drr_pass_type(self):
+        return "access_topo_drr_pass_type"
+
+
+class register_drr_pass:
+    def __init__(self, pass_name, tag):
+        self.pass_name = pass_name
+        self.tag = tag
+
+    def __call__(self, drr_pass_cls):
+        Registry.access_topo_drr_pass(  # noqa: F821
+            self.pass_name, self.tag, drr_pass_cls
+        )
+        return drr_pass_cls

python/paddle/apy/matmul_pass/code_gen_value_util.py

@@ -0,0 +1,29 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+class CodeGenValue:
+    def __init__(self, pir_type, var_name):
+        self.pir_type = pir_type
+        self.var_name = var_name
+        self.const_value = None
+
+    def get_dtype(self):
+        def convert_to_dtype(pir_dtype, shape, data_layout):
+            return pir_dtype.convert_to_dtype()
+
+        return self.pir_type.match(t_dtensor=convert_to_dtype)
+
+    def is_dense_tensor_type(self):
+        return self.pir_type.get_type_name() == "t_dtensor"

python/paddle/apy/matmul_pass/index_code_gen_value_util.py

@@ -0,0 +1,19 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+class IndexCodeGenValue:
+    def __init__(self, iter_var_names):
+        self.iter_var_names = iter_var_names
+        self.const_data = None

python/paddle/apy/matmul_pass/index_drr_pass_util.py

@@ -0,0 +1,33 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import access_topo_drr
+import ap
+import pir
+
+
+class InsertReshapeBeforeYieldPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.yield_op = o.ap_native_op("cf.yield")
+        o.yield_op([t.output], [])
+
+    def result_pattern(self, o, t):
+        t.declare_internal_native_ir_value("reshaped_output")
+        o.reshape_op = o.ap_native_op("cinn_op.reshape")
+        o.reshape_op.shape = lambda o, t: pir.a_array(
+            [pir.a_i32(ap.DataValue.int32("-1"))]
+        )
+        o.reshape_op([t.output], [t.reshaped_output])
+        o.yield_op = o.ap_native_op("cf.yield")
+        o.yield_op([t.reshaped_output], [])

python/paddle/apy/matmul_pass/index_program_translator_util.py

@@ -0,0 +1,137 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import ap
+import index_drr_pass_util
+import ir_tools
+import op_index_translator_util
+
+
+class IndexProgramTranslatorMap:
+    def __init__(
+        self,
+        index_func_unique_id2index_program,
+        kernel_arg_translator,
+        anchor_iter_var_names,
+    ):
+        self.kernel_arg_translator = kernel_arg_translator
+        self.anchor_iter_var_names = anchor_iter_var_names
+        items = index_func_unique_id2index_program.items()
+        self.index_func_unique_id2translator = ap.OrderedDict(
+            ap.map(
+                lambda i: [items[i][0], self.make_translator(i, items[i][1])],
+                range(len(items)),
+            )
+        )
+
+    def get_offset_var_name(
+        self,
+        index_func_unique_id,
+        mut_kernel_arg_id_registry,
+        mut_lir_code_gen_ctx,
+    ):
+        translator = self.index_func_unique_id2translator[index_func_unique_id]
+        ret = translator.translate(
+            mut_kernel_arg_id_registry=mut_kernel_arg_id_registry,
+            mut_lir_code_gen_ctx=mut_lir_code_gen_ctx,
+        )
+        return ret.iter_var_names[0]
+
+    def make_translator(self, program_id, index_program):
+        pass_manager = ir_tools.create_pass_manager()
+        drr_pass = index_drr_pass_util.InsertReshapeBeforeYieldPass()
+        pass_manager.add_pass(
+            ir_tools.create_access_topo_drr_one_step_pass(drr_pass)
+        )
+        pass_manager.add_pass(ir_tools.create_dce_pass())
+        pass_manager.run(index_program)
+        return IndexProgramTranslator(
+            index_program,
+            program_id=program_id,
+            kernel_arg_translator=self.kernel_arg_translator,
+            index_op_translator_maker=op_index_translator_util.OpIndexTranslatorFactory(),
+            anchor_iter_var_names=self.anchor_iter_var_names,
+        )
+
+
+class IndexProgramTranslator:
+    def __init__(
+        self,
+        index_program,
+        program_id,
+        kernel_arg_translator,
+        index_op_translator_maker,
+        anchor_iter_var_names,
+    ):
+        self.program_id = program_id
+        self.program_property = index_program.copy_to_const_program_data()
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_op_translator_maker = index_op_translator_maker
+        self.anchor_iter_var_names = anchor_iter_var_names
+        self.ir_value_index2translated_value = ap.MutableList()
+
+        def PushNone(x):
+            self.ir_value_index2translated_value.append(None)
+
+        ap.map(PushNone, self.program_property.values)
+
+    def translate(
+        self,
+        mut_kernel_arg_id_registry,
+        mut_lir_code_gen_ctx,
+    ):
+        def TranslateOp(op_property):
+            self._translate_op(
+                op_property, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+            )
+
+        ap.map(TranslateOp, self.program_property.ops)
+        return self.ir_value_index2translated_value[-1]
+
+    def _translate_op(
+        self, op_property, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        index_op_translator = self.index_op_translator_maker(
+            index_program_id=self.program_id,
+            op_property=op_property,
+            input_properties=ap.map(
+                self._get_value_property, op_property.input_value_indexes
+            ),
+            output_properties=ap.map(
+                self._get_value_property, op_property.output_value_indexes
+            ),
+            kernel_arg_translator=self.kernel_arg_translator,
+            anchor_iter_var_names=self.anchor_iter_var_names,
+        )
+        inputs = ap.map(
+            self._get_translated_value, op_property.input_value_indexes
+        )
+        outputs = index_op_translator(
+            inputs,
+            mut_kernel_arg_id_registry=mut_kernel_arg_id_registry,
+            mut_lir_code_gen_ctx=mut_lir_code_gen_ctx,
+        )
+        ap.map(
+            self._set_translated_value,
+            zip(op_property.output_value_indexes, outputs),
+        )
+
+    def _get_value_property(self, i):
+        return self.program_property.values[i]
+
+    def _get_translated_value(self, i):
+        return self.ir_value_index2translated_value[i]
+
+    def _set_translated_value(self, pair):
+        self.ir_value_index2translated_value[pair[0]] = pair[1]

python/paddle/apy/matmul_pass/kernel_arg_id_util.py

@@ -0,0 +1,88 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import ap
+
+
+class KernelArgIdNameRegistry:
+    def __init__(self, code_gen_ctx, tensor_match_ctx, name_prefix):
+        self.code_gen_ctx = code_gen_ctx
+        self.tensor_match_ctx = tensor_match_ctx
+        self.name_prefix = name_prefix
+        self.generated_kernel_arg_id2unique_name = ap.MutableOrderedDict()
+        self.all_kernel_arg_id2unique_name = ap.MutableOrderedDict()
+        self.in_tensor_data_ptr_seq_no = 0
+        self.out_tensor_data_ptr_seq_no = 0
+        self.dim_expr_seq_no = 0
+
+    def get_or_create_kernel_arg_id_manul_var_name(
+        self, kernel_arg_id, cpp_var_name
+    ):
+        create = lambda: 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):
+        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
+        )
+        return self.generated_kernel_arg_id2unique_name.get_or_create(
+            kernel_arg_id, create
+        )
+
+    def _get_creator(self, kernel_arg_id, backend_creator):
+        return lambda: self.all_kernel_arg_id2unique_name.get_or_create(
+            kernel_arg_id, backend_creator
+        )
+
+    def _create_in_tensor_data_ptr_var_name(self):
+        name = f"{self.name_prefix}in_ptr_{self.in_tensor_data_ptr_seq_no}"
+        self.in_tensor_data_ptr_seq_no = self.in_tensor_data_ptr_seq_no + 1
+        return name
+
+    def get_out_tensor_data_ptr_var_name(self, out_ir_value_name):
+        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
+        )
+        return self.generated_kernel_arg_id2unique_name.get_or_create(
+            kernel_arg_id, create
+        )
+
+    def _create_out_tensor_data_ptr_var_name(self):
+        name = f"{self.name_prefix}out_ptr_{self.out_tensor_data_ptr_seq_no}"
+        self.out_tensor_data_ptr_seq_no = self.out_tensor_data_ptr_seq_no + 1
+        return name
+
+    def get_dim_expr_var_name(self, dim_expr):
+        kernel_arg_id = self.code_gen_ctx.dim_expr_kernel_arg_id(dim_expr)
+        create = self._get_creator(
+            kernel_arg_id, self._create_dim_expr_var_name
+        )
+        return self.generated_kernel_arg_id2unique_name.get_or_create(
+            kernel_arg_id, create
+        )
+
+    def _create_dim_expr_var_name(self):
+        name = f"{self.name_prefix}dim_{self.dim_expr_seq_no}"
+        self.dim_expr_seq_no = self.dim_expr_seq_no + 1
+        return name

python/paddle/apy/matmul_pass/kernel_arg_translator_util.py

@@ -0,0 +1,30 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+class KernelArgTranslator:
+    def __init__(self, param_struct_name):
+        self.param_struct_name = param_struct_name
+
+    def get_kernel_arg_name(self, var_name):
+        return var_name
+
+    def get_param_struct_field_name(self, var_name):
+        return var_name
+
+    def get_param_struct_init_name(self, var_name):
+        return f"{self.param_struct_name}.{var_name}"
+
+    def get_use_name(self, var_name):
+        return f"{self.param_struct_name}.{var_name}"

python/paddle/apy/matmul_pass/low_level_ir_code_gen_ctx_util.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import ap
+
+
+class CudaLikeIrCodeGenCtx:
+    def __init__(self, compute_dtype):
+        self.stmts = ap.MutableList()
+        self.dtype2type_name = ap.OrderedDict(
+            [
+                [ap.DataType.float, "float"],
+                [ap.DataType.float16, "half"],
+                [ap.DataType.bfloat16, "nv_bfloat16"],
+                [ap.DataType.int32, "int"],
+                [ap.DataType.int64, "int64_t"],
+            ]
+        )
+        self.compute_dtype = compute_dtype
+        self.compute_dtype_name = self.dtype2type_name[self.compute_dtype]
+        self.type_cast_str_list = [
+            "",
+            f"static_cast<{self.compute_dtype_name}>",
+        ]
+
+    def assign(self, dst, src):
+        self.stmts.append(f"{dst.var_name} = {src.var_name};")
+
+    def let(self, var, val_name):
+        var_dtype_name = self.dtype2type_name[var.get_dtype()]
+        is_same = self.compute_dtype == var.get_dtype()
+        type_name = (
+            f"{var_dtype_name}" if is_same else f"{self.compute_dtype_name}"
+        )
+        type_cast_str = (
+            "" if is_same else f"static_cast<{self.compute_dtype_name}>"
+        )
+        self.stmts.append(
+            f"{type_name} {var.var_name} = {type_cast_str}({val_name});"
+        )
+
+    def store(self, dtype, dst, offset_var_name, src):
+        is_same = dtype == self.dtype2type_name[self.compute_dtype]
+        type_cast_str = "" if is_same else f"static_cast<{dtype}>"
+        self.stmts.append(f"{dst}[{offset_var_name}] = {type_cast_str}({src});")
+
+    def get_stmts_joined_str(self, indent):
+        return f"\n{indent}".join([*self.stmts])

python/paddle/apy/matmul_pass/matmul/.apy_ignore

@@ -0,0 +1,3 @@
+cutlass
+generate_configs.py
+hytlass

python/paddle/apy/matmul_pass/matmul/cutlass_patch/backend.h

@@ -0,0 +1,24 @@
+// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#ifdef __NVCC__
+#include "cutlass/cutlass.h"
+#define GPUStream_t cudaStream_t
+#elif defined(__HIPCC__)
+#include "hytlass/hytlass.h"
+namespace cutlass = hytlass;
+#define GPUStream_t hipStream_t
+#endif

python/paddle/apy/matmul_pass/matmul/cutlass_patch/batched_matrix_coord.h

@@ -0,0 +1,33 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "cutlass_patch/backend.h"
+
+namespace cutlass_patch {
+
+struct BatchedMatrixCoord {
+  int batch;
+  int row;
+  int column;
+
+  CUTLASS_HOST_DEVICE
+  BatchedMatrixCoord() : batch(0), row(0), column(0) {}
+
+  CUTLASS_HOST_DEVICE
+  BatchedMatrixCoord(int b, int r, int c) : batch(b), row(r), column(c) {}
+};
+
+};  // namespace cutlass_patch

python/paddle/apy/matmul_pass/matmul/cutlass_patch/cuda/all_tuning_configs.h

@@ -0,0 +1,496 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// auto-generated by generate_configs.py
+
+#pragma once
+
+#include "cutlass/gemm_coord.h"
+
+namespace ap {
+
+constexpr int kNumConfigsHalf = 23;
+constexpr int kNumConfigsFloat = 13;
+
+template <int SwizzleFactor, bool Batched>
+struct SwizzleWrapper {
+  using Type =
+      cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<SwizzleFactor>;
+};
+
+// template <int SwizzleFactor>
+// struct SwizzleWrapper<SwizzleFactor, true> {
+//   using Type =
+//       cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle;
+// };
+
+#define AP_AUTOTUNE(func, stream_ptr, count, ...)                             \
+  {                                                                           \
+    using FuncType = decltype(func<0>);                                       \
+    static int selected_config_id = -1;                                       \
+    static std::vector<std::function<FuncType>> matmul_functions =            \
+        []<std::size_t... Is>(std::index_sequence<Is...>) {                   \
+      return std::vector<std::function<FuncType>>{func<Is>...};               \
+    }                                                                         \
+    (std::make_index_sequence<count>());                                      \
+                                                                              \
+    if (selected_config_id == -1) {                                           \
+      selected_config_id =                                                    \
+          ap::ProfileBestConfig(matmul_functions, stream_ptr, ##__VA_ARGS__); \
+    }                                                                         \
+                                                                              \
+    matmul_functions[selected_config_id](__VA_ARGS__);                        \
+  }
+
+#define AP_AUTOTUNE_half(func, stream_ptr, ...) \
+  AP_AUTOTUNE(func, stream_ptr, ap::kNumConfigsHalf, __VA_ARGS__)
+#define AP_AUTOTUNE_float(func, stream_ptr, ...) \
+  AP_AUTOTUNE(func, stream_ptr, ap::kNumConfigsFloat, __VA_ARGS__)
+#define AP_AUTOTUNE_bfloat16(func, stream_ptr, ...) \
+  AP_AUTOTUNE_half(func, stream_ptr, __VA_ARGS__)
+
+template <typename ElementT, int SwizzleFactor, bool Batched, int Id = 0>
+struct GemmTuningConfigs {
+  using TShape = cutlass::gemm::GemmShape<256, 128, 64>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 2;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = Id;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 1> {
+  using TShape = cutlass::gemm::GemmShape<64, 128, 64>;
+  using WShape = cutlass::gemm::GemmShape<32, 64, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 1;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 2> {
+  using TShape = cutlass::gemm::GemmShape<64, 128, 64>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 2;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 3> {
+  using TShape = cutlass::gemm::GemmShape<128, 64, 64>;
+  using WShape = cutlass::gemm::GemmShape<64, 32, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 3;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 4> {
+  using TShape = cutlass::gemm::GemmShape<128, 128, 32>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 4;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 5> {
+  using TShape = cutlass::gemm::GemmShape<128, 128, 64>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 5;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 6> {
+  using TShape = cutlass::gemm::GemmShape<256, 64, 32>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 6;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 7> {
+  using TShape = cutlass::gemm::GemmShape<256, 64, 64>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 7;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 8> {
+  using TShape = cutlass::gemm::GemmShape<256, 128, 32>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 8;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 9> {
+  using TShape = cutlass::gemm::GemmShape<256, 128, 64>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 9;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 10> {
+  using TShape = cutlass::gemm::GemmShape<128, 32, 64>;
+  using WShape = cutlass::gemm::GemmShape<32, 32, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 10;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 11> {
+  using TShape = cutlass::gemm::GemmShape<128, 128, 32>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 11;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 12> {
+  using TShape = cutlass::gemm::GemmShape<128, 128, 64>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 12;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 13> {
+  using TShape = cutlass::gemm::GemmShape<256, 64, 64>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 13;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 14> {
+  using TShape = cutlass::gemm::GemmShape<256, 64, 32>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 14;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 15> {
+  using TShape = cutlass::gemm::GemmShape<32, 64, 64>;
+  using WShape = cutlass::gemm::GemmShape<16, 32, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 5;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 15;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 16> {
+  using TShape = cutlass::gemm::GemmShape<64, 64, 64>;
+  using WShape = cutlass::gemm::GemmShape<32, 32, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 5;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 16;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 17> {
+  using TShape = cutlass::gemm::GemmShape<128, 128, 32>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 5;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 17;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 18> {
+  using TShape = cutlass::gemm::GemmShape<128, 128, 64>;
+  using WShape = cutlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 5;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 18;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 19> {
+  using TShape = cutlass::gemm::GemmShape<64, 128, 32>;
+  using WShape = cutlass::gemm::GemmShape<32, 64, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 6;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 19;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 20> {
+  using TShape = cutlass::gemm::GemmShape<128, 64, 32>;
+  using WShape = cutlass::gemm::GemmShape<64, 32, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 6;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 20;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 21> {
+  using TShape = cutlass::gemm::GemmShape<128, 32, 32>;
+  using WShape = cutlass::gemm::GemmShape<32, 32, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 7;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 21;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 22> {
+  using TShape = cutlass::gemm::GemmShape<64, 64, 32>;
+  using WShape = cutlass::gemm::GemmShape<32, 32, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 16>;
+  static constexpr int kNumStages = 10;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 22;
+};
+
+// Specialization for float
+template <int SwizzleFactor, bool Batched, int Id>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, Id> {
+  using TShape = cutlass::gemm::GemmShape<64, 64, 16>;
+  using WShape = cutlass::gemm::GemmShape<32, 32, 16>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = Id;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 1> {
+  using TShape = cutlass::gemm::GemmShape<64, 64, 32>;
+  using WShape = cutlass::gemm::GemmShape<32, 32, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 1;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 2> {
+  using TShape = cutlass::gemm::GemmShape<64, 128, 32>;
+  using WShape = cutlass::gemm::GemmShape<32, 64, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 2;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 3> {
+  using TShape = cutlass::gemm::GemmShape<64, 256, 16>;
+  using WShape = cutlass::gemm::GemmShape<32, 64, 16>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 3;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 4> {
+  using TShape = cutlass::gemm::GemmShape<64, 256, 32>;
+  using WShape = cutlass::gemm::GemmShape<32, 64, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 4;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 5> {
+  using TShape = cutlass::gemm::GemmShape<128, 64, 32>;
+  using WShape = cutlass::gemm::GemmShape<64, 32, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 5;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 6> {
+  using TShape = cutlass::gemm::GemmShape<128, 128, 16>;
+  using WShape = cutlass::gemm::GemmShape<32, 64, 16>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 6;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 7> {
+  using TShape = cutlass::gemm::GemmShape<128, 128, 32>;
+  using WShape = cutlass::gemm::GemmShape<32, 64, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 7;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 8> {
+  using TShape = cutlass::gemm::GemmShape<256, 64, 16>;
+  using WShape = cutlass::gemm::GemmShape<64, 32, 16>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 8;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 9> {
+  using TShape = cutlass::gemm::GemmShape<256, 64, 32>;
+  using WShape = cutlass::gemm::GemmShape<64, 32, 32>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 9;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 10> {
+  using TShape = cutlass::gemm::GemmShape<64, 128, 16>;
+  using WShape = cutlass::gemm::GemmShape<32, 64, 16>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 10;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 11> {
+  using TShape = cutlass::gemm::GemmShape<128, 64, 16>;
+  using WShape = cutlass::gemm::GemmShape<64, 32, 16>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 11;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 12> {
+  using TShape = cutlass::gemm::GemmShape<128, 128, 16>;
+  using WShape = cutlass::gemm::GemmShape<32, 64, 16>;
+  using IShape = cutlass::gemm::GemmShape<16, 8, 8>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 12;
+};
+
+}  // namespace ap

python/paddle/apy/matmul_pass/matmul/cutlass_patch/cuda/cutlass_matmul.cuh

@@ -0,0 +1,273 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+
+#include "cutlass/cutlass.h"
+#include "cutlass/gemm_coord.h"
+#include "cutlass/layout/matrix.h"
+
+#include "cutlass/epilogue/thread/linear_combination_bias_elementwise.h"
+#include "cutlass/util/device_memory.h"
+
+#include "cutlass/gemm/device/gemm_universal.h"
+#include "cutlass/gemm/device/gemm_universal_with_broadcast.h"
+
+#include "cutlass_patch/batched_matrix_coord.h"
+#include "cutlass_patch/cuda/default_config_id.h"
+#include "cutlass_patch/epilogue/thread/linear_combination_unary.h"
+#include "cutlass_patch/epilogue/thread/linear_combination_variadic.h"
+#include "cutlass_patch/gemm/device/gemm_universal_with_variadic.h"
+
+#include "params.h"  // NOLINT
+
+#define CHECK_CUTLASS(status)                                             \
+  {                                                                       \
+    cutlass::Status error = status;                                       \
+    if (error != cutlass::Status::kSuccess) {                             \
+      std::cerr << "Got cutlass error: " << cutlassGetStatusString(error) \
+                << " at: " << __LINE__ << std::endl;                      \
+      exit(EXIT_FAILURE);                                                 \
+    }                                                                     \
+  }
+
+namespace ap {
+using bfloat16 = nv_bfloat16;
+
+template <typename T, int N>
+using Array = cutlass::Array<T, N>;
+
+using MatrixCoord = cutlass_patch::BatchedMatrixCoord;
+
+// Convert CUDA data type to cutlass data type
+template <typename T>
+struct CutlassDataType {
+  using Type = T;
+};
+
+template <>
+struct CutlassDataType<half> {
+  using Type = cutlass::half_t;
+};
+
+template <>
+struct CutlassDataType<__nv_bfloat16> {
+  using Type = cutlass::bfloat16_t;
+};
+
+// Convert to cutlass layout
+template <bool Transposed>
+struct MatrixLayout {
+  using Type = cutlass::layout::RowMajor;
+};
+
+template <>
+struct MatrixLayout<true> {
+  using Type = cutlass::layout::ColumnMajor;
+};
+
+// Operation performed by GEMM
+template <typename ElementT>
+struct GemmOperation {
+  using Type = cutlass::arch::OpMultiplyAdd;
+};
+
+template <>
+struct GemmOperation<float> {
+  using Type = cutlass::arch::OpMultiplyAddFastF32;
+};
+
+static cutlass::gemm::GemmUniversalMode GetGemmMode(int batch_count) {
+  return batch_count > 1 ? cutlass::gemm::GemmUniversalMode::kBatched
+                         : cutlass::gemm::GemmUniversalMode::kGemm;
+}
+
+static void *GetWorkspace(size_t workspace_size) {
+  static cutlass::device_memory::allocation<uint8_t> workspace;
+  if (workspace.size() < workspace_size) {
+    workspace.reset(workspace_size);
+  }
+  return workspace.get();
+}
+
+template <typename GemmFunc>
+cutlass::Status SetMaxDynamicSharedMemorySize() {
+  cudaError_t cudart_result;
+
+  // If requires more than 48KB: configure for extended, dynamic shared memory
+  if constexpr (GemmFunc::kSharedStorageSize >= (48 << 10)) {
+    cudart_result =
+        cudaFuncSetAttribute(cutlass::Kernel2<typename GemmFunc::GemmKernel>,
+                             cudaFuncAttributeMaxDynamicSharedMemorySize,
+                             GemmFunc::kSharedStorageSize);
+    if (cudart_result != cudaSuccess) {
+      CUTLASS_TRACE_HOST("cudaFuncSetAttribute() returned error "
+                         << cudaGetErrorString(cudart_result));
+      return cutlass::Status::kErrorInternal;
+    }
+  }
+
+#if AP_ENABLE_DEBUG
+  // Update SM occupancy member
+  int sm_occupancy = -1;
+  cudart_result = cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(
+      &sm_occupancy,
+      cutlass::Kernel2<typename GemmFunc::GemmKernel>,
+      GemmFunc::GemmKernel::kThreadCount,
+      GemmFunc::kSharedStorageSize,
+      cudaOccupancyDisableCachingOverride);
+  if (cudart_result != cudaSuccess) {
+    CUTLASS_TRACE_HOST(
+        "cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags() returned "
+        "error "
+        << cudaGetErrorString(cudart_result));
+    return cutlass::Status::kErrorInternal;
+  }
+  CUTLASS_TRACE_HOST("sm_occupancy: (" << sm_occupancy
+                                       << ") "
+                                          "smem_size: ("
+                                       << GemmFunc::kSharedStorageSize
+                                       << ") "
+                                          "GemmKernel::kThreadCount: ("
+                                       << GemmFunc::GemmKernel::kThreadCount
+                                       << ")");
+#endif
+  return cutlass::Status::kSuccess;
+}
+
+template <typename ElementT,
+          typename ElementComputeT,
+          template <typename T>
+          class VariadicFunctor,
+          int AlignA = 128 / cutlass::sizeof_bits<ElementT>::value,
+          int AlignB = 128 / cutlass::sizeof_bits<ElementT>::value,
+          int ConfigId = DefaultConfig::kConfigId,
+          int SwizzleFactor = DefaultConfig::kSwizzleFactor,
+          bool Batched = DefaultConfig::kBatched>
+void MatmulAddVariadic(
+    const GemmEpilogueParams &params,
+    const typename VariadicFunctor<ElementComputeT>::Arguments &variadic_args) {
+  using ElementAccumulator =
+      typename CutlassDataType<ElementComputeT>::Type;  // <- data type of
+                                                        // accumulator
+  using ElementComputeEpilogue =
+      ElementAccumulator;  // <- data type of epilogue operations
+  using ElementInputA =
+      typename CutlassDataType<ElementT>::Type;  // <- data type of elements in
+                                                 // input matrix A
+  using ElementInputB =
+      typename CutlassDataType<ElementT>::Type;  // <- data type of elements in
+                                                 // input matrix B
+  using ElementOutput =
+      typename CutlassDataType<ElementT>::Type;  // <- data type of elements in
+                                                 // output matrix D
+
+  constexpr int AlignC = AlignB;
+
+  // Epilogue operation as LinearCombination:
+  //  alpha * accumulator + beta * source
+  using EpilogueOutputOp =
+      cutlass_patch::epilogue::thread::LinearCombinationVariadic<
+          VariadicFunctor,
+          ElementOutput,
+          AlignC,
+          ElementAccumulator,
+          ElementComputeEpilogue,
+          cutlass::epilogue::thread::ScaleType::NoBetaScaling>;  // <- alpha x
+                                                                 // AB + bias
+
+  using GemmFunc = cutlass_patch::gemm::device::GemmUniversalWithVariadic<
+      ElementInputA,
+      cutlass::layout::RowMajor,
+      ElementInputB,
+      cutlass::layout::RowMajor,
+      ElementOutput,
+      cutlass::layout::RowMajor,
+      ElementAccumulator,
+      cutlass::arch::OpClassTensorOp,
+      cutlass::arch::Sm80,
+      typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
+          TShape,
+      typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
+          WShape,
+      typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
+          IShape,
+      EpilogueOutputOp,
+      typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
+          SwizzleThreadBlock,
+      GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::kNumStages,
+      AlignA,
+      AlignB,
+      typename GemmOperation<ElementT>::Type>;
+
+  CHECK_CUTLASS(SetMaxDynamicSharedMemorySize<GemmFunc>());
+
+  /// Arguments
+  cutlass::gemm::GemmCoord problem_size{params.m, params.n, params.k};
+
+  const ElementInputA *input =
+      reinterpret_cast<const ElementInputA *>(params.input);
+  const ElementInputB *weight =
+      reinterpret_cast<const ElementInputB *>(params.weight);
+  const ElementOutput *bias =
+      reinterpret_cast<const ElementOutput *>(params.bias);
+  ElementOutput *output = reinterpret_cast<ElementOutput *>(params.output);
+
+  ElementComputeEpilogue alpha = static_cast<ElementComputeEpilogue>(1);
+  ElementComputeEpilogue beta = bias ? static_cast<ElementComputeEpilogue>(1)
+                                     : static_cast<ElementComputeEpilogue>(0);
+
+  typename GemmFunc::Arguments arguments{
+      GetGemmMode(params.batch_count),
+      problem_size,        // <- problem size of matrix multiplication
+      params.batch_count,  // <- batch_count or k-dimension split factor
+      {alpha, beta, variadic_args},  // <- epilogue params, alpha, beta
+      input,                         // <- input, ptr_A, A, shape={M, K}
+      weight,                        // <- input, ptr_B, B, shape={K, N}
+      bias,                          // <- input, ptr_C, shape={M, N} or {1, N}
+      output,                        // <- output, ptr_D, Z, shape={M, N}
+      params.shape_args.batch_stride_A,
+      params.shape_args.batch_stride_B,
+      params.shape_args.batch_stride_C,
+      params.shape_args.batch_stride_D,
+      params.shape_args.lda,
+      params.shape_args.ldb,
+      params.shape_args.ldc_bias,
+      params.shape_args.ldd};
+
+  size_t workspace_size = GemmFunc::get_workspace_size(arguments);
+  void *workspace = workspace_size > 0 ? GetWorkspace(workspace_size) : nullptr;
+
+  GemmFunc device_gemm;
+
+  cudaStream_t *stream_ptr =
+      reinterpret_cast<cudaStream_t *>(params.stream_ptr);
+
+  CHECK_CUTLASS(device_gemm.can_implement(arguments));
+  CHECK_CUTLASS(device_gemm.initialize(arguments, workspace, *stream_ptr));
+
+  //
+  // Run the GEMM
+  //
+  CHECK_CUTLASS(device_gemm(*stream_ptr));
+#if AP_ENABLE_DEBUG
+  CHECK_CUDA(cudaStreamSynchronize(*stream_ptr));
+#endif
+}
+
+}  // namespace ap

python/paddle/apy/matmul_pass/matmul/cutlass_patch/cuda/default_config_id.h

@@ -0,0 +1,27 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "all_tuning_configs.h"  // NOLINT
+
+namespace ap {
+
+struct DefaultConfig {
+  static constexpr int kConfigId = 0;
+  static constexpr int kSwizzleFactor = 1;
+  static constexpr bool kBatched = false;
+};
+
+}  // namespace ap

python/paddle/apy/matmul_pass/matmul/cutlass_patch/epilogue/thread/linear_combination_unary.h

@@ -0,0 +1,302 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/*! \file
+  \brief Functor performing linear combination operations used by epilogues.
+  It is modified from LinearCombinationGeneric.
+*/
+
+#pragma once
+
+#include "cutlass_patch/backend.h"
+
+#ifdef __NVCC__
+#include "cutlass/array.h"
+#include "cutlass/epilogue/thread/scale_type.h"
+#include "cutlass/functional.h"
+#include "cutlass/numeric_conversion.h"
+#include "cutlass/numeric_types.h"
+#elif defined(__HIPCC__)
+#include "hytlass/array.h"
+#include "hytlass/epilogue/thread/scale_type.h"
+#include "hytlass/functional.h"
+#include "hytlass/numeric_conversion.h"
+#include "hytlass/numeric_types.h"
+#endif
+
+namespace cutlass_patch {
+namespace epilogue {
+namespace thread {
+
+template <class UnaryOp, class = void>
+struct GenericUnaryTraits {
+  static constexpr bool IsArgumentsNeeded = false;
+  struct Arguments {};
+};
+
+template <class UnaryOp>
+struct GenericUnaryTraits<UnaryOp,
+                          decltype(typename UnaryOp::Arguments(), void())> {
+  static constexpr bool IsArgumentsNeeded = true;
+  using Arguments = typename UnaryOp::Arguments;
+};
+
+/// Applies a linear combination operator followed by an unary function to an
+/// array of elements.
+///
+/// D = unary_op(alpha * accumulator + beta * source)
+///
+template <
+    template <typename T>
+    class UnaryOp,
+    typename ElementOutput_,  ///< Data type used to load and store tensors
+    int ElementsPerAccess,    ///< Number of elements computed per operation
+                            ///< Usually it is 128/sizeof_bits<ElementOutput_>,
+                            ///< but we use 64 or 32 sometimes when there are
+                            ///< not enough data to store
+    typename ElementAccumulator_ = ElementOutput_,  ///< Accumulator data type
+    typename ElementCompute_ =
+        ElementOutput_,  ///< Data type used to compute linear combination
+    cutlass::epilogue::thread::ScaleType::Kind Scale =
+        cutlass::epilogue::thread::ScaleType::Default,  ///< Control Alpha and
+                                                        ///< Beta scaling
+    cutlass::FloatRoundStyle Round = cutlass::FloatRoundStyle::round_to_nearest,
+    bool IsHeavy = false>
+class LinearCombinationUnary {
+ public:
+  using ElementOutput = ElementOutput_;
+  using ElementAccumulator = ElementAccumulator_;
+  using ElementCompute = ElementCompute_;
+  using UnaryArguments =
+      typename GenericUnaryTraits<UnaryOp<ElementCompute>>::Arguments;
+
+  static bool const kIsHeavy = IsHeavy;
+  static int const kElementsPerAccess = ElementsPerAccess;
+  static int const kCount = ElementsPerAccess;
+  static const cutlass::epilogue::thread::ScaleType::Kind kScale = Scale;
+
+  using FragmentOutput = cutlass::Array<ElementOutput, kElementsPerAccess>;
+  using FragmentAccumulator =
+      cutlass::Array<ElementAccumulator, kElementsPerAccess>;
+  using FragmentSource = cutlass::Array<ElementOutput, kElementsPerAccess>;
+  using FragmentCompute = cutlass::Array<ElementCompute, kElementsPerAccess>;
+
+  static cutlass::FloatRoundStyle const kRound = Round;
+
+  /// Host-constructable parameters structure
+  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
+    ElementCompute const *beta_ptr;   ///< pointer to source scalar - if not
+                                      ///< null, loads it from memory
+    UnaryArguments unary_args;
+
+    //
+    // Methods
+    //
+
+    CUTLASS_HOST_DEVICE
+    Params()
+        : alpha(ElementCompute(1)),
+          beta(ElementCompute(0)),
+          alpha_ptr(nullptr),
+          beta_ptr(nullptr) {}
+
+    CUTLASS_HOST_DEVICE
+    Params(ElementCompute alpha,
+           ElementCompute beta = ElementCompute(0),
+           UnaryArguments unary_args_ = UnaryArguments{})
+        : alpha(alpha),
+          beta(beta),
+          alpha_ptr(nullptr),
+          beta_ptr(nullptr),
+          unary_args(unary_args_) {}
+
+    CUTLASS_HOST_DEVICE
+    Params(ElementCompute const *alpha_ptr,
+           ElementCompute const *beta_ptr = nullptr)
+        : alpha(0), beta(0), alpha_ptr(alpha_ptr), beta_ptr(beta_ptr) {}
+  };
+
+ private:
+  //
+  // Data members
+  //
+
+  Params params_;
+  bool skip_elementwise_;
+
+ public:
+  /// Constructs the function object, possibly loading from pointers in host
+  /// memory
+  CUTLASS_HOST_DEVICE
+  explicit LinearCombinationUnary(Params const &params) {
+    params_ = params;
+    params_.alpha = (params.alpha_ptr ? *params.alpha_ptr : params.alpha);
+    params_.beta = (params.beta_ptr ? *params.beta_ptr : params.beta);
+    skip_elementwise_ = false;
+  }
+
+  /// Returns true if source is needed
+  CUTLASS_HOST_DEVICE
+  bool is_source_needed() const {
+    if (Scale == cutlass::epilogue::thread::ScaleType::NoBetaScaling)
+      return params_.beta != ElementCompute(0);
+
+    if (Scale == cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling)
+      return false;
+
+    if (Scale == cutlass::epilogue::thread::ScaleType::Nothing) return false;
+
+    return params_.beta != ElementCompute(0);
+  }
+
+  /// Functionally required for serial reduction in the epilogue
+  CUTLASS_HOST_DEVICE
+  void set_k_partition(int k_partition, int k_partition_count) {
+    if (k_partition) {
+      params_.beta = ElementCompute(1);
+    }
+
+    if (k_partition != k_partition_count - 1) {
+      skip_elementwise_ = true;
+    }
+  }
+
+  /// Computes linear scaling: D = alpha * accumulator + beta * source
+  CUTLASS_HOST_DEVICE
+  FragmentOutput operator()(FragmentAccumulator const &accumulator,
+                            FragmentOutput const &source) const {
+    // Convert source to internal compute numeric type
+    cutlass::NumericArrayConverter<ElementCompute,
+                                   ElementOutput,
+                                   kElementsPerAccess,
+                                   Round>
+        source_converter;
+    cutlass::NumericArrayConverter<ElementCompute,
+                                   ElementAccumulator,
+                                   kElementsPerAccess,
+                                   Round>
+        accumulator_converter;
+
+    FragmentCompute converted_source = source_converter(source);
+    FragmentCompute converted_accumulator = accumulator_converter(accumulator);
+
+    // Perform binary operations
+    FragmentCompute intermediate;
+
+    cutlass::multiplies<FragmentCompute> mul_add_source;
+    cutlass::multiply_add<FragmentCompute> mul_add_accumulator;
+    UnaryOp<ElementCompute> unary_op;
+
+    if (Scale == cutlass::epilogue::thread::ScaleType::NoBetaScaling) {
+      intermediate = converted_source;
+      // D = alpha * Accum + X
+      intermediate = mul_add_accumulator(
+          params_.alpha, converted_accumulator, intermediate);
+    } else if (Scale == cutlass::epilogue::thread::ScaleType::Nothing) {
+      intermediate = converted_accumulator;
+    } else {
+      // X =  beta * C + uniform
+      intermediate = mul_add_source(params_.beta, converted_source);
+      // D = alpha * Accum + X
+      intermediate = mul_add_accumulator(
+          params_.alpha, converted_accumulator, intermediate);
+    }
+
+    if constexpr (GenericUnaryTraits<
+                      UnaryOp<ElementCompute>>::IsArgumentsNeeded) {
+      if (!skip_elementwise_) {
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < kElementsPerAccess; ++i) {
+          intermediate[i] = unary_op(intermediate[i], params_.unary_args);
+        }
+      }
+    } else {
+      if (!skip_elementwise_) {
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < kElementsPerAccess; ++i) {
+          intermediate[i] = unary_op(intermediate[i]);
+        }
+      }
+    }
+
+    // Convert to destination numeric type
+    cutlass::NumericArrayConverter<ElementOutput,
+                                   ElementCompute,
+                                   kElementsPerAccess,
+                                   Round>
+        destination_converter;
+
+    return destination_converter(intermediate);
+  }
+
+  /// Computes linear scaling: D = alpha * accumulator
+  CUTLASS_HOST_DEVICE
+  FragmentOutput operator()(FragmentAccumulator const &accumulator) const {
+    // Convert source to internal compute numeric type
+    cutlass::NumericArrayConverter<ElementCompute,
+                                   ElementAccumulator,
+                                   kElementsPerAccess,
+                                   Round>
+        accumulator_converter;
+
+    FragmentCompute converted_accumulator = accumulator_converter(accumulator);
+
+    // Perform binary operations
+    FragmentCompute intermediate;
+
+    cutlass::multiplies<FragmentCompute> mul_add_accumulator;
+    UnaryOp<ElementCompute> unary_op;
+
+    if (Scale == cutlass::epilogue::thread::ScaleType::Nothing) {
+      intermediate = converted_accumulator;
+    } else {
+      // D = alpha * Accum
+      intermediate = mul_add_accumulator(params_.alpha, converted_accumulator);
+    }
+
+    if constexpr (GenericUnaryTraits<
+                      UnaryOp<FragmentCompute>>::IsArgumentsNeeded) {
+      if (!skip_elementwise_) {
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < kElementsPerAccess; ++i) {
+          intermediate[i] = unary_op(intermediate[i], params_.unary_args);
+        }
+      }
+    } else {
+      if (!skip_elementwise_) {
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < kElementsPerAccess; ++i) {
+          intermediate[i] = unary_op(intermediate[i]);
+        }
+      }
+    }
+
+    // Convert to destination numeric type
+    cutlass::NumericArrayConverter<ElementOutput,
+                                   ElementCompute,
+                                   kElementsPerAccess,
+                                   Round>
+        destination_converter;
+
+    return destination_converter(intermediate);
+  }
+};
+
+}  // namespace thread
+}  // namespace epilogue
+}  // namespace cutlass_patch

python/paddle/apy/matmul_pass/matmul/cutlass_patch/epilogue/thread/linear_combination_variadic.h

@@ -0,0 +1,337 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/*! \file
+  \brief Functor performing linear combination operations used by epilogues.
+*/
+
+#pragma once
+
+#include "cutlass_patch/backend.h"
+
+#ifdef __NVCC__
+#include "cutlass/array.h"
+#include "cutlass/epilogue/thread/scale_type.h"
+#include "cutlass/functional.h"
+#include "cutlass/numeric_conversion.h"
+#include "cutlass/numeric_types.h"
+#elif defined(__HIPCC__)
+#include "hytlass/array.h"
+#include "hytlass/epilogue/thread/scale_type.h"
+#include "hytlass/functional.h"
+#include "hytlass/numeric_conversion.h"
+#include "hytlass/numeric_types.h"
+#endif
+
+#include "cutlass_patch/batched_matrix_coord.h"
+#include "cutlass_patch/trace_device.h"
+
+namespace cutlass_patch {
+namespace epilogue {
+namespace thread {
+
+template <class VariadicOp, class = void>
+struct GenericVariadicTraits {
+  static constexpr bool IsArgumentsNeeded = false;
+  struct Arguments {};
+};
+
+template <class VariadicOp>
+struct GenericVariadicTraits<VariadicOp,
+                             decltype(typename VariadicOp::Arguments(),
+                                      void())> {
+  static constexpr bool IsArgumentsNeeded = true;
+  using Arguments = typename VariadicOp::Arguments;
+};
+
+/// Applies a linear combination operator to an array of elements.
+///
+/// D = VariadicOp(alpha * accumulator + beta * source)
+///
+template <
+    template <typename T>
+    class VariadicOp,
+    typename ElementOutput_,  ///< Data type used to load and store tensors
+    int ElementsPerAccess,    ///< Number of elements computed per operation.
+                            ///< Usually it is 128/sizeof_bits<ElementOutput_>,
+                            ///< but we use 64 or 32 sometimes when there are
+                            ///< not enough data to store
+    typename ElementAccumulator_ = ElementOutput_,  ///< Accumulator data type
+    typename ElementCompute_ =
+        ElementOutput_,  ///< Data type used to compute linear combination
+    cutlass::epilogue::thread::ScaleType::Kind Scale =
+        cutlass::epilogue::thread::ScaleType::Default,  ///< Control Alpha and
+                                                        ///< Beta scaling
+    cutlass::FloatRoundStyle Round = cutlass::FloatRoundStyle::round_to_nearest,
+    bool IsHeavy = false>
+class LinearCombinationVariadic {
+ public:
+  using ElementOutput = ElementOutput_;
+  using ElementAccumulator = ElementAccumulator_;
+  using ElementCompute = ElementCompute_;
+  using VariadicArguments =
+      typename GenericVariadicTraits<VariadicOp<ElementCompute>>::Arguments;
+
+  static bool const kIsHeavy = IsHeavy;
+  static int const kElementsPerAccess = ElementsPerAccess;
+  static int const kCount = ElementsPerAccess;
+  static const cutlass::epilogue::thread::ScaleType::Kind kScale = Scale;
+
+  using FragmentOutput = cutlass::Array<ElementOutput, kElementsPerAccess>;
+  using FragmentAccumulator =
+      cutlass::Array<ElementAccumulator, kElementsPerAccess>;
+  using FragmentSource = cutlass::Array<ElementOutput, kElementsPerAccess>;
+  using FragmentCompute = cutlass::Array<ElementCompute, kElementsPerAccess>;
+
+  static cutlass::FloatRoundStyle const kRound = Round;
+
+  /// Host-constructable parameters structure
+  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
+    ElementCompute const *beta_ptr;   ///< pointer to source scalar - if not
+                                      ///< null, loads it from memory
+    VariadicArguments variadic_args;
+
+    CUTLASS_HOST_DEVICE
+    Params()
+        : alpha(ElementCompute(1)),
+          beta(ElementCompute(0)),
+          alpha_ptr(nullptr),
+          beta_ptr(nullptr) {}
+
+    CUTLASS_HOST_DEVICE
+    Params(ElementCompute alpha,
+           ElementCompute beta,
+           VariadicArguments variadic_args_ = VariadicArguments{})
+        : alpha(alpha),
+          beta(beta),
+          alpha_ptr(nullptr),
+          beta_ptr(nullptr),
+          variadic_args(variadic_args_) {}
+  };
+
+ private:
+  //
+  // Data members
+  //
+
+  Params params_;
+  bool skip_elementwise_;
+
+ public:
+  /// Constructs the function object, possibly loading from pointers in host
+  /// memory
+  CUTLASS_HOST_DEVICE
+  LinearCombinationVariadic(Params const &params) {
+    params_ = params;
+    params_.alpha = (params.alpha_ptr ? *params.alpha_ptr : params.alpha);
+    params_.beta = (params.beta_ptr ? *params.beta_ptr : params.beta);
+    skip_elementwise_ = false;
+  }
+
+  /// Returns true if source is needed
+  CUTLASS_HOST_DEVICE
+  bool is_source_needed() const {
+    if (Scale == cutlass::epilogue::thread::ScaleType::NoBetaScaling)
+      return params_.beta != ElementCompute(0);
+
+    if (Scale == cutlass::epilogue::thread::ScaleType::OnlyAlphaScaling)
+      return false;
+
+    if (Scale == cutlass::epilogue::thread::ScaleType::Nothing) return false;
+
+    return params_.beta != ElementCompute(0);
+  }
+
+  /// Functionally required for serial reduction in the epilogue
+  CUTLASS_HOST_DEVICE
+  void set_k_partition(int k_partition, int k_partition_count) {
+    if (k_partition) {
+      params_.beta = ElementCompute(1);
+    }
+
+    if (k_partition != k_partition_count - 1) {
+      skip_elementwise_ = true;
+    }
+  }
+
+  /// Computes linear scaling with source: D = alpha * accumulator + beta *
+  /// source
+  CUTLASS_HOST_DEVICE
+  FragmentOutput operator()(FragmentAccumulator const &accumulator,
+                            FragmentSource const &source,
+                            int row_offset,
+                            int column_offset) const {
+    CUTLASS_TRACE_DEVICE(
+        "kElementsPerAccess: %d, row_offset: %d, column_offset: %d",
+        kElementsPerAccess,
+        row_offset,
+        column_offset);
+
+    // Convert source to internal compute numeric type
+    cutlass::NumericArrayConverter<ElementCompute,
+                                   ElementOutput,
+                                   kElementsPerAccess,
+                                   Round>
+        source_converter;
+    cutlass::NumericArrayConverter<ElementCompute,
+                                   ElementAccumulator,
+                                   kElementsPerAccess,
+                                   Round>
+        accumulator_converter;
+
+    FragmentCompute converted_source = source_converter(source);
+    FragmentCompute converted_accumulator = accumulator_converter(accumulator);
+
+    // Perform binary operations
+    FragmentCompute intermediate;
+
+    cutlass::multiplies<FragmentCompute> mul_add_source;
+    cutlass::multiply_add<FragmentCompute> mul_add_accumulator;
+    VariadicOp<ElementCompute> variadic_op;
+
+    if (Scale == cutlass::epilogue::thread::ScaleType::NoBetaScaling) {
+      intermediate = converted_source;
+      // D = alpha * Accum + X
+      intermediate = mul_add_accumulator(
+          params_.alpha, converted_accumulator, intermediate);
+    } else if (Scale == cutlass::epilogue::thread::ScaleType::Nothing) {
+      intermediate = converted_accumulator;
+    } else {
+      // X =  beta * C + uniform
+      intermediate = mul_add_source(params_.beta, converted_source);
+      // D = alpha * Accum + X
+      intermediate = mul_add_accumulator(
+          params_.alpha, converted_accumulator, intermediate);
+    }
+
+    if constexpr (GenericVariadicTraits<
+                      VariadicOp<ElementCompute>>::IsArgumentsNeeded) {
+      if (!skip_elementwise_) {
+#if CUTLASS_EPILOGUE_ENABLE_VECTORIZE
+        intermediate = variadic_op.Compute<kElementsPerAccess>(
+            intermediate,
+            params_.variadic_args,
+            BatchedMatrixCoord(blockIdx.z, row_offset, column_offset));
+#else
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < kElementsPerAccess; ++i) {
+          intermediate[i] = variadic_op(
+              intermediate[i],
+              params_.variadic_args,
+              BatchedMatrixCoord(blockIdx.z, row_offset, column_offset + i));
+        }
+#endif
+      }
+    } else {
+      if (!skip_elementwise_) {
+#if CUTLASS_EPILOGUE_ENABLE_VECTORIZE
+        intermediate = variadic_op.Compute<kElementsPerAccess>(intermediate);
+#else
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < kElementsPerAccess; ++i) {
+          intermediate[i] = variadic_op(intermediate[i]);
+        }
+#endif
+      }
+    }
+
+    // Convert to destination numeric type
+    cutlass::NumericArrayConverter<ElementOutput,
+                                   ElementCompute,
+                                   kElementsPerAccess,
+                                   Round>
+        destination_converter;
+
+    return destination_converter(intermediate);
+  }
+
+  /// Computes linear scaling: D = alpha * accumulator
+  CUTLASS_HOST_DEVICE
+  FragmentOutput operator()(FragmentAccumulator const &accumulator,
+                            int row_offset,
+                            int column_offset) const {
+    CUTLASS_TRACE_DEVICE(
+        "kElementsPerAccess: %d, row_offset: %d, column_offset: %d",
+        kElementsPerAccess,
+        row_offset,
+        column_offset);
+
+    // Convert source to internal compute numeric type
+    cutlass::NumericArrayConverter<ElementCompute,
+                                   ElementAccumulator,
+                                   kElementsPerAccess,
+                                   Round>
+        accumulator_converter;
+
+    FragmentCompute converted_accumulator = accumulator_converter(accumulator);
+
+    // Perform binary operations
+    FragmentCompute intermediate;
+
+    cutlass::multiplies<FragmentCompute> mul_accumulator;
+    VariadicOp<ElementCompute> variadic_op;
+
+    if (Scale == cutlass::epilogue::thread::ScaleType::Nothing) {
+      intermediate = converted_accumulator;
+    } else {
+      // D = alpha * Accum
+      intermediate = mul_accumulator(params_.alpha, converted_accumulator);
+    }
+
+    if constexpr (GenericVariadicTraits<
+                      VariadicOp<FragmentCompute>>::IsArgumentsNeeded) {
+      if (!skip_elementwise_) {
+#if CUTLASS_EPILOGUE_ENABLE_VECTORIZE
+        intermediate = variadic_op.Compute<kElementsPerAccess>(
+            intermediate,
+            params_.variadic_args,
+            BatchedMatrixCoord(blockIdx.z, row_offset, column_offset));
+#else
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < kElementsPerAccess; ++i) {
+          intermediate[i] = variadic_op(
+              intermediate[i],
+              params_.variadic_args,
+              BatchedMatrixCoord(blockIdx.z, row_offset, column_offset + i));
+        }
+#endif
+      }
+    } else {
+      if (!skip_elementwise_) {
+#if CUTLASS_EPILOGUE_ENABLE_VECTORIZE
+        intermediate = variadic_op.Compute<kElementsPerAccess>(intermediate);
+#else
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < kElementsPerAccess; ++i) {
+          intermediate[i] = variadic_op(intermediate[i]);
+        }
+#endif
+      }
+    }
+
+    // Convert to destination numeric type
+    cutlass::NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round>
+        destination_converter;
+
+    return destination_converter(intermediate);
+  }
+};
+
+}  // namespace thread
+}  // namespace epilogue
+}  // namespace cutlass_patch

python/paddle/apy/matmul_pass/matmul/cutlass_patch/epilogue/threadblock/default_epilogue_with_variadic.h

@@ -0,0 +1,243 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/*! \file
+  \brief Epilogue for threadblock scoped GEMMs using Tensor Ops.
+
+  The epilogue rearranges the result of a matrix product through shared memory
+  to match canonical tensor layouts in global memory. Epilogues support
+  conversion and reduction operations.
+
+*/
+
+#pragma once
+
+#include "cutlass_patch/backend.h"
+
+#ifdef __NVCC__
+#include "cutlass/array.h"
+#include "cutlass/numeric_types.h"
+
+#include "cutlass/gemm/gemm.h"
+
+#include "cutlass/epilogue/threadblock/default_epilogue_tensor_op.h"
+#include "cutlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h"
+#include "cutlass/epilogue/threadblock/epilogue.h"
+
+#include "cutlass/layout/permute.h"
+#elif defined(__HIPCC__)
+#include "hytlass/array.h"
+#include "hytlass/numeric_types.h"
+
+#include "hytlass/gemm/gemm.h"
+
+#include "hytlass/epilogue/threadblock/default_epilogue_tensor_op.h"
+#include "hytlass/epilogue/threadblock/default_epilogue_volta_tensor_op.h"
+#include "hytlass/epilogue/threadblock/epilogue.h"
+
+#include "hytlass/layout/permute.h"
+#endif
+
+#include "cutlass_patch/epilogue/threadblock/epilogue_with_variadic.h"
+// #include "cutlass/epilogue/threadblock/epilogue_streamk_with_broadcast.h"
+
+namespace cutlass_patch {
+namespace epilogue {
+namespace threadblock {
+
+/// Defines sensible defaults for epilogues for SimtOps.
+template <typename Shape,
+          typename WarpMmaSimt,
+          typename ElementOutput,
+          typename OutputOp,
+          int ElementsPerAccess,
+          bool ScatterD = false,
+          typename PermuteDLayout = cutlass::layout::NoPermute,
+          cutlass::conv::StrideSupport StrideSupport =
+              cutlass::conv::StrideSupport::kUnity,
+          int Rank = 4>
+struct DefaultEpilogueWithVariadicSimt {
+  static cutlass::conv::StrideSupport const kStrideSupport = StrideSupport;
+  static int const kRank = Rank;
+
+  static bool const UseCUDAStore =
+      cutlass::platform::is_same<ElementOutput, double>::value;
+
+  /// Use defaults related to the existing epilogue
+  using Base = cutlass::epilogue::threadblock::
+      DefaultEpilogueSimt<Shape, WarpMmaSimt, OutputOp, ElementsPerAccess>;
+
+  using PackedOutputTileIterator =
+      cutlass::epilogue::threadblock::PredicatedTileIterator<
+          typename Base::OutputTileThreadMap,
+          ElementOutput,
+          ScatterD,
+          PermuteDLayout,
+          UseCUDAStore>;
+
+  using StridedOutputTileIterator =
+      cutlass::epilogue::threadblock::PredicatedTileIteratorConv<
+          typename Base::OutputTileThreadMap,
+          ElementOutput,
+          ScatterD,
+          PermuteDLayout,
+          UseCUDAStore,
+          kRank>;
+
+  //
+  // Stores the result z = (y = GEMM(A, B, C), variadic)
+  //
+  using OutputTileIterator = typename cutlass::platform::conditional<
+      StrideSupport == cutlass::conv::StrideSupport::kUnity,
+      PackedOutputTileIterator,
+      StridedOutputTileIterator>::type;
+
+  //
+  // Define the epilogue
+  //
+  using Epilogue = cutlass_patch::epilogue::threadblock::EpilogueWithVariadic<
+      Shape,
+      WarpMmaSimt,
+      Base::kPartitionsK,
+      OutputTileIterator,
+      typename Base::AccumulatorFragmentIterator,
+      typename Base::WarpTileIterator,
+      typename Base::SharedLoadIterator,
+      OutputOp,
+      typename Base::Padding>;
+};
+
+/// Defines sensible defaults for strided dgrad epilogues for SimtOps.
+template <typename Shape,
+          typename WarpMmaSimt,
+          typename ElementOutput,
+          typename OutputOp,
+          int ElementsPerAccess,
+          bool ScatterD = false,
+          typename PermuteDLayout = cutlass::layout::NoPermute>
+struct DefaultEpilogueWithVariadicSimtStridedDgrad {
+  /// Use defaults related to the existing epilogue
+  using Base = cutlass::epilogue::threadblock::DefaultEpilogueSimtStridedDgrad<
+      Shape,
+      WarpMmaSimt,
+      OutputOp,
+      ElementsPerAccess>;
+
+  //
+  // Stores the result z = (y = GEMM(A, B, C), variadic)
+  //
+  using OutputTileIterator =
+      cutlass::epilogue::threadblock::PredicatedTileIteratorStridedDgrad<
+          typename Base::OutputTileThreadMap,
+          ElementOutput>;
+
+  //
+  // Define the epilogue
+  //
+  using Epilogue = cutlass_patch::epilogue::threadblock::EpilogueWithVariadic<
+      Shape,
+      WarpMmaSimt,
+      Base::kPartitionsK,
+      OutputTileIterator,
+      typename Base::AccumulatorFragmentIterator,
+      typename Base::WarpTileIterator,
+      typename Base::SharedLoadIterator,
+      OutputOp,
+      typename Base::Padding>;
+};
+
+/// Defines sensible defaults for epilogues for TensorOps.
+template <typename Shape,
+          typename WarpMmaTensorOp,
+          int PartitionsK,
+          typename ElementOutput,
+          typename OutputOp,
+          int ElementsPerAccess,
+          bool ScatterD = false,
+          typename PermuteDLayout = cutlass::layout::NoPermute>
+struct DefaultEpilogueWithVariadicTensorOp {
+  /// Use defaults related to the existing epilogue
+  using Base = cutlass::epilogue::threadblock::DefaultEpilogueTensorOp<
+      Shape,
+      WarpMmaTensorOp,
+      PartitionsK,
+      OutputOp,
+      ElementsPerAccess>;
+
+  //
+  // Stores the result z = (y = GEMM(A, B, C), variadic)
+  //
+  using OutputTileIterator =
+      cutlass::epilogue::threadblock::PredicatedTileIterator<
+          typename Base::OutputTileThreadMap,
+          ElementOutput,
+          ScatterD,
+          PermuteDLayout>;
+
+  //
+  // Define the epilogue
+  //
+  using Epilogue = cutlass_patch::epilogue::threadblock::EpilogueWithVariadic<
+      Shape,
+      WarpMmaTensorOp,
+      PartitionsK,
+      OutputTileIterator,
+      typename Base::AccumulatorFragmentIterator,
+      typename Base::WarpTileIterator,
+      typename Base::SharedLoadIterator,
+      OutputOp,
+      typename Base::Padding,
+      Base::kFragmentsPerIteration>;
+};
+
+/// Defines sensible defaults for epilogues for VoltaTensorOps.
+template <typename Shape,
+          typename WarpMmaTensorOp,
+          int PartitionsK,
+          typename ElementOutput,
+          typename OutputOp,
+          int ElementsPerAccess>
+struct DefaultEpilogueWithVariadicVoltaTensorOp {
+  /// Use defaults related to the existing epilogue
+  using Base = cutlass::epilogue::threadblock::DefaultEpilogueVoltaTensorOp<
+      Shape,
+      WarpMmaTensorOp,
+      PartitionsK,
+      OutputOp,
+      ElementsPerAccess>;
+
+  //
+  // Stores the result z = (y = GEMM(A, B, C), variadic)
+  //
+  using OutputTileIterator = cutlass::epilogue::threadblock::
+      PredicatedTileIterator<typename Base::OutputTileThreadMap, ElementOutput>;
+
+  //
+  // Define the epilogue
+  //
+  using Epilogue = cutlass_patch::epilogue::threadblock::EpilogueWithVariadic<
+      Shape,
+      WarpMmaTensorOp,
+      PartitionsK,
+      OutputTileIterator,
+      typename Base::AccumulatorFragmentIterator,
+      typename Base::WarpTileIterator,
+      typename Base::SharedLoadIterator,
+      OutputOp,
+      typename Base::Padding>;
+};
+
+}  // namespace threadblock
+}  // namespace epilogue
+}  // namespace cutlass_patch

python/paddle/apy/matmul_pass/matmul/cutlass_patch/epilogue/threadblock/epilogue_with_variadic.h

@@ -0,0 +1,666 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/*! \file
+  \brief Epilogue for threadblock scoped GEMMs using Tensor Ops.
+
+  The epilogue rearranges the result of a matrix product through shared memory
+  to match canonical tensor layouts in global memory. Epilogues support
+  conversion and reduction operations.
+
+  The shared memory resource is time-sliced across warps.
+*/
+
+#pragma once
+
+#include "cutlass_patch/backend.h"
+
+#ifdef __NVCC__
+#include <cuda/std/cassert>
+
+#include "cutlass/aligned_buffer.h"
+#include "cutlass/array.h"
+#include "cutlass/functional.h"
+#include "cutlass/layout/tensor.h"
+#include "cutlass/layout/vector.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/tensor_coord.h"
+
+#include "cutlass/gemm/gemm.h"
+
+#include "cutlass/transform/pitch_linear_thread_map.h"
+#include "cutlass/transform/threadblock/regular_tile_iterator.h"
+
+#include "cutlass/epilogue/threadblock/epilogue_base.h"
+#include "cutlass/epilogue/threadblock/epilogue_base_streamk.h"
+#include "cutlass/epilogue/threadblock/predicated_tile_iterator.h"
+#elif defined(__HIPCC__)
+#include "hytlass/aligned_buffer.h"
+#include "hytlass/array.h"
+#include "hytlass/functional.h"
+#include "hytlass/layout/tensor.h"
+#include "hytlass/layout/vector.h"
+#include "hytlass/numeric_types.h"
+#include "hytlass/tensor_coord.h"
+
+#include "hytlass/gemm/gemm.h"
+
+#include "hytlass/transform/pitch_linear_thread_map.h"
+#include "hytlass/transform/threadblock/regular_tile_iterator.h"
+
+#include "hytlass/epilogue/threadblock/epilogue_base.h"
+#include "hytlass/epilogue/threadblock/epilogue_base_streamk.h"
+#include "hytlass/epilogue/threadblock/predicated_tile_iterator.h"
+#endif
+
+#include "cutlass_patch/trace_device.h"
+
+namespace cutlass_patch {
+namespace epilogue {
+namespace threadblock {
+
+/// Epilogue operator
+template <
+    typename Shape_,  ///< Shape of threadblock tile (concept: GemmShape)
+    typename WarpMmaOperator_,     ///< Warp-level MMA operator (concept:
+                                   ///< gemm::warp::MmaTensorOp)
+    int PartitionsK,               ///< Number of partitions of the K dimension
+    typename OutputTileIterator_,  ///< Tile iterator reading and writing
+                                   ///< output tensors
+    typename AccumulatorFragmentIterator_,  ///< Fragment iterator
+                                            ///< selecting accumulators
+    typename WarpTileIterator_,    ///< Warp-scoped tile iterator writing
+                                   ///< accumulators to SMEM
+    typename SharedLoadIterator_,  ///< Threadblock-scoped tile iterator
+                                   ///< loading from SMEM
+    typename OutputOp_,            ///< Output operator
+    typename Padding_,  ///< Padding added to SMEM allocation to avoid
+                        ///< bank conflicts (concept: MatrixShape)
+    int FragmentsPerPartition =
+        1,                  ///< Used to coarsten the epilogue granularity
+    int IterationsUnroll =  ///< Used to reduce binary size when epilogue
+                            ///< op is large
+    (!cutlass::epilogue::threadblock::IsEpilogueFunctorHeavy<OutputOp_>::value)>
+class EpilogueWithVariadic
+    : public cutlass::epilogue::threadblock::EpilogueBase<
+          Shape_,
+          typename WarpMmaOperator_::Shape,
+          PartitionsK,
+          AccumulatorFragmentIterator_,
+          WarpTileIterator_,
+          Padding_,
+          FragmentsPerPartition>,
+      public cutlass::epilogue::threadblock::EpilogueBaseStreamK<
+          Shape_,
+          PartitionsK,
+          WarpMmaOperator_,
+          AccumulatorFragmentIterator_> {
+ public:
+  using Base = cutlass::epilogue::threadblock::EpilogueBase<
+      Shape_,
+      typename WarpMmaOperator_::Shape,
+      PartitionsK,
+      AccumulatorFragmentIterator_,
+      WarpTileIterator_,
+      Padding_,
+      FragmentsPerPartition>;
+
+  using BaseStreamK = cutlass::epilogue::threadblock::EpilogueBaseStreamK<
+      Shape_,
+      PartitionsK,
+      WarpMmaOperator_,
+      AccumulatorFragmentIterator_>;
+
+  using Shape = Shape_;
+  using WarpMmaOperator = WarpMmaOperator_;
+  static int const kPartitionsK = PartitionsK;
+  using OutputTileIterator = OutputTileIterator_;
+  using AccumulatorFragmentIterator = AccumulatorFragmentIterator_;
+  using WarpTileIterator = WarpTileIterator_;
+  using SharedLoadIterator = SharedLoadIterator_;
+  using OutputOp = OutputOp_;
+  using Padding = Padding_;
+  using Layout = cutlass::layout::RowMajor;
+  using LongIndex = typename Layout::LongIndex;
+
+  /// Number of warps per block
+  using WarpCount = typename Base::WarpCount;
+
+  /// Number of threads per block
+  static int const kBlockThreads = 32 * WarpCount::kCount;
+
+  /// Per-thread accumulator tile type
+  using AccumulatorTile = typename Base::AccumulatorTile;
+
+  /// Numerical accumulation element type
+  using ElementAccumulator = typename WarpMmaOperator::ElementC;
+
+  /// Fragment type used by the accumulator tile's fragment iterator
+  using AccumulatorFragment = typename AccumulatorFragmentIterator::Fragment;
+
+  /// Output element
+  using ElementOutput = typename OutputTileIterator::Element;
+
+  /// Output access size
+  static int const kElementsPerAccess = OutputTileIterator::kElementsPerAccess;
+
+  /// Tensor reference to destination tensor
+  using TensorRef = typename OutputTileIterator::TensorRef;
+
+  /// Tensor reference to sync tensor
+  using SyncTensorRef =
+      typename cutlass::TensorRef<int, cutlass::layout::PackedVectorLayout>;
+
+  /// Const tensor reference to source tensor
+  using ConstTensorRef = typename OutputTileIterator::ConstTensorRef;
+
+  /// Vector type used by the global output iterator
+  using OutputAccessType =
+      cutlass::Array<typename OutputTileIterator::Element,
+                     OutputTileIterator::kElementsPerAccess>;
+
+  /// Vector type used by the shared output iterator
+  using AccumulatorAccessType =
+      cutlass::Array<typename WarpTileIterator::Element,
+                     OutputTileIterator::kElementsPerAccess>;
+
+  static int constexpr kSmemTiles = Base::kFragmentsPerIteration > 1
+                                        ? Base::kFragmentsPerIteration
+                                        : kPartitionsK;
+
+  static int constexpr kSmemPointerOffset =
+      Base::SharedStorage::StorageShape::kCount / kSmemTiles;
+
+ public:
+  static_assert(
+      SharedLoadIterator::Fragment::kElements ==
+          OutputTileIterator::Fragment::kElements,
+      "Mismatch between shared load iterator and output tile iterator.");
+
+  static_assert(OutputTileIterator::kElementsPerAccess,
+                "OutputTileIterator::kElementsPerAccess must not be zero.");
+
+  static_assert(!(OutputTileIterator::Fragment::kElements %
+                  OutputTileIterator::kElementsPerAccess),
+                "Divisibility");
+
+  static_assert(kPartitionsK == 1 || Base::kFragmentsPerIteration == 1,
+                "One of these must be exactly 1.");
+
+ public:
+  /// Aspect for when epilogue source is not needed
+  struct SourceAspectNotNeeded {
+    /// Constructor
+    CUTLASS_DEVICE
+    SourceAspectNotNeeded() {}
+
+    // No-op
+    CUTLASS_DEVICE
+    void load() {}
+
+    /// Invoke the output functor over each vector of output
+    CUTLASS_DEVICE
+    void apply_output_operator(
+        const OutputTileIterator &output_iterator,
+        typename OutputTileIterator::Fragment &output_fragment,  // NOLINT
+        OutputOp const &output_op,
+        typename SharedLoadIterator::Fragment const &aligned_accum_fragment) {
+      CUTLASS_TRACE_DEVICE("");
+
+      OutputAccessType *output_frag_ptr =
+          reinterpret_cast<OutputAccessType *>(&output_fragment);
+
+      AccumulatorAccessType const *compute_frag_ptr =
+          reinterpret_cast<AccumulatorAccessType const *>(
+              &aligned_accum_fragment);
+
+      const int32_t thread_start_row = output_iterator.thread_start_row();
+      const int32_t thread_start_column = output_iterator.thread_start_column();
+
+      const typename OutputTileIterator::Index extent_row =
+          output_iterator.extent_row();
+      const typename OutputTileIterator::Index extent_column =
+          output_iterator.extent_column();
+
+      using ThreadMap = typename OutputTileIterator::ThreadMap;
+
+      typename OutputTileIterator::Mask mask;
+      output_iterator.get_mask(mask);
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster;
+           ++cluster) {
+        CUTLASS_PRAGMA_UNROLL
+        for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+          CUTLASS_PRAGMA_UNROLL
+          for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+            int frag_row_idx =
+                (row + ThreadMap::Iterations::kRow *
+                           (group + ThreadMap::Iterations::kGroup * cluster));
+
+            int row_offset = thread_start_row + row * ThreadMap::Delta::kRow +
+                             group * ThreadMap::Delta::kGroup +
+                             cluster * ThreadMap::Delta::kCluster;
+
+            bool row_guard = row_offset < extent_row;
+
+            CUTLASS_PRAGMA_UNROLL
+            for (int column = 0; column < ThreadMap::Iterations::kColumn;
+                 ++column) {
+              bool guard = row_guard && mask.predicates[column];
+              if (!guard) {
+                continue;
+              }
+
+              int column_offset =
+                  thread_start_column + column * ThreadMap::Delta::kColumn;
+              int frag_offset =
+                  frag_row_idx * ThreadMap::Iterations::kColumn + column;
+
+              output_frag_ptr[frag_offset] = output_op(
+                  compute_frag_ptr[frag_offset], row_offset, column_offset);
+            }
+          }
+        }
+      }
+    }
+  };
+
+  /// Aspect for when epilogue source is needed
+  struct SourceAspectNeeded {
+    OutputTileIterator source_iterator;
+
+    typename OutputTileIterator::Fragment source_fragment;
+
+    /// Invoke the output functor over each vector of output
+    CUTLASS_DEVICE
+    static void apply_output_operator(
+        const OutputTileIterator &output_iterator,
+        typename OutputTileIterator::Fragment &output_fragment,  // NOLINT
+        OutputOp const &output_op,
+        typename SharedLoadIterator::Fragment const &aligned_accum_fragment,
+        typename OutputTileIterator::Fragment const &source_fragment) {
+      CUTLASS_TRACE_DEVICE("");
+
+      OutputAccessType *output_frag_ptr =
+          reinterpret_cast<OutputAccessType *>(&output_fragment);
+
+      AccumulatorAccessType const *compute_frag_ptr =
+          reinterpret_cast<AccumulatorAccessType const *>(
+              &aligned_accum_fragment);
+
+      OutputAccessType const *source_frag_ptr =
+          reinterpret_cast<OutputAccessType const *>(&source_fragment);
+
+      typename OutputTileIterator::Element const *source_ptr =
+          reinterpret_cast<typename OutputTileIterator::Element const *>(
+              &source_fragment);
+
+      const int32_t thread_start_row = output_iterator.thread_start_row();
+      const int32_t thread_start_column = output_iterator.thread_start_column();
+
+      const typename OutputTileIterator::Index extent_row =
+          output_iterator.extent_row();
+      const typename OutputTileIterator::Index extent_column =
+          output_iterator.extent_column();
+
+      using ThreadMap = typename OutputTileIterator::ThreadMap;
+
+      typename OutputTileIterator::Mask mask;
+      output_iterator.get_mask(mask);
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int cluster = 0; cluster < ThreadMap::Iterations::kCluster;
+           ++cluster) {
+        CUTLASS_PRAGMA_UNROLL
+        for (int group = 0; group < ThreadMap::Iterations::kGroup; ++group) {
+          CUTLASS_PRAGMA_UNROLL
+          for (int row = 0; row < ThreadMap::Iterations::kRow; ++row) {
+            int frag_row_idx =
+                (row + ThreadMap::Iterations::kRow *
+                           (group + ThreadMap::Iterations::kGroup * cluster));
+
+            int row_offset = thread_start_row + row * ThreadMap::Delta::kRow +
+                             group * ThreadMap::Delta::kGroup +
+                             cluster * ThreadMap::Delta::kCluster;
+
+            bool row_guard = row_offset < extent_row;
+
+            CUTLASS_PRAGMA_UNROLL
+            for (int column = 0; column < ThreadMap::Iterations::kColumn;
+                 ++column) {
+              bool guard = row_guard && mask.predicates[column];
+              if (!guard) {
+                continue;
+              }
+
+              int column_offset =
+                  thread_start_column + column * ThreadMap::Delta::kColumn;
+              int frag_offset =
+                  frag_row_idx * ThreadMap::Iterations::kColumn + column;
+
+              output_frag_ptr[frag_offset] =
+                  output_op(compute_frag_ptr[frag_offset],
+                            source_frag_ptr[frag_offset],
+                            row_offset,
+                            column_offset);
+            }
+          }
+        }
+      }
+    }
+
+    /// Constructor
+    CUTLASS_DEVICE
+    explicit SourceAspectNeeded(OutputTileIterator source_iterator)
+        : source_iterator(source_iterator) {
+      source_fragment.clear();
+    }
+
+    // Load addend source fragment from global memory
+    CUTLASS_DEVICE
+    void load() {
+      source_iterator.load(source_fragment);
+      ++source_iterator;
+    }
+
+    /// Invoke the output functor over each vector of output
+    CUTLASS_DEVICE
+    void apply_output_operator(
+        const OutputTileIterator &output_iterator,
+        typename OutputTileIterator::Fragment &output_fragment,  // NOLINT
+        OutputOp const &output_op,
+        typename SharedLoadIterator::Fragment const &aligned_accum_fragment) {
+      apply_output_operator(output_iterator,
+                            output_fragment,
+                            output_op,
+                            aligned_accum_fragment,
+                            source_fragment);
+    }
+  };
+
+ private:
+  /// Loads fragment from shared memory aligned with output tensor
+  SharedLoadIterator shared_load_iterator_;
+
+  /// Thread index in the threadblock
+  int thread_idx;
+
+  /// Warp index in the threadblock
+  int warp_idx;
+
+ public:
+  /// Constructor
+  CUTLASS_DEVICE
+  EpilogueWithVariadic(
+      typename Base::SharedStorage
+          &shared_storage,  // NOLINT  ///< Shared storage object
+      int thread_idx,       ///< ID of a thread within the threadblock
+      int warp_idx,         ///< ID of warp within threadblock
+      int lane_idx)         ///< Id of thread within warp
+      : Base(shared_storage, thread_idx, warp_idx, lane_idx),
+        BaseStreamK(thread_idx),
+        shared_load_iterator_(shared_storage.reference(), thread_idx),
+        thread_idx(thread_idx),
+        warp_idx(warp_idx) {}
+
+  /// Aggregates the accumulator sets shared by peer blocks in the global
+  /// workspace, performing epilogue computations, writing to output
+  CUTLASS_DEVICE
+  void reduce(int peer_idx_begin,
+              int peer_idx_end,
+              int reduce_fragment_idx,
+              void *element_workspace,
+              OutputOp const &output_op,  ///< Output operator
+              OutputTileIterator
+                  destination_iterator,  ///< Tile iterator for destination
+              OutputTileIterator
+                  source_iterator) {  ///< Threadblock tile coordinate in GEMM
+                                      ///< (in units of threadblock tiles)
+    CUTLASS_TRACE_DEVICE("");
+
+    // Reduce peer accumulator fragments into one fragment
+    AccumulatorFragment accum_fragment;
+    BaseStreamK::reduce(accum_fragment,
+                        peer_idx_begin,
+                        peer_idx_end,
+                        reduce_fragment_idx,
+                        element_workspace);
+
+    // Store fragment to shared memory
+    this->warp_tile_iterator_.store(accum_fragment);
+
+    __syncthreads();
+
+    // Initialize/load source-fragment data
+    typename OutputTileIterator::Fragment source_fragment;
+    source_fragment.clear();
+
+    if (output_op.is_source_needed()) {
+      source_iterator += reduce_fragment_idx;
+      source_iterator.load(source_fragment);
+    }
+
+    // Load fragment from shared memory
+    typename SharedLoadIterator::Fragment aligned_accum_fragment;
+    shared_load_iterator_.load(aligned_accum_fragment);
+
+    // Add fragments shared by other k partitions
+    if (kPartitionsK > 1) {
+      cutlass::plus<typename SharedLoadIterator::Fragment> add_fragments;
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 1; i < kPartitionsK; ++i) {
+        typename SharedLoadIterator::Fragment aligned_addend_fragment;
+        shared_load_iterator_.add_pointer_offset(kSmemPointerOffset);
+        shared_load_iterator_.load(aligned_addend_fragment);
+        aligned_accum_fragment =
+            add_fragments(aligned_accum_fragment, aligned_addend_fragment);
+      }
+    }
+
+    // Compute the output result
+    typename OutputTileIterator::Fragment output_fragment;
+
+    // Apply the output operator
+    SourceAspectNeeded::apply_output_operator(
+        output_fragment, output_op, aligned_accum_fragment, source_fragment);
+
+    // Store the final result
+    destination_iterator += reduce_fragment_idx;
+    destination_iterator.store(output_fragment);
+  }
+
+  /// Perform the epilogue computations and stream the result to global memory.
+  CUTLASS_DEVICE
+  void operator()(OutputOp const &output_op,  ///< Output operator
+                  OutputTileIterator
+                      destination_iterator,  ///< Tile iterator for destination
+                  AccumulatorTile const &
+                      accumulators) {  ///< Complete warp-level accumulator tile
+    CUTLASS_TRACE_DEVICE("");
+    operator()(
+        output_op, destination_iterator, accumulators, SourceAspectNotNeeded());
+  }
+
+  /// Perform the epilogue computations and stream the result to global memory.
+  /// Implements two alternative codepaths, depending on whether the output op
+  /// requires addend data to be loaded.
+  CUTLASS_DEVICE
+  void operator()(OutputOp const &output_op,  ///< Output operator
+                  OutputTileIterator
+                      destination_iterator,  ///< Tile iterator for destination
+                  AccumulatorTile const
+                      &accumulators,  ///< Complete warp-level accumulator tile
+                  OutputTileIterator
+                      source_iterator) {  ///< Tile iterator for addend source
+    CUTLASS_TRACE_DEVICE("");
+    if (output_op.is_source_needed()) {
+      operator()(output_op,
+                 destination_iterator,
+                 accumulators,
+                 SourceAspectNeeded(source_iterator));
+    } else {
+      operator()(output_op,
+                 destination_iterator,
+                 accumulators,
+                 SourceAspectNotNeeded());
+    }
+  }
+
+  /// Perform the epilogue computations and stream the result to global memory.
+  /// Implements a single codepath, regardless of whether the output op requires
+  /// addend data to be loaded
+  CUTLASS_DEVICE
+  void unified(OutputOp const &output_op,  ///< Output operator
+               OutputTileIterator
+                   destination_iterator,  ///< Tile iterator for destination
+               AccumulatorTile const
+                   &accumulators,  ///< Complete warp-level accumulator tile
+               OutputTileIterator
+                   source_iterator) {  ///< Tile iterator for addend source
+    CUTLASS_TRACE_DEVICE("");
+    if (!output_op.is_source_needed()) {
+      source_iterator.clear_mask();
+      __syncthreads();  // Dummy (CUDA 11.0)
+    }
+
+    operator()(output_op,
+               destination_iterator,
+               accumulators,
+               SourceAspectNeeded(source_iterator));
+  }
+
+  template <class Seq>
+  struct acc2smem;
+
+  template <size_t... Seq>
+  struct acc2smem<cutlass::index_sequence<Seq...>> {
+    template <int Advance>
+    CUTLASS_DEVICE static void helper(
+        AccumulatorFragmentIterator accum_fragment_iterator,
+        WarpTileIterator &warp_tile_iterator) {  // NOLINT
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < Advance; i++) {
+        ++accum_fragment_iterator;
+      }
+
+      typename AccumulatorFragmentIterator::Fragment accum_fragment;
+
+      accum_fragment_iterator.load(accum_fragment);
+      ++accum_fragment_iterator;
+      warp_tile_iterator.store(accum_fragment);
+    }
+
+    CUTLASS_DEVICE
+    static void push(size_t pos,
+                     AccumulatorFragmentIterator const &iterator_begin,
+                     WarpTileIterator &warp_tile_iterator) {  // NOLINT
+      int dummy[] = {(pos == Seq) &&
+                     (helper<Seq>(iterator_begin, warp_tile_iterator), 0)...};
+    }
+  };
+
+  /// Streams the result to global memory
+  template <typename SourceAspect>
+  CUTLASS_DEVICE void operator()(
+      OutputOp const &output_op,  ///< Output operator
+      OutputTileIterator
+          destination_iterator,  ///< Tile iterator for destination
+      AccumulatorTile const
+          &accumulators,  ///< Complete warp-level accumulator tile
+      SourceAspect source) {
+    CUTLASS_TRACE_DEVICE("");
+
+    // Iterator over warp-level accumulator fragment
+    AccumulatorFragmentIterator accum_fragment_iterator(accumulators);
+
+    //
+    // Iterate over accumulator tile
+    //
+
+#ifdef __clang__
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wcuda-compat"
+// Turn off clangs warning about loop unroll argument using parens.
+#endif
+
+#pragma unroll(IterationsUnroll ? OutputTileIterator::kIterations : 1)
+    for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter) {
+      //
+      // Load the source
+      //
+
+      source.load();
+      //
+      // Convert and store fragment
+      //
+
+      __syncthreads();
+
+      acc2smem<cutlass::make_index_sequence<OutputTileIterator::kIterations>>::
+          push(iter, accum_fragment_iterator, this->warp_tile_iterator_);
+
+      __syncthreads();
+
+      //
+      // Load fragments from shared memory
+      //
+
+      typename SharedLoadIterator::Fragment
+          aligned_accum_fragment[kPartitionsK];
+      shared_load_iterator_.load(aligned_accum_fragment[0]);
+
+      if (kPartitionsK > 1) {
+        cutlass::plus<typename SharedLoadIterator::Fragment> add_fragments;
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 1; i < kPartitionsK; ++i) {
+          shared_load_iterator_.add_pointer_offset(kSmemPointerOffset);
+          shared_load_iterator_.load(aligned_accum_fragment[i]);
+          aligned_accum_fragment[0] = add_fragments(aligned_accum_fragment[0],
+                                                    aligned_accum_fragment[i]);
+        }
+
+        shared_load_iterator_.add_pointer_offset((1 - kPartitionsK) *
+                                                 kSmemPointerOffset);
+      }
+
+      //
+      // Compute the output result
+      //
+
+      typename OutputTileIterator::Fragment output_fragment;
+      source.apply_output_operator(destination_iterator,
+                                   output_fragment,
+                                   output_op,
+                                   aligned_accum_fragment[0]);
+
+      //
+      // Store the final result
+      //
+
+      destination_iterator.store(output_fragment);
+      ++destination_iterator;
+    }
+
+#ifdef __clang__
+#pragma clang diagnostic pop
+#endif
+  }
+};
+
+}  // namespace threadblock
+}  // namespace epilogue
+}  // namespace cutlass_patch

python/paddle/apy/matmul_pass/matmul/cutlass_patch/gemm/device/gemm_universal_with_variadic.h

@@ -0,0 +1,416 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/*! \file
+    \brief
+*/
+
+#pragma once
+
+#include "cutlass_patch/backend.h"
+
+#ifdef __NVCC__
+#include "cutlass/arch/arch.h"
+#include "cutlass/arch/mma.h"
+#include "cutlass/device_kernel.h"
+#include "cutlass/numeric_types.h"
+
+#include "cutlass/gemm/gemm.h"
+#include "cutlass/gemm/kernel/gemm_universal.h"
+#include "cutlass/gemm/threadblock/threadblock_swizzle.h"
+
+#include "cutlass/gemm/device/default_gemm_configuration.h"
+#include "cutlass/gemm/device/gemm_universal_base.h"
+#include "cutlass/gemm/kernel/default_gemm_universal.h"
+
+#include "cutlass/layout/permute.h"
+#elif defined(__HIPCC__)
+#include "hytlass/arch/arch.h"
+#include "hytlass/arch/mma.h"
+#include "hytlass/device_kernel.h"
+#include "hytlass/numeric_types.h"
+
+#include "hytlass/gemm/gemm.h"
+#include "hytlass/gemm/kernel/gemm_universal.h"
+#include "hytlass/gemm/threadblock/threadblock_swizzle.h"
+
+#include "hytlass/gemm/device/default_gemm_configuration.h"
+#include "hytlass/gemm/device/gemm_universal_base.h"
+#include "hytlass/gemm/kernel/default_gemm_universal.h"
+
+#include "hytlass/layout/permute.h"
+#endif
+
+#include "cutlass_patch/gemm/kernel/default_gemm_with_variadic.h"
+
+namespace cutlass_patch {
+namespace gemm {
+namespace device {
+
+/*!
+  GemmUniversal with variadic epilogues.
+*/
+template <
+    /// Element type for A matrix operand
+    typename ElementA_,
+    /// Layout type for A matrix operand
+    typename LayoutA_,
+    /// Element type for B matrix operand
+    typename ElementB_,
+    /// Layout type for B matrix operand
+    typename LayoutB_,
+    /// Element type for C and D matrix operands
+    typename ElementC_,
+    /// Layout type for C and D matrix operands
+    typename LayoutC_,
+    /// Element type for internal accumulation
+    typename ElementAccumulator_ = ElementC_,
+    /// Operator class tag
+    typename OperatorClass_ = cutlass::arch::OpClassSimt,
+    /// Tag indicating architecture to tune for.  This is the minimum SM that
+    /// supports the intended feature. The device kernel can be built
+    /// targeting any SM larger than this number.
+    typename ArchTag_ = cutlass::arch::Sm70,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape_ = typename cutlass::gemm::device::
+        DefaultGemmConfiguration<OperatorClass_,
+                                 ArchTag_,
+                                 ElementA_,
+                                 ElementB_,
+                                 ElementC_,
+                                 ElementAccumulator_>::ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape_ = typename cutlass::gemm::device::
+        DefaultGemmConfiguration<OperatorClass_,
+                                 ArchTag_,
+                                 ElementA_,
+                                 ElementB_,
+                                 ElementC_,
+                                 ElementAccumulator_>::WarpShape,
+    /// Instruction-level tile size (concept: GemmShape)
+    typename InstructionShape_ = typename cutlass::gemm::device::
+        DefaultGemmConfiguration<OperatorClass_,
+                                 ArchTag_,
+                                 ElementA_,
+                                 ElementB_,
+                                 ElementC_,
+                                 ElementAccumulator_>::InstructionShape,
+    /// Epilogue output operator
+    typename EpilogueOutputOp_ = typename cutlass::gemm::device::
+        DefaultGemmConfiguration<OperatorClass_,
+                                 ArchTag_,
+                                 ElementA_,
+                                 ElementB_,
+                                 ElementC_,
+                                 ElementAccumulator_>::EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle_ =
+        cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
+    /// Number of stages used in the pipelined mainloop
+    int Stages = cutlass::gemm::device::DefaultGemmConfiguration<
+        OperatorClass_,
+        ArchTag_,
+        ElementA_,
+        ElementB_,
+        ElementC_,
+        ElementAccumulator_>::kStages,
+    /// Access granularity of A matrix in units of elements
+    int AlignmentA = cutlass::gemm::device::DefaultGemmConfiguration<
+        OperatorClass_,
+        ArchTag_,
+        ElementA_,
+        ElementB_,
+        ElementC_,
+        ElementAccumulator_>::kAlignmentA,
+    /// Access granularity of B matrix in units of elements
+    int AlignmentB = cutlass::gemm::device::DefaultGemmConfiguration<
+        OperatorClass_,
+        ArchTag_,
+        ElementA_,
+        ElementB_,
+        ElementC_,
+        ElementAccumulator_>::kAlignmentB,
+    /// Operation performed by GEMM
+    typename Operator_ = typename cutlass::gemm::device::
+        DefaultGemmConfiguration<OperatorClass_,
+                                 ArchTag_,
+                                 ElementA_,
+                                 ElementB_,
+                                 ElementC_,
+                                 ElementAccumulator_>::Operator,
+    /// Complex elementwise transformation on A operand
+    cutlass::ComplexTransform TransformA = cutlass::ComplexTransform::kNone,
+    /// Complex elementwise transformation on B operand
+    cutlass::ComplexTransform TransformB = cutlass::ComplexTransform::kNone>
+class GemmUniversalWithVariadic
+    : public cutlass::gemm::device::GemmUniversalBase<
+          typename cutlass_patch::gemm::kernel::DefaultGemmWithVariadic<
+              ElementA_,
+              LayoutA_,
+              TransformA,
+              AlignmentA,
+              ElementB_,
+              LayoutB_,
+              TransformB,
+              AlignmentB,
+              ElementC_,
+              LayoutC_,
+              ElementAccumulator_,
+              OperatorClass_,
+              ArchTag_,
+              ThreadblockShape_,
+              WarpShape_,
+              InstructionShape_,
+              EpilogueOutputOp_,
+              ThreadblockSwizzle_,
+              Stages,
+              Operator_>::GemmKernel> {
+ public:
+  using ElementAccumulator = ElementAccumulator_;
+  using OperatorClass = OperatorClass_;
+  using ArchTag = ArchTag_;
+  using ThreadblockShape = ThreadblockShape_;
+  using WarpShape = WarpShape_;
+  using InstructionShape = InstructionShape_;
+  using EpilogueOutputOp = EpilogueOutputOp_;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
+  using Operator = Operator_;
+  static int const kStages = Stages;
+  static int const kAlignmentA = AlignmentA;
+  static int const kAlignmentB = AlignmentB;
+  static int const kAlignmentC = EpilogueOutputOp::kCount;
+  static cutlass::ComplexTransform const kTransformA = TransformA;
+  static cutlass::ComplexTransform const kTransformB = TransformB;
+
+  using Base = cutlass::gemm::device::GemmUniversalBase<
+      typename cutlass_patch::gemm::kernel::DefaultGemmWithVariadic<
+          ElementA_,
+          LayoutA_,
+          TransformA,
+          AlignmentA,
+          ElementB_,
+          LayoutB_,
+          TransformB,
+          AlignmentB,
+          ElementC_,
+          LayoutC_,
+          ElementAccumulator_,
+          OperatorClass_,
+          ArchTag_,
+          ThreadblockShape_,
+          WarpShape_,
+          InstructionShape_,
+          EpilogueOutputOp_,
+          ThreadblockSwizzle_,
+          Stages,
+          Operator_>::GemmKernel>;
+
+  using Arguments = typename Base::Arguments;
+  using GemmKernel = typename Base::GemmKernel;
+};
+
+/// Partial specialization for column-major output exchanges problem size and
+/// operand.
+template <
+    /// Element type for A matrix operand
+    typename ElementA_,
+    /// Layout type for A matrix operand
+    typename LayoutA_,
+    /// Element type for B matrix operand
+    typename ElementB_,
+    /// Layout type for B matrix operand
+    typename LayoutB_,
+    /// Element type for C and D matrix operands
+    typename ElementC_,
+    /// Element type for internal accumulation
+    typename ElementAccumulator_,
+    /// Operator class tag
+    typename OperatorClass_,
+    /// Tag indicating architecture to tune for.  This is the minimum SM that
+    /// supports the intended feature. The device kernel can be built
+    /// targeting any SM larger than this number.
+    typename ArchTag_,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape_,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape_,
+    /// Instruction-level tile size (concept: GemmShape)
+    typename InstructionShape_,
+    /// Epilogue output operator
+    typename EpilogueOutputOp_,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle_,
+    /// Number of stages used in the pipelined mainloop
+    int Stages,
+    /// Access granularity of A matrix in units of elements
+    int AlignmentA,
+    /// Access granularity of B matrix in units of elements
+    int AlignmentB,
+    /// Operation performed by GEMM
+    typename Operator_,
+    /// Complex elementwise transformation on A operand
+    cutlass::ComplexTransform TransformA,
+    /// Complex elementwise transformation on B operand
+    cutlass::ComplexTransform TransformB>
+class GemmUniversalWithVariadic<ElementA_,
+                                LayoutA_,
+                                ElementB_,
+                                LayoutB_,
+                                ElementC_,
+                                cutlass::layout::ColumnMajor,  // partially
+                                                               // specialized on
+                                                               // LayoutC
+                                ElementAccumulator_,
+                                OperatorClass_,
+                                ArchTag_,
+                                ThreadblockShape_,
+                                WarpShape_,
+                                InstructionShape_,
+                                EpilogueOutputOp_,
+                                ThreadblockSwizzle_,
+                                Stages,
+                                AlignmentA,
+                                AlignmentB,
+                                Operator_,
+                                TransformA,
+                                TransformB> {
+ public:
+  using ElementA = ElementA_;
+  using LayoutA = LayoutA_;
+  using TensorRefA = cutlass::TensorRef<ElementA const, LayoutA>;
+  using ElementB = ElementB_;
+  using LayoutB = LayoutB_;
+  using TensorRefB = cutlass::TensorRef<ElementB const, LayoutB>;
+  using ElementC = ElementC_;
+  using LayoutC = cutlass::layout::ColumnMajor;
+  using TensorRefC = cutlass::TensorRef<ElementC const, LayoutC>;
+  using TensorRefD = cutlass::TensorRef<ElementC, LayoutC>;
+  using ElementAccumulator = ElementAccumulator_;
+  using OperatorClass = OperatorClass_;
+  using ArchTag = ArchTag_;
+  using ThreadblockShape = ThreadblockShape_;
+  using WarpShape = WarpShape_;
+  using InstructionShape = InstructionShape_;
+  using EpilogueOutputOp = EpilogueOutputOp_;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
+  using Operator = Operator_;
+  static int const kStages = Stages;
+  static int const kAlignmentA = AlignmentA;
+  static int const kAlignmentB = AlignmentB;
+  static cutlass::ComplexTransform const kTransformA = TransformA;
+  static cutlass::ComplexTransform const kTransformB = TransformB;
+
+  using UnderlyingOperator = typename GemmUniversalWithVariadic<
+      ElementB,
+      typename cutlass::layout::LayoutTranspose<LayoutB>::type,
+      ElementA,
+      typename cutlass::layout::LayoutTranspose<LayoutA>::type,
+      ElementC,
+      cutlass::layout::RowMajor,
+      ElementAccumulator,
+      OperatorClass,
+      ArchTag,
+      ThreadblockShape,
+      WarpShape,
+      InstructionShape,
+      EpilogueOutputOp,
+      ThreadblockSwizzle,
+      Stages,
+      kAlignmentB,
+      kAlignmentA,
+      Operator,
+      kTransformB,
+      kTransformA>::Base;
+
+  using GemmKernel = typename UnderlyingOperator::GemmKernel;
+  static int const kAlignmentC = EpilogueOutputOp::kCount;
+
+  /// Argument structure
+  using Arguments = typename UnderlyingOperator::Arguments;
+
+ private:
+  UnderlyingOperator underlying_operator_;
+
+ public:
+  /// Constructs the GEMM.
+  GemmUniversalWithVariadic() {}
+
+  /// Helper to construct a transposed equivalent for the underlying GEMM
+  /// operator
+  static Arguments to_underlying_arguments(Arguments const &args) {
+    return args.transposed_problem();
+  }
+
+  /// Determines whether the GEMM can execute the given problem.
+  static cutlass::Status can_implement(Arguments const &args) {
+    return UnderlyingOperator::can_implement(to_underlying_arguments(args));
+  }
+
+  /// Gets the workspace size
+  static size_t get_workspace_size(Arguments const &args) {
+    return UnderlyingOperator::get_workspace_size(
+        to_underlying_arguments(args));
+  }
+
+  /// Computes the grid shape
+  static dim3 get_grid_shape(Arguments const &args) {
+    return UnderlyingOperator::get_grid_shape(to_underlying_arguments(args));
+  }
+
+  /// Computes the maximum number of active blocks per multiprocessor
+  static int maximum_active_blocks(int smem_capacity = -1) {
+    return UnderlyingOperator::maximum_active_blocks(smem_capacity);
+  }
+
+  /// Initializes GEMM state from arguments.
+  cutlass::Status initialize(Arguments const &args,
+                             void *workspace = nullptr,
+                             GPUStream_t stream = nullptr) {
+    return underlying_operator_.initialize(
+        to_underlying_arguments(args), workspace, stream);
+  }
+
+  /// Lightweight update given a subset of arguments
+  cutlass::Status update(Arguments const &args, void *workspace = nullptr) {
+    return underlying_operator_.update(to_underlying_arguments(args),
+                                       workspace);
+  }
+
+  /// Runs the kernel using initialized state.
+  cutlass::Status run(GPUStream_t stream = nullptr) {
+    return underlying_operator_.run(stream);
+  }
+
+  /// Runs the kernel using initialized state.
+  cutlass::Status operator()(GPUStream_t stream = nullptr) {
+    return run(stream);
+  }
+
+  /// Runs the kernel using initialized state.
+  cutlass::Status operator()(Arguments const &args,
+                             void *workspace = nullptr,
+                             GPUStream_t stream = nullptr) {
+    cutlass::Status status = initialize(args, workspace, stream);
+
+    if (status == cutlass::Status::kSuccess) {
+      status = run(stream);
+    }
+
+    return status;
+  }
+};
+
+}  // namespace device
+}  // namespace gemm
+}  // namespace cutlass_patch

python/paddle/apy/matmul_pass/matmul/cutlass_patch/gemm/kernel/default_gemm_with_variadic.h

@@ -0,0 +1,227 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/*! \file
+  \brief
+    Defines a GEMM with Reduction based on an existing UniversalGemm kernel.
+
+*/
+
+#pragma once
+
+#include "cutlass_patch/backend.h"
+
+#ifdef __NVCC__
+#include "cutlass/gemm/kernel/default_gemm_universal.h"
+#include "cutlass/gemm/kernel/gemm_universal.h"
+#elif defined(__HIPCC__)
+#include "hytlass/gemm/kernel/default_gemm_universal.h"
+#include "hytlass/gemm/kernel/gemm_universal.h"
+#endif
+
+#include "cutlass_patch/epilogue/threadblock/default_epilogue_with_variadic.h"
+#include "cutlass_patch/epilogue/threadblock/epilogue_with_variadic.h"
+
+namespace cutlass_patch {
+namespace gemm {
+namespace kernel {
+
+template <
+    /// Element type for A matrix operand
+    typename ElementA_,
+    /// Layout type for A matrix operand
+    typename LayoutA_,
+    /// Complex elementwise transformation on A operand
+    cutlass::ComplexTransform TransformA,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentA,
+    /// Element type for B matrix operand
+    typename ElementB_,
+    /// Layout type for B matrix operand
+    typename LayoutB_,
+    /// Complex elementwise transformation on B operand
+    cutlass::ComplexTransform TransformB,
+    /// Access granularity of B matrix in units of elements
+    int kAlignmentB,
+    /// Element type for C and D matrix operands
+    typename ElementC_,
+    /// Layout type for C and D matrix operands
+    typename LayoutC_,
+    /// Element type for internal accumulation
+    typename ElementAccumulator,
+    /// Operator class tag
+    typename OperatorClass,
+    /// Tag indicating architecture to tune for
+    typename ArchTag,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename InstructionShape,
+    /// Epilogue output operator      - must satisfy concept of
+    /// 'EpilogueWithVariadicOp'
+    typename EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle,
+    /// Number of stages used in the pipelined mainloop
+    int Stages,
+    /// Operation performed by GEMM
+    typename Operator,
+    ///
+    typename Enable = void>
+struct DefaultGemmWithVariadic {
+  using GemmBase = typename cutlass::gemm::kernel::DefaultGemmUniversal<
+      ElementA_,
+      LayoutA_,
+      TransformA,
+      kAlignmentA,
+      ElementB_,
+      LayoutB_,
+      TransformB,
+      kAlignmentB,
+      ElementC_,
+      LayoutC_,
+      ElementAccumulator,
+      OperatorClass,
+      ArchTag,
+      ThreadblockShape,
+      WarpShape,
+      InstructionShape,
+      EpilogueOutputOp,
+      ThreadblockSwizzle,
+      Stages,
+      Operator>::GemmKernel;
+
+  // Define epilogue
+  using Epilogue = typename cutlass_patch::epilogue::threadblock::
+      DefaultEpilogueWithVariadicTensorOp<
+          typename GemmBase::Epilogue::Shape,
+          typename GemmBase::Epilogue::WarpMmaOperator,
+          GemmBase::Epilogue::kPartitionsK,
+          ElementC_,
+          EpilogueOutputOp,
+          GemmBase::Epilogue::kElementsPerAccess>::Epilogue;
+
+  // Compose the GEMM kernel
+  using GemmKernel = cutlass::gemm::kernel::
+      GemmUniversal<typename GemmBase::Mma, Epilogue, ThreadblockSwizzle>;
+};
+
+/// Partial specialization: ArchTag = cutlass::arch::Sm70
+///
+///
+template <
+    /// Element type for A matrix operand
+    typename ElementA_,
+    /// Layout type for A matrix operand
+    typename LayoutA_,
+    /// Complex elementwise transformation on A operand
+    cutlass::ComplexTransform TransformA,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentA,
+    /// Element type for B matrix operand
+    typename ElementB_,
+    /// Layout type for B matrix operand
+    typename LayoutB_,
+    /// Complex elementwise transformation on B operand
+    cutlass::ComplexTransform TransformB,
+    /// Access granularity of B matrix in units of elements
+    int kAlignmentB,
+    /// Element type for C and D matrix operands
+    typename ElementC_,
+    /// Layout type for C and D matrix operands
+    typename LayoutC_,
+    /// Element type for internal accumulation
+    typename ElementAccumulator,
+    /// Operator class tag
+    typename OperatorClass,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename InstructionShape,
+    /// Epilogue output operator      - must satisfy concept of
+    /// 'EpilogueWithVariadicOp'
+    typename EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle,
+    /// Number of stages used in the pipelined mainloop
+    int Stages,
+    /// Operation performed by GEMM
+    typename Operator,
+    ///
+    typename Enable>
+struct DefaultGemmWithVariadic<ElementA_,
+                               LayoutA_,
+                               TransformA,
+                               kAlignmentA,
+                               ElementB_,
+                               LayoutB_,
+                               TransformB,
+                               kAlignmentB,
+                               ElementC_,
+                               LayoutC_,
+                               ElementAccumulator,
+                               OperatorClass,
+                               cutlass::arch::Sm70,
+                               ThreadblockShape,
+                               WarpShape,
+                               InstructionShape,
+                               EpilogueOutputOp,
+                               ThreadblockSwizzle,
+                               Stages,
+                               Operator,
+                               Enable> {
+  using GemmBase = typename cutlass::gemm::kernel::DefaultGemmUniversal<
+      ElementA_,
+      LayoutA_,
+      TransformA,
+      kAlignmentA,
+      ElementB_,
+      LayoutB_,
+      TransformB,
+      kAlignmentB,
+      ElementC_,
+      LayoutC_,
+      ElementAccumulator,
+      OperatorClass,
+      cutlass::arch::Sm70,
+      ThreadblockShape,
+      WarpShape,
+      InstructionShape,
+      EpilogueOutputOp,
+      ThreadblockSwizzle,
+      Stages,
+      Operator>::GemmKernel;
+
+  // Define epilogue
+  using Epilogue = typename cutlass_patch::epilogue::threadblock::
+      DefaultEpilogueWithVariadicVoltaTensorOp<
+          typename GemmBase::Epilogue::Shape,
+          typename GemmBase::Epilogue::WarpMmaOperator,
+          GemmBase::Epilogue::kPartitionsK,
+          ElementC_,
+          EpilogueOutputOp,
+          GemmBase::Epilogue::kElementsPerAccess>::Epilogue;
+
+  // Compose the GEMM kernel
+  using GemmKernel = cutlass::gemm::kernel::
+      GemmUniversal<typename GemmBase::Mma, Epilogue, ThreadblockSwizzle>;
+};
+
+}  // namespace kernel
+}  // namespace gemm
+}  // namespace cutlass_patch

python/paddle/apy/matmul_pass/matmul/cutlass_patch/hip/all_tuning_configs.h

@@ -0,0 +1,552 @@
+// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "hytlass/gemm_coord.h"
+
+namespace ap {
+
+constexpr int kNumConfigsHalf = 28;
+constexpr int kNumConfigsFloat = 13;
+
+template <int SwizzleFactor, bool Batched>
+struct SwizzleWrapper {
+  using Type =
+      hytlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<SwizzleFactor>;
+};
+
+#define AP_AUTOTUNE(func, stream_ptr, count, ...)                             \
+  {                                                                           \
+    using FuncType = decltype(func<0>);                                       \
+    static int selected_config_id = -1;                                       \
+    static std::vector<std::function<FuncType>> matmul_functions =            \
+        []<std::size_t... Is>(std::index_sequence<Is...>) {                   \
+      return std::vector<std::function<FuncType>>{func<Is>...};               \
+    }                                                                         \
+    (std::make_index_sequence<count>());                                      \
+    if (selected_config_id == -1) {                                           \
+      selected_config_id =                                                    \
+          ap::ProfileBestConfig(matmul_functions, stream_ptr, ##__VA_ARGS__); \
+    }                                                                         \
+    matmul_functions[selected_config_id](__VA_ARGS__);                        \
+  }
+
+#define AP_AUTOTUNE_half(func, stream_ptr, ...) \
+  AP_AUTOTUNE(func, stream_ptr, ap::kNumConfigsHalf, __VA_ARGS__)
+#define AP_AUTOTUNE_float(func, stream_ptr, ...) \
+  AP_AUTOTUNE(func, stream_ptr, ap::kNumConfigsFloat, __VA_ARGS__)
+#define AP_AUTOTUNE_bfloat16(func, stream_ptr, ...) \
+  AP_AUTOTUNE_half(func, stream_ptr, __VA_ARGS__)
+
+template <typename ElementT, int SwizzleFactor, bool Batched, int Id = 0>
+struct GemmTuningConfigs {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 2;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = Id;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 1> {
+  using TShape = hytlass::gemm::GemmShape<64, 128, 64>;
+  using WShape = hytlass::gemm::GemmShape<32, 64, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 1;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 2> {
+  using TShape = hytlass::gemm::GemmShape<64, 128, 64>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 2;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 3> {
+  using TShape = hytlass::gemm::GemmShape<128, 64, 64>;
+  using WShape = hytlass::gemm::GemmShape<64, 32, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 3;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 4> {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 4;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 5> {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 64>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 5;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 6> {
+  using TShape = hytlass::gemm::GemmShape<256, 64, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 6;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 7> {
+  using TShape = hytlass::gemm::GemmShape<256, 64, 64>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 7;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 8> {
+  using TShape = hytlass::gemm::GemmShape<256, 128, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 8;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 9> {
+  using TShape = hytlass::gemm::GemmShape<256, 128, 64>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 9;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 10> {
+  using TShape = hytlass::gemm::GemmShape<128, 32, 64>;
+  using WShape = hytlass::gemm::GemmShape<32, 32, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 10;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 11> {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 11;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 12> {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 64>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 12;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 13> {
+  using TShape = hytlass::gemm::GemmShape<256, 64, 64>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 13;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 14> {
+  using TShape = hytlass::gemm::GemmShape<256, 64, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 14;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 15> {
+  using TShape = hytlass::gemm::GemmShape<32, 64, 64>;
+  using WShape = hytlass::gemm::GemmShape<16, 32, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 5;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 15;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 16> {
+  using TShape = hytlass::gemm::GemmShape<64, 64, 64>;
+  using WShape = hytlass::gemm::GemmShape<32, 32, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 5;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 16;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 17> {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 5;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 17;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 18> {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 64>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 5;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 18;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 19> {
+  using TShape = hytlass::gemm::GemmShape<64, 128, 32>;
+  using WShape = hytlass::gemm::GemmShape<32, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 6;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 19;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 20> {
+  using TShape = hytlass::gemm::GemmShape<128, 64, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 32, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 6;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 20;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 21> {
+  using TShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using WShape = hytlass::gemm::GemmShape<32, 32, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 10;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 21;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 22> {
+  using TShape = hytlass::gemm::GemmShape<128, 256, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 2;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 22;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 23> {
+  using TShape = hytlass::gemm::GemmShape<128, 256, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 23;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 24> {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 32>;
+  using WShape = hytlass::gemm::GemmShape<32, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 24;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 25> {
+  using TShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using WShape = hytlass::gemm::GemmShape<32, 32, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 25;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 26> {
+  using TShape = hytlass::gemm::GemmShape<64, 128, 64>;
+  using WShape = hytlass::gemm::GemmShape<32, 32, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 26;
+};
+
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, 27> {
+  using TShape = hytlass::gemm::GemmShape<128, 64, 64>;
+  using WShape = hytlass::gemm::GemmShape<64, 32, 64>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 16>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 27;
+};
+
+// Specialization for float
+template <int SwizzleFactor, bool Batched, int Id>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, Id> {
+  using TShape = hytlass::gemm::GemmShape<64, 64, 16>;
+  using WShape = hytlass::gemm::GemmShape<32, 32, 16>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = Id;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 1> {
+  using TShape = hytlass::gemm::GemmShape<64, 64, 32>;
+  using WShape = hytlass::gemm::GemmShape<32, 32, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 1;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 2> {
+  using TShape = hytlass::gemm::GemmShape<64, 128, 32>;
+  using WShape = hytlass::gemm::GemmShape<32, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 2;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 3> {
+  using TShape = hytlass::gemm::GemmShape<64, 256, 16>;
+  using WShape = hytlass::gemm::GemmShape<32, 64, 16>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 3;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 4> {
+  using TShape = hytlass::gemm::GemmShape<64, 256, 32>;
+  using WShape = hytlass::gemm::GemmShape<32, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 4;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 5> {
+  using TShape = hytlass::gemm::GemmShape<128, 64, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 32, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 5;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 6> {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 16>;
+  using WShape = hytlass::gemm::GemmShape<32, 64, 16>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 6;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 7> {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 32>;
+  using WShape = hytlass::gemm::GemmShape<32, 64, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 7;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 8> {
+  using TShape = hytlass::gemm::GemmShape<256, 64, 16>;
+  using WShape = hytlass::gemm::GemmShape<64, 32, 16>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 8;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 9> {
+  using TShape = hytlass::gemm::GemmShape<256, 64, 32>;
+  using WShape = hytlass::gemm::GemmShape<64, 32, 32>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 3;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 9;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 10> {
+  using TShape = hytlass::gemm::GemmShape<64, 128, 16>;
+  using WShape = hytlass::gemm::GemmShape<32, 64, 16>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 10;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 11> {
+  using TShape = hytlass::gemm::GemmShape<128, 64, 16>;
+  using WShape = hytlass::gemm::GemmShape<64, 32, 16>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 11;
+};
+
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, 12> {
+  using TShape = hytlass::gemm::GemmShape<128, 128, 16>;
+  using WShape = hytlass::gemm::GemmShape<32, 64, 16>;
+  using IShape = hytlass::gemm::GemmShape<16, 16, 8>;
+  static constexpr int kNumStages = 4;
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = 12;
+};
+
+struct DefaultConfig {
+  static constexpr int kConfigId = 0;
+  static constexpr int kSwizzleFactor = 1;
+  static constexpr bool kBatched = false;
+};
+
+}  // namespace ap

python/paddle/apy/matmul_pass/matmul/cutlass_patch/hip/hytlass_matmul.h

@@ -0,0 +1,254 @@
+// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "hytlass/epilogue/thread/linear_combination_bias_elementwise.h"
+#include "hytlass/gemm/device/gemm_universal.h"
+#include "hytlass/gemm/device/gemm_universal_with_broadcast.h"
+#include "hytlass/util/device_memory.h"
+
+#include "cutlass_patch/batched_matrix_coord.h"
+#include "cutlass_patch/epilogue/thread/linear_combination_unary.h"
+#include "cutlass_patch/epilogue/thread/linear_combination_variadic.h"
+#include "cutlass_patch/gemm/device/gemm_universal_with_variadic.h"
+#include "cutlass_patch/hip/all_tuning_configs.h"
+
+#include "params.h"  // NOLINT
+
+#define CHECK_HYTLASS(status)                               \
+  {                                                         \
+    auto error = status;                                    \
+    if (error != hytlass::Status::kSuccess) {               \
+      std::cerr << "HYTLASS error = " << int(error) << " (" \
+                << hytlassGetStatusString(error) << ")"     \
+                << " at line " << __LINE__ << std::endl;    \
+      std::abort();                                         \
+    }                                                       \
+  }
+
+namespace ap {
+
+using MatrixCoord = cutlass_patch::BatchedMatrixCoord;
+using bfloat16 = __hip_bfloat16;
+
+// Operation performed by GEMM
+template <typename ElementT>
+struct GemmOperation {
+  using Type = hytlass::arch::OpMultiplyAdd;
+};
+
+template <>
+struct GemmOperation<float> {
+  using Type = hytlass::arch::OpMultiplyAddFastF32;
+};
+
+static hytlass::gemm::GemmUniversalMode GetGemmMode(int batch_count) {
+  return batch_count > 1 ? hytlass::gemm::GemmUniversalMode::kBatched
+                         : hytlass::gemm::GemmUniversalMode::kGemm;
+}
+
+static void *GetWorkspace(size_t workspace_size) {
+  static hytlass::device_memory::allocation<uint8_t> workspace;
+  if (workspace.size() < workspace_size) {
+    workspace.reset(workspace_size);
+  }
+  return workspace.get();
+}
+
+template <typename GemmFunc>
+hytlass::Status SetMaxDynamicSharedMemorySize() {
+  hipError_t hiprt_result;
+
+  // If requires more than 48KB: configure for extended, dynamic shared memory
+  if constexpr (GemmFunc::kSharedStorageSize >= (48 << 10)) {
+    hiprt_result = hipFuncSetAttribute(
+        (const void *)hytlass::Kernel2<typename GemmFunc::GemmKernel>,
+        hipFuncAttributeMaxDynamicSharedMemorySize,
+        GemmFunc::kSharedStorageSize);
+    if (hiprt_result != hipSuccess) {
+      HYTLASS_TRACE_HOST("hipFuncSetAttribute() returned error "
+                         << hipGetErrorString(hiprt_result));
+      return hytlass::Status::kErrorInternal;
+    }
+  }
+
+#if AP_ENABLE_DEBUG
+  // Update SM occupancy member
+  int sm_occupancy = -1;
+  hiprt_result = hipOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(
+      &sm_occupancy,
+      hytlass::Kernel2<typename GemmFunc::GemmKernel>,
+      GemmFunc::GemmKernel::kThreadCount,
+      GemmFunc::kSharedStorageSize,
+      hipOccupancyDisableCachingOverride);
+  if (hiprt_result != hipSuccess) {
+    HYTLASS_TRACE_HOST(
+        "hipOccupancyMaxActiveBlocksPerMultiprocessorWithFlags() returned "
+        "error "
+        << hipGetErrorString(hiprt_result));
+    return hytlass::Status::kErrorInternal;
+  }
+  HYTLASS_TRACE_HOST("sm_occupancy: (" << sm_occupancy
+                                       << ") "
+                                          "smem_size: ("
+                                       << GemmFunc::kSharedStorageSize
+                                       << ") "
+                                          "GemmKernel::kThreadCount: ("
+                                       << GemmFunc::GemmKernel::kThreadCount
+                                       << ")");
+#endif
+  return hytlass::Status::kSuccess;
+}
+
+// Convert HIP data type to hytlass data type
+template <typename T>
+struct HytlassDataType {
+  using Type = T;
+};
+
+template <>
+struct HytlassDataType<half> {
+  using Type = hytlass::half_t;
+};
+
+template <>
+struct HytlassDataType<__hip_bfloat16> {
+  using Type = hytlass::bfloat16_t;
+};
+
+// Convert to hytlass layout
+template <bool Transposed>
+struct MatrixLayout {
+  using Type = hytlass::layout::RowMajor;
+};
+
+template <>
+struct MatrixLayout<true> {
+  using Type = hytlass::layout::ColumnMajor;
+};
+
+template <typename ElementT,
+          typename ElementComputeT,
+          template <typename T>
+          class VariadicFunctor,
+          int AlignA = 128 / hytlass::sizeof_bits<ElementT>::value,
+          int AlignB = 128 / hytlass::sizeof_bits<ElementT>::value,
+          int ConfigId = DefaultConfig::kConfigId,
+          int SwizzleFactor = DefaultConfig::kSwizzleFactor,
+          bool Batched = DefaultConfig::kBatched>
+void MatmulAddVariadic(
+    const GemmEpilogueParams &params,
+    const typename VariadicFunctor<ElementComputeT>::Arguments &variadic_args) {
+  // <- data type of accumulator
+  using ElementAccumulator = typename HytlassDataType<ElementComputeT>::Type;
+  // <- data type of epilogue operations
+  using ElementComputeEpilogue = ElementAccumulator;
+  // <- data type of elements in input matrix A
+  using ElementInputA = typename HytlassDataType<ElementT>::Type;
+  // <- data type of elements in input matrix B
+  using ElementInputB = typename HytlassDataType<ElementT>::Type;
+  // <- data type of elements in output matrix D
+  using ElementOutput = typename HytlassDataType<ElementT>::Type;
+
+  constexpr int AlignC = AlignB;
+
+  // Epilogue operation as LinearCombination:
+  //  alpha * accumulator + beta * source
+  using EpilogueOutputOp =
+      cutlass_patch::epilogue::thread::LinearCombinationVariadic<
+          VariadicFunctor,
+          ElementOutput,
+          AlignC,
+          ElementAccumulator,
+          ElementComputeEpilogue,
+          hytlass::epilogue::thread::ScaleType::NoBetaScaling>;
+
+  using GemmFunc = cutlass_patch::gemm::device::GemmUniversalWithVariadic<
+      ElementInputA,
+      hytlass::layout::RowMajor,
+      ElementInputB,
+      hytlass::layout::RowMajor,
+      ElementOutput,
+      hytlass::layout::RowMajor,
+      ElementAccumulator,
+      hytlass::arch::OpClassTensorOp,
+      hytlass::arch::Gfx928,
+      typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
+          TShape,
+      typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
+          WShape,
+      typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
+          IShape,
+      EpilogueOutputOp,
+      typename GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::
+          SwizzleThreadBlock,
+      GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ConfigId>::kNumStages,
+      AlignA,
+      AlignB,
+      typename GemmOperation<ElementT>::Type>;
+
+  CHECK_HYTLASS(SetMaxDynamicSharedMemorySize<GemmFunc>());
+
+  /// Arguments
+  hytlass::gemm::GemmCoord problem_size{params.m, params.n, params.k};
+
+  const ElementInputA *input =
+      reinterpret_cast<const ElementInputA *>(params.input);
+  const ElementInputB *weight =
+      reinterpret_cast<const ElementInputB *>(params.weight);
+  const ElementOutput *bias =
+      reinterpret_cast<const ElementOutput *>(params.bias);
+  ElementOutput *output = reinterpret_cast<ElementOutput *>(params.output);
+
+  ElementComputeEpilogue alpha = static_cast<ElementComputeEpilogue>(1);
+  ElementComputeEpilogue beta = bias ? static_cast<ElementComputeEpilogue>(1)
+                                     : static_cast<ElementComputeEpilogue>(0);
+
+  typename GemmFunc::Arguments arguments{
+      GetGemmMode(params.batch_count),
+      problem_size,        // <- problem size of matrix multiplication
+      params.batch_count,  // <- batch_count or k-dimension split factor
+      {alpha, beta, variadic_args},  // <- epilogue params, alpha, beta
+      input,                         // <- input, ptr_A, A, shape={M, K}
+      weight,                        // <- input, ptr_B, B, shape={K, N}
+      bias,                          // <- input, ptr_C, shape={M, N} or {1, N}
+      output,                        // <- output, ptr_D, Z, shape={M, N}
+      params.shape_args.batch_stride_A,
+      params.shape_args.batch_stride_B,
+      params.shape_args.batch_stride_C,
+      params.shape_args.batch_stride_D,
+      params.shape_args.lda,
+      params.shape_args.ldb,
+      params.shape_args.ldc_bias,
+      params.shape_args.ldd};
+
+  size_t workspace_size = GemmFunc::get_workspace_size(arguments);
+  void *workspace = workspace_size > 0 ? GetWorkspace(workspace_size) : nullptr;
+
+  GemmFunc device_gemm;
+
+  hipStream_t *stream_ptr = reinterpret_cast<hipStream_t *>(params.stream_ptr);
+
+  CHECK_HYTLASS(device_gemm.can_implement(arguments));
+  CHECK_HYTLASS(device_gemm.initialize(arguments, workspace, *stream_ptr));
+
+  // Run the GEMM
+  CHECK_HYTLASS(device_gemm(*stream_ptr));
+#if AP_ENABLE_DEBUG
+  CHECK_HIP(hipStreamSynchronize(*stream_ptr));
+#endif
+}
+
+}  // namespace ap

python/paddle/apy/matmul_pass/matmul/cutlass_patch/trace_device.h

@@ -0,0 +1,73 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#if CUTLASS_DEBUG_TRACE_LEVEL
+
+#ifndef CUTLASS_TRACE_DEVICE
+#define CUTLASS_TRACE_DEVICE(format, ...)                        \
+  {                                                              \
+    if (blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0 && \
+        threadIdx.x == 0 && threadIdx.y == 0) {                  \
+      printf("[DEVICE][%s:%d, %s]" format "\n",                  \
+             __FILE__,                                           \
+             __LINE__,                                           \
+             __FUNCTION__,                                       \
+             ##__VA_ARGS__);                                     \
+    }                                                            \
+  }
+#endif
+
+#ifndef CUTLASS_TRACE_DEVICE_TID_DETAIL
+#define CUTLASS_TRACE_DEVICE_TID_DETAIL(bidz, bidx, tidx, format, ...)    \
+  {                                                                       \
+    if (blockIdx.x == bidx && blockIdx.y == 0 && blockIdx.z == bidz &&    \
+        threadIdx.x == tidx && threadIdx.y == 0) {                        \
+      printf("[DEVICE][%s:%d, %s][bid={%d,%d,%d}, tid={%d,%d,%d}]" format \
+             "\n",                                                        \
+             __FILE__,                                                    \
+             __LINE__,                                                    \
+             __FUNCTION__,                                                \
+             blockIdx.x,                                                  \
+             blockIdx.y,                                                  \
+             blockIdx.z,                                                  \
+             threadIdx.x,                                                 \
+             threadIdx.y,                                                 \
+             threadIdx.z,                                                 \
+             ##__VA_ARGS__);                                              \
+    }                                                                     \
+  }
+#endif
+
+#ifndef CUTLASS_TRACE_DEVICE_TID
+#define CUTLASS_TRACE_DEVICE_TID(format, ...)                       \
+  {                                                                 \
+    CUTLASS_TRACE_DEVICE_TID_DETAIL(0, 0, 0, format, ##__VA_ARGS__) \
+    CUTLASS_TRACE_DEVICE_TID_DETAIL(0, 0, 1, format, ##__VA_ARGS__) \
+    CUTLASS_TRACE_DEVICE_TID_DETAIL(0, 1, 0, format, ##__VA_ARGS__) \
+  }
+#endif
+
+#else
+
+#ifndef CUTLASS_TRACE_DEVICE
+#define CUTLASS_TRACE_DEVICE(format, ...)
+#endif
+
+#ifndef CUTLASS_TRACE_DEVICE_TID
+#define CUTLASS_TRACE_DEVICE_TID(format, ...)
+#endif
+
+#endif

python/paddle/apy/matmul_pass/matmul/generate_configs.py

@@ -0,0 +1,298 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+head_template = """// auto-generated by generate_configs.py
+
+#pragma once
+
+#include "cutlass/gemm_coord.h"
+
+namespace ap {
+
+constexpr int kNumConfigsHalf = ${num_configs_fp16};
+constexpr int kNumConfigsFloat = ${num_configs_fp32};
+
+template <int SwizzleFactor, bool Batched> struct SwizzleWrapper {
+  using Type =
+      cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<SwizzleFactor>;
+};
+
+// template <int SwizzleFactor>
+// struct SwizzleWrapper<SwizzleFactor, true> {
+//   using Type =
+//       cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle;
+// };
+"""
+
+autotune_wrapper_template = """
+#define AP_AUTOTUNE_${datatype}(func, stream, ...) { \\
+  using FuncType = decltype(func<0>); \\
+  static int selected_config_id = -1; \\
+  static std::vector<std::function<FuncType>> \\
+      matmul_functions = { \\
+          ${repeat_functions} \\
+          };  \\
+  if (selected_config_id == -1) { \\
+    selected_config_id = ap::ProfileBestConfig(matmul_functions, stream, ##__VA_ARGS__); \\
+  } \\
+  matmul_functions[selected_config_id](__VA_ARGS__); \\
+}
+"""
+
+fp16_config_template_0 = """
+template <typename ElementT, int SwizzleFactor, bool Batched, int Id = 0>
+struct GemmTuningConfigs {
+  using TShape = cutlass::gemm::GemmShape<${tshape}>;
+  using WShape = cutlass::gemm::GemmShape<${wshape}>;
+  using IShape = cutlass::gemm::GemmShape<${ishape}>;
+  static constexpr int kNumStages = ${stages};
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = Id;
+};
+"""
+
+fp16_config_template = """
+template <typename ElementT, int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<ElementT, SwizzleFactor, Batched, ${config_id}> {
+  using TShape = cutlass::gemm::GemmShape<${tshape}>;
+  using WShape = cutlass::gemm::GemmShape<${wshape}>;
+  using IShape = cutlass::gemm::GemmShape<${ishape}>;
+  static constexpr int kNumStages = ${stages};
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = ${config_id};
+};
+"""
+
+fp32_config_template_0 = """
+// Specialization for float
+template <int SwizzleFactor, bool Batched, int Id>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, Id> {
+  using TShape = cutlass::gemm::GemmShape<${tshape}>;
+  using WShape = cutlass::gemm::GemmShape<${wshape}>;
+  using IShape = cutlass::gemm::GemmShape<${ishape}>;
+  static constexpr int kNumStages = ${stages};
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = Id;
+};
+"""
+
+fp32_config_template = """
+template <int SwizzleFactor, bool Batched>
+struct GemmTuningConfigs<float, SwizzleFactor, Batched, ${config_id}> {
+  using TShape = cutlass::gemm::GemmShape<${tshape}>;
+  using WShape = cutlass::gemm::GemmShape<${wshape}>;
+  using IShape = cutlass::gemm::GemmShape<${ishape}>;
+  static constexpr int kNumStages = ${stages};
+
+  using SwizzleThreadBlock =
+      typename SwizzleWrapper<SwizzleFactor, Batched>::Type;
+  static constexpr int kId = ${config_id};
+};
+"""
+
+tail_code_str = """
+} // namespace ap
+"""
+
+
+class GemmTuningConfig:
+    def __init__(self, tshape, wshape, ishape, stages, level):
+        self.tshape = tshape
+        self.wshape = wshape
+        self.ishape = ishape
+        self.stages = stages
+        self.level = level
+
+    def __eq__(self, other):
+        def check_shape(s1, s2):
+            assert len(s1) == len(s2)
+            res = True
+            for i in range(len(s1)):
+                if s1[i] != s2[i]:
+                    res = False
+                    break
+            return res
+
+        res = check_shape(self.tshape, other.tshape)
+        res = res and check_shape(self.wshape, other.wshape)
+        res = res and check_shape(self.ishape, other.ishape)
+        res = res and self.stages == other.stages
+        return res
+
+
+def all_configs_sm80_fp16():
+    all_tuning_configs_fp16 = [
+        GemmTuningConfig([16, 128, 64], [16, 32, 64], [16, 8, 16], 2, 2),
+        GemmTuningConfig([32, 128, 64], [32, 32, 64], [16, 8, 16], 2, 2),
+        GemmTuningConfig([64, 128, 64], [64, 64, 64], [16, 8, 16], 2, 2),
+        GemmTuningConfig([128, 128, 64], [64, 64, 64], [16, 8, 16], 2, 2),
+        GemmTuningConfig([128, 128, 64], [128, 32, 64], [16, 8, 16], 2, 2),
+        GemmTuningConfig([128, 256, 64], [64, 64, 64], [16, 8, 16], 2, 2),
+        GemmTuningConfig([256, 128, 64], [64, 64, 64], [16, 8, 16], 2, 1),
+        GemmTuningConfig([16, 128, 64], [16, 32, 64], [16, 8, 16], 3, 3),
+        GemmTuningConfig([32, 128, 64], [32, 32, 64], [16, 8, 16], 3, 2),
+        GemmTuningConfig([64, 128, 64], [32, 64, 64], [16, 8, 16], 3, 1),
+        GemmTuningConfig([64, 128, 64], [64, 64, 64], [16, 8, 16], 3, 1),
+        GemmTuningConfig([64, 256, 64], [64, 64, 64], [16, 8, 16], 3, 3),
+        GemmTuningConfig([128, 64, 64], [64, 32, 64], [16, 8, 16], 3, 1),
+        GemmTuningConfig([128, 128, 32], [64, 64, 32], [16, 8, 16], 3, 1),
+        GemmTuningConfig([128, 128, 64], [64, 64, 64], [16, 8, 16], 3, 1),
+        GemmTuningConfig([128, 128, 64], [128, 32, 64], [16, 8, 16], 3, 2),
+        GemmTuningConfig([128, 256, 32], [64, 64, 32], [16, 8, 16], 3, 2),
+        GemmTuningConfig([128, 256, 64], [64, 64, 64], [16, 8, 16], 3, 2),
+        GemmTuningConfig([256, 64, 32], [64, 64, 32], [16, 8, 16], 3, 1),
+        GemmTuningConfig([256, 64, 64], [64, 64, 64], [16, 8, 16], 3, 1),
+        GemmTuningConfig([256, 128, 32], [64, 64, 32], [16, 8, 16], 3, 1),
+        GemmTuningConfig([256, 128, 64], [64, 64, 64], [16, 8, 16], 3, 1),
+        GemmTuningConfig([16, 128, 64], [16, 32, 64], [16, 8, 16], 4, 3),
+        GemmTuningConfig([32, 128, 64], [32, 32, 64], [16, 8, 16], 4, 2),
+        GemmTuningConfig([64, 128, 64], [64, 64, 64], [16, 8, 16], 4, 2),
+        GemmTuningConfig([64, 256, 32], [64, 64, 32], [16, 8, 16], 4, 2),
+        GemmTuningConfig([64, 256, 64], [64, 64, 64], [16, 8, 16], 4, 3),
+        GemmTuningConfig([128, 32, 64], [32, 32, 64], [16, 8, 16], 4, 1),
+        GemmTuningConfig([128, 128, 32], [64, 64, 32], [16, 8, 16], 4, 1),
+        GemmTuningConfig([128, 128, 64], [64, 64, 64], [16, 8, 16], 4, 1),
+        GemmTuningConfig([128, 128, 64], [128, 32, 64], [16, 8, 16], 4, 2),
+        GemmTuningConfig([256, 64, 64], [64, 64, 64], [16, 8, 16], 4, 1),
+        GemmTuningConfig([256, 64, 32], [64, 64, 32], [16, 8, 16], 4, 1),
+        GemmTuningConfig([16, 64, 64], [16, 32, 64], [16, 8, 16], 5, 2),
+        GemmTuningConfig([16, 128, 64], [16, 32, 64], [16, 8, 16], 5, 3),
+        GemmTuningConfig([32, 64, 64], [16, 32, 64], [16, 8, 16], 5, 1),
+        GemmTuningConfig([32, 128, 64], [32, 32, 64], [16, 8, 16], 5, 3),
+        GemmTuningConfig([64, 64, 64], [32, 32, 64], [16, 8, 16], 5, 1),
+        GemmTuningConfig([64, 128, 64], [64, 64, 64], [16, 8, 16], 5, 3),
+        GemmTuningConfig([128, 128, 32], [64, 64, 32], [16, 8, 16], 5, 1),
+        GemmTuningConfig([128, 128, 64], [64, 64, 64], [16, 8, 16], 5, 1),
+        GemmTuningConfig([128, 128, 64], [128, 32, 64], [16, 8, 16], 5, 2),
+        GemmTuningConfig([64, 128, 32], [32, 64, 32], [16, 8, 16], 6, 1),
+        GemmTuningConfig([128, 64, 32], [64, 32, 32], [16, 8, 16], 6, 1),
+        GemmTuningConfig([128, 32, 32], [32, 32, 32], [16, 8, 16], 7, 1),
+        GemmTuningConfig([64, 64, 32], [32, 32, 32], [16, 8, 16], 10, 1),
+    ]
+    return all_tuning_configs_fp16
+
+
+def all_configs_sm80_fp32():
+    all_tuning_configs_fp32 = [
+        GemmTuningConfig([64, 64, 16], [32, 32, 16], [16, 8, 8], 3, 1),
+        GemmTuningConfig([64, 64, 32], [32, 32, 32], [16, 8, 8], 3, 1),
+        GemmTuningConfig([64, 128, 32], [32, 64, 32], [16, 8, 8], 3, 1),
+        GemmTuningConfig([64, 256, 16], [32, 64, 16], [16, 8, 8], 3, 1),
+        GemmTuningConfig([64, 256, 32], [32, 64, 32], [16, 8, 8], 3, 1),
+        GemmTuningConfig([128, 64, 32], [64, 32, 32], [16, 8, 8], 3, 1),
+        GemmTuningConfig([128, 128, 16], [32, 64, 16], [16, 8, 8], 3, 1),
+        GemmTuningConfig([128, 128, 32], [32, 64, 32], [16, 8, 8], 3, 1),
+        GemmTuningConfig([256, 64, 16], [64, 32, 16], [16, 8, 8], 3, 1),
+        GemmTuningConfig([256, 64, 32], [64, 32, 32], [16, 8, 8], 3, 1),
+        GemmTuningConfig([64, 128, 16], [32, 64, 16], [16, 8, 8], 4, 1),
+        GemmTuningConfig([128, 64, 16], [64, 32, 16], [16, 8, 8], 4, 1),
+        GemmTuningConfig([128, 128, 16], [32, 64, 16], [16, 8, 8], 4, 1),
+    ]
+    return all_tuning_configs_fp32
+
+
+def generate_autotune_wrapper(datatype, num_configs):
+    repeat_func_strs = []
+    for i in range(num_configs):
+        repeat_func_strs.append(f"func<{i}>")
+    code_str = autotune_wrapper_template.replace(
+        "${datatype}", datatype
+    ).replace("${repeat_functions}", ", \\\n          ".join(repeat_func_strs))
+    return code_str
+
+
+def get_configs(all_configs_list, level=3):
+    consigs_list = []
+    for i in range(len(all_configs_list)):
+        if all_configs_list[i].level <= level:
+            already_have = False
+            for config in consigs_list:
+                if config == all_configs_list[i]:
+                    print(f"-- The {i}-th config is repeat.")
+                    already_have = True
+                    break
+            if not already_have:
+                consigs_list.append(all_configs_list[i])
+    return consigs_list
+
+
+def generate_configs(configs_list, config_template_0, config_template):
+    code_str = ""
+    config_id = 0
+    for config in configs_list:
+        if config_id == 0:
+            config_code_str = (
+                config_template_0.replace(
+                    "${tshape}", ", ".join(map(str, config.tshape))
+                )
+                .replace("${wshape}", ", ".join(map(str, config.wshape)))
+                .replace("${ishape}", ", ".join(map(str, config.ishape)))
+                .replace("${stages}", str(config.stages))
+            )
+        else:
+            config_code_str = (
+                config_template.replace(
+                    "${tshape}", ", ".join(map(str, config.tshape))
+                )
+                .replace("${wshape}", ", ".join(map(str, config.wshape)))
+                .replace("${ishape}", ", ".join(map(str, config.ishape)))
+                .replace("${stages}", str(config.stages))
+                .replace("${config_id}", str(config_id))
+            )
+        code_str += config_code_str
+        config_id += 1
+    return len(configs_list), code_str
+
+
+def main():
+    level = 1
+    num_fp16_configs, fp16_configs_code_str = generate_configs(
+        configs_list=get_configs(all_configs_sm80_fp16(), level=level),
+        config_template_0=fp16_config_template_0,
+        config_template=fp16_config_template,
+    )
+    num_fp32_configs, fp32_configs_code_str = generate_configs(
+        configs_list=get_configs(all_configs_sm80_fp32(), level=level),
+        config_template_0=fp32_config_template_0,
+        config_template=fp32_config_template,
+    )
+    print(
+        f"-- Total {num_fp16_configs} fp16 configs, {num_fp32_configs} fp32 configs"
+    )
+    head_code_str = head_template.replace(
+        "${num_configs_fp16}", str(num_fp16_configs)
+    ).replace("${num_configs_fp32}", str(num_fp32_configs))
+    fp16_autotune_wrapper_code_str = generate_autotune_wrapper(
+        "half", num_fp16_configs
+    )
+    fp32_autotune_wrapper_code_str = generate_autotune_wrapper(
+        "float", num_fp32_configs
+    )
+    with open("all_tuning_configs.h", "w") as f:
+        f.write(head_code_str)
+        f.write(fp16_autotune_wrapper_code_str)
+        f.write(fp32_autotune_wrapper_code_str)
+        f.write(fp16_configs_code_str)
+        f.write(fp32_configs_code_str)
+        f.write(tail_code_str)
+
+
+if __name__ == "__main__":
+    main()

python/paddle/apy/matmul_pass/matmul/math_function.h

@@ -0,0 +1,39 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#ifdef __NVCC__
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#endif
+#ifdef __HIPCC__
+#include <hip/hip_bfloat16.h>
+#include <hip/hip_fp16.h>
+#include <hip/hip_runtime.h>
+#endif
+
+namespace ap {
+
+template <typename T>
+__forceinline__ __host__ __device__ T ComputePow(T base, T exponent) {
+  T res = (exponent == static_cast<T>(2))
+              ? (base * base)
+              : ((exponent == static_cast<T>(3)) ? (base * base * base)
+                                                 : (powf(base, exponent)));
+  return res;
+}
+
+}  // namespace ap

python/paddle/apy/matmul_pass/matmul/matmul.h

@@ -0,0 +1,63 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <iostream>
+#include <map>
+#include <vector>
+
+#ifdef __NVCC__
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+
+#define CHECK_CUDA(func)                                                      \
+  {                                                                           \
+    cudaError_t err = func;                                                   \
+    if (err != cudaSuccess) {                                                 \
+      std::cerr << "[" << __FILE__ << ":" << __LINE__ << ", " << __FUNCTION__ \
+                << "] "                                                       \
+                << "CUDA error(" << err << "), " << cudaGetErrorString(err)   \
+                << " when call " << #func << std::endl;                       \
+      exit(EXIT_FAILURE);                                                     \
+    }                                                                         \
+  }
+
+#include "cutlass_patch/cuda/cutlass_matmul.cuh"  // NOLINT
+#include "math_function.h"                        // NOLINT
+#include "profile.h"                              // NOLINT
+#endif
+
+#ifdef __HIPCC__
+#include <hip/hip_bfloat16.h>
+#include <hip/hip_fp16.h>
+#include <hip/hip_runtime.h>
+
+#define CHECK_HIP(func)                                                       \
+  {                                                                           \
+    hipError_t err = func;                                                    \
+    if (err != hipSuccess) {                                                  \
+      std::cerr << "[" << __FILE__ << ":" << __LINE__ << ", " << __FUNCTION__ \
+                << "] "                                                       \
+                << "HIP error(" << err << "), " << hipGetErrorString(err)     \
+                << " when call " << #func << std::endl;                       \
+      exit(EXIT_FAILURE);                                                     \
+    }                                                                         \
+  }
+
+#include "cutlass_patch/hip/hytlass_matmul.h"  // NOLINT
+#include "math_function.h"                     // NOLINT
+#include "profile.h"                           // NOLINT
+#endif

python/paddle/apy/matmul_pass/matmul/params.h

@@ -0,0 +1,194 @@
+// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <cstdint>
+#include <iostream>
+#include <limits>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+#define ASSERT_CHECK(__cond)                                            \
+  do {                                                                  \
+    const bool __cond_var = (__cond);                                   \
+    if (!__cond_var) {                                                  \
+      ::std::string __err_msg = ::std::string("`") + #__cond +          \
+                                "` check failed at " + __FILE__ + ":" + \
+                                ::std::to_string(__LINE__);             \
+      throw std::runtime_error(__err_msg);                              \
+    }                                                                   \
+  } while (0)
+
+namespace ap {
+
+inline int CheckedCastToInt(int64_t value) {
+  ASSERT_CHECK(value >= 0);
+  ASSERT_CHECK(value <= static_cast<int64_t>(std::numeric_limits<int>::max()));
+  return static_cast<int>(value);
+}
+
+template <typename T, int Dim>
+struct Alignment {
+  static constexpr int kValue =
+      ((Dim % 8) == 0) ? 8
+                       : (((Dim % 4) == 0) ? 4 : (((Dim % 2) == 0) ? 2 : 1));
+};
+
+template <int Dim>
+struct Alignment<float, Dim> {
+  static constexpr int kValue =
+      ((Dim % 4) == 0) ? 4 : (((Dim % 2) == 0) ? 2 : 1);
+};
+
+struct GemmEpilogueParams {
+  int batch_count;
+  int m;
+  int n;
+  int k;
+
+  bool transpose_a;
+  bool transpose_b;
+
+  // Shape related aruguments
+  struct ShapeArguments {
+    int64_t batch_stride_A;
+    int64_t batch_stride_B;
+    int64_t batch_stride_C;
+    int64_t batch_stride_D;
+    int64_t lda;
+    int64_t ldb;
+    int64_t ldc_bias;
+    int64_t ldd;
+  };
+
+  ShapeArguments shape_args;
+
+  const void *input;
+  const void *weight;
+  const void *bias;
+  void *output;
+
+  void *stream_ptr;
+
+  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(void *stream_ptr,
+                     const void *input,
+                     const void *weight,
+                     const void *bias,
+                     void *output,
+                     const std::vector<int64_t> &input_shape,
+                     const std::vector<int64_t> &weight_shape,
+                     const std::vector<int64_t> &bias_shape,
+                     bool transpose_a = false,
+                     bool transpose_b = false)
+      : stream_ptr(stream_ptr),
+        input(input),
+        weight(weight),
+        bias(bias),
+        output(output),
+        transpose_a(transpose_a),
+        transpose_b(transpose_b) {
+    ASSERT_CHECK(input_shape.size() >= 2U);
+    ASSERT_CHECK(weight_shape.size() >= 2U);
+
+    input0_shape = input_shape;
+    input1_shape = weight_shape;
+
+    int64_t batch_count_i64 = 1;
+    for (size_t i = 0; i < input_shape.size() - 2; ++i) {
+      batch_count_i64 *= input_shape[i];
+    }
+    batch_count = CheckedCastToInt(batch_count_i64);
+
+    int64_t m_i64;
+    int64_t n_i64;
+    int64_t k_i64;
+    if (transpose_a) {
+      m_i64 = input_shape[input_shape.size() - 1];
+      k_i64 = input_shape[input_shape.size() - 2];
+    } else {
+      m_i64 = input_shape[input_shape.size() - 2];
+      k_i64 = input_shape[input_shape.size() - 1];
+    }
+    if (transpose_b) {
+      ASSERT_CHECK(weight_shape[weight_shape.size() - 1] == k_i64);
+      n_i64 = weight_shape[weight_shape.size() - 2];
+    } else {
+      ASSERT_CHECK(weight_shape[weight_shape.size() - 2] == k_i64);
+      n_i64 = weight_shape[weight_shape.size() - 1];
+    }
+    m = CheckedCastToInt(m_i64);
+    n = CheckedCastToInt(n_i64);
+    k = CheckedCastToInt(k_i64);
+
+    if (bias) {
+      ASSERT_CHECK(bias_shape.size() >= 1U);
+      ASSERT_CHECK(bias_shape[bias_shape.size() - 1] == n_i64);
+    }
+
+#if AP_ENABLE_DEBUG
+    std::cout << "-- [GemmEpilogueParams] batch_count: " << batch_count
+              << ", m: " << m << ", n: " << n << ", k: " << k << std::endl;
+    std::cout << "-- [GemmEpilogueParams] input: " << input << std::endl;
+    std::cout << "-- [GemmEpilogueParams] weight: " << weight << std::endl;
+    std::cout << "-- [GemmEpilogueParams] bias: " << bias << std::endl;
+    std::cout << "-- [GemmEpilogueParams] output: " << output << std::endl;
+    std::cout << "-- [GemmEpilogueParams] stream: " << stream << std::endl;
+#endif
+
+    shape_args.batch_stride_A = m_i64 * k_i64;
+    shape_args.batch_stride_B = (weight_shape.size() == 2) ? 0 : n_i64 * k_i64;
+    shape_args.batch_stride_D = m_i64 * n_i64;
+
+    shape_args.lda = transpose_a ? m_i64 : k_i64;
+    shape_args.ldb = transpose_b ? k_i64 : n_i64;
+    shape_args.ldd = n_i64;
+
+    bool is_C_bias = bias_shape.size() == 1UL;
+
+    /// Only available in RRR format
+    shape_args.batch_stride_C = (!bias || is_C_bias) ? 0 : m_i64 * n_i64;
+    shape_args.ldc_bias = (!bias || is_C_bias) ? 0 : n_i64;
+  }
+
+  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<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;
+  }
+};
+
+}  // namespace ap

python/paddle/apy/matmul_pass/matmul/profile.h

@@ -0,0 +1,151 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <functional>
+#include <utility>
+
+#ifdef __NVCC__
+#include <cuda_profiler_api.h>
+
+#define GPUEvent_t cudaEvent_t
+
+#ifndef GPUStream_t
+#define GPUStream_t cudaStream_t
+#endif
+
+#define GPUEventCreate(e) cudaEventCreate(e)
+#define GPUEventDestroy(e) cudaEventDestroy(e)
+#define GPUEventRecord(e, s) cudaEventRecord(e, s)
+#define GPUEventSynchronize(e) cudaEventSynchronize(e)
+#define GPUEventElapsedTime(ms, s, e) cudaEventElapsedTime(ms, s, e)
+#define GPUProfilerStart() cudaProfilerStart()
+#define GPUProfilerStop() cudaProfilerStop()
+#define GPUStreamSynchronize(s) cudaStreamSynchronize(s)
+#define CHECK_GPU CHECK_CUDA
+#endif
+
+#ifdef __HIPCC__
+#include <hip/hip_runtime.h>
+#include <hip/hip_runtime_api.h>
+
+#define GPUEvent_t hipEvent_t
+
+#ifndef GPUStream_t
+#define GPUStream_t hipStream_t
+#endif
+
+#define GPUEventCreate(e) hipEventCreate(e)
+#define GPUEventDestroy(e) hipEventDestroy(e)
+#define GPUEventRecord(e, s) hipEventRecord(e, s)
+#define GPUEventSynchronize(e) hipEventSynchronize(e)
+#define GPUEventElapsedTime(ms, s, e) hipEventElapsedTime(ms, s, e)
+#define GPUProfilerStart() hipProfilerStart()
+#define GPUProfilerStop() hipProfilerStop()
+#define GPUStreamSynchronize(s) hipStreamSynchronize(s)
+#define CHECK_GPU CHECK_HIP
+#endif
+
+namespace ap {
+
+class GpuTimer {
+ public:
+  explicit GpuTimer(bool profile) : profile_(profile) {
+    CHECK_GPU(GPUEventCreate(&start_));
+    CHECK_GPU(GPUEventCreate(&stop_));
+  }
+
+  ~GpuTimer() {
+    CHECK_GPU(GPUEventDestroy(start_));
+    CHECK_GPU(GPUEventDestroy(stop_));
+  }
+
+  void Start(GPUStream_t stream) {
+    CHECK_GPU(GPUEventRecord(start_, stream));
+    if (profile_) {
+      CHECK_GPU(GPUProfilerStart());
+    }
+  }
+
+  void Stop(GPUStream_t stream) {
+    CHECK_GPU(GPUEventRecord(stop_, stream));
+    if (profile_) {
+      CHECK_GPU(GPUProfilerStop());
+    }
+  }
+
+  float ElapsedTime() {
+    float milliseconds = 0;
+    CHECK_GPU(GPUEventSynchronize(stop_));
+    CHECK_GPU(GPUEventElapsedTime(&milliseconds, start_, stop_));
+    return milliseconds;
+  }
+
+ private:
+  bool profile_{false};
+  GPUEvent_t start_{nullptr};
+  GPUEvent_t stop_{nullptr};
+};
+
+template <typename FuncType, typename... Args>
+int ProfileBestConfig(const std::vector<FuncType> &funcs,
+                      void *stream_ptr,
+                      Args &&...args) {
+  std::cout
+      << "=================================================================="
+      << std::endl;
+
+  constexpr int kWarmupIters = 1;
+  constexpr int kRepeatIters = 100;
+
+  GpuTimer gpu_timer(false);
+  float min_time_ms = 100000.f;
+  int min_time_idx = -1;
+
+  GPUStream_t stream = *reinterpret_cast<GPUStream_t *>(stream_ptr);
+
+  for (int idx = 0; idx < funcs.size(); ++idx) {
+    auto func = funcs[idx];
+    for (int i = 0; i < kWarmupIters; i++) {
+      func(std::forward<Args>(args)...);
+    }
+    if (stream) {
+      CHECK_GPU(GPUStreamSynchronize(stream));
+    }
+
+    gpu_timer.Start(stream);
+    for (int i = 0; i < kRepeatIters; i++) {
+      func(std::forward<Args>(args)...);
+    }
+    gpu_timer.Stop(stream);
+
+    float elapsed_time_ms = gpu_timer.ElapsedTime();
+    std::cout << "-- [ProfileBestConfig] No " << idx
+              << ", elapsed_time: " << elapsed_time_ms << " ms" << std::endl;
+    if (elapsed_time_ms < min_time_ms) {
+      min_time_ms = elapsed_time_ms;
+      min_time_idx = idx;
+    }
+  }
+
+  std::cout << "-- [ProfileBestConfig] best config idx: " << min_time_idx
+            << std::endl;
+  std::cout
+      << "=================================================================="
+      << std::endl;
+  return min_time_idx;
+}
+
+}  // namespace ap

python/paddle/apy/matmul_pass/matmul_epilogue_pass.py

@@ -0,0 +1,140 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import access_topo_drr
+import pir
+
+
+class RemoveDataOpPairPass(access_topo_drr.DrrPass):
+    def __init__(self, src_data_op_name, dst_data_op_name):
+        self.src_data_op_name = pir.a_str(src_data_op_name)
+        self.dst_data_op_name = pir.a_str(dst_data_op_name)
+
+    def source_pattern(self, o, t):
+        o.src_data_op = o.ap_native_op("pd_op.data")
+        o.src_data_op([], [t.input0])
+        o.dst_data_op = o.ap_native_op("pd_op.data")
+        o.dst_data_op([], [t.input1])
+        o.up_spider_op = o.ap_native_op("ap_op.up_spider")
+        o.up_spider_op([t.input0, t.input1], [])
+
+    def constraint(self, o, t):
+        return [o.src_data_op.name, o.dst_data_op.name] == [
+            self.src_data_op_name,
+            self.dst_data_op_name,
+        ]
+
+    def result_pattern(self, o, t):
+        pass
+
+
+class RemoveDataOp2SumOp2DataOpPass(access_topo_drr.DrrPass):
+    def __init__(self, src_data_op_name, dst_data_op_name):
+        self.src_data_op_name = pir.a_str(src_data_op_name)
+        self.dst_data_op_name = pir.a_str(dst_data_op_name)
+
+    def source_pattern(self, o, t):
+        o.src_data_op = o.ap_native_op("pd_op.data")
+        o.src_data_op.name = self.src_data_op_name
+        o.src_data_op([], [t.input0])
+        o.full_int_array_op = o.ap_native_op("pd_op.full_int_array")
+        o.full_int_array_op([], [t.axis])
+        o.sum_op = o.ap_native_op("pd_op.sum")
+        o.sum_op([t.input0, t.axis], [t.sum_out])
+        o.dst_data_op = o.ap_native_op("pd_op.data")
+        o.dst_data_op.name = self.dst_data_op_name
+        o.dst_data_op([], [t.input1])
+        o.up_spider_op = o.ap_native_op("ap_op.up_spider")
+        o.up_spider_op([t.sum_out, t.input1], [])
+
+    def result_pattern(self, o, t):
+        pass
+
+
+class RemoveElementInputIndexPass(access_topo_drr.DrrPass):
+    def __init__(self, src_data_op_name, dst_load_from_global_op_name):
+        self.src_data_op_name = pir.a_str(src_data_op_name)
+        self.dst_load_from_global_op_name = pir.a_str(
+            dst_load_from_global_op_name
+        )
+
+    def source_pattern(self, o, t):
+        o.src_data_op = o.ap_native_op("pd_op.data")
+        o.src_data_op.name = self.src_data_op_name
+        o.src_data_op([], [t.src_input])
+
+        o.dst_load_from_global_op = o.ap_native_op("ap_op.load_from_global")
+        o.dst_load_from_global_op.index_func_unique_id = (
+            self.dst_load_from_global_op_name
+        )
+        o.dst_load_from_global_op(
+            [t.dst_input], [t.dst_load_from_global_output]
+        )
+        o.up_spider_op = o.ap_native_op("ap_op.up_spider")
+        o.up_spider_op([t.src_input, t.dst_load_from_global_output], [])
+
+    def result_pattern(self, o, t):
+        pass
+
+
+class RemoveBroadcastInputIndexPass(access_topo_drr.DrrPass):
+    def __init__(self, src_data_op_name, dst_load_from_global_op_name):
+        self.src_data_op_name = pir.a_str(src_data_op_name)
+        self.dst_load_from_global_op_name = pir.a_str(
+            dst_load_from_global_op_name
+        )
+
+    def source_pattern(self, o, t):
+        o.src_data_op = o.ap_native_op("pd_op.data")
+        o.src_data_op.name = self.src_data_op_name
+        o.src_data_op([], [t.input0])
+        o.full_int_array_op = o.ap_native_op("pd_op.full_int_array")
+        o.full_int_array_op([], [t.axis])
+        o.sum_op = o.ap_native_op("pd_op.sum")
+        o.sum_op([t.input0, t.axis], [t.sum_out])
+        o.dst_load_from_global_op = o.ap_native_op("ap_op.load_from_global")
+        o.dst_load_from_global_op.index_func_unique_id = (
+            self.dst_load_from_global_op_name
+        )
+        o.dst_load_from_global_op(
+            [t.dst_input], [t.dst_load_from_global_output]
+        )
+        o.up_spider_op = o.ap_native_op("ap_op.up_spider")
+        o.up_spider_op([t.sum_out, t.dst_load_from_global_output], [])
+
+    def result_pattern(self, o, t):
+        pass
+
+
+class RemoveOutputIndexPass(access_topo_drr.DrrPass):
+    def __init__(self, src_data_op_name, dst_store_to_global_op_name):
+        self.src_data_op_name = pir.a_str(src_data_op_name)
+        self.dst_store_to_global_op_name = pir.a_str(
+            dst_store_to_global_op_name
+        )
+
+    def source_pattern(self, o, t):
+        o.src_data_op = o.ap_native_op("pd_op.data")
+        o.src_data_op.name = self.src_data_op_name
+        o.src_data_op([], [t.src_input])
+        o.down_spider_op = o.ap_native_op("ap_op.down_spider")
+        o.down_spider_op([t.src_input], [t.dst_output_val])
+        o.dst_store_to_global_op = o.ap_native_op("ap_op.store_to_global")
+        o.dst_store_to_global_op.index_func_unique_id = (
+            self.dst_store_to_global_op_name
+        )
+        o.dst_store_to_global_op([t.dst_output, t.dst_output_val], [])
+
+    def result_pattern(self, o, t):
+        pass

python/paddle/apy/matmul_pass/matmul_variadic_ptn.py

@@ -0,0 +1,673 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import abstract_drr
+import access_topo_drr  # noqa: F401
+import ap
+import index_program_translator_util
+import ir_tools
+import kernel_arg_id_util
+import kernel_arg_translator_util  # noqa: F401
+import low_level_ir_code_gen_ctx_util  # noqa: F401
+import matmul_epilogue_pass
+import matmul_variadic_tpl
+import op_compute_translator_util
+import op_conversion_drr_pass  # noqa: F401
+import pir  # noqa: F401
+import program_translator_util
+import topo_drr_pass
+import umprime  # noqa: F401
+
+
+class MatmulEpilogueFusion(abstract_drr.DrrPass):
+    def source_pattern(self, o, t):
+        in_num = self.number_of_inputs()
+        out_num = self.number_of_outputs()
+        o.matmul_op = o.ap_native_op("pd_op.matmul")
+        o.matmul_op([t.input0, t.input1], [t.mm_out])
+        o.trivial_op = o.ap_trivial_fusion_op()
+        o.trivial_op(
+            [
+                t.mm_out,
+                *ap.map(
+                    lambda index: getattr(t, f"input{index + 2}"),
+                    range(in_num - 2),
+                ),
+            ],
+            ap.map(lambda index: getattr(t, f"output{index}"), range(out_num)),
+        )
+
+    def result_pattern(self, o, t):
+        in_num = self.number_of_inputs()
+        out_num = self.number_of_outputs()
+        o.fustion_op = o.ap_pattern_fusion_op(self.code_gen)
+        o.fustion_op(
+            ap.map(lambda index: getattr(t, f"input{index}"), range(in_num)),
+            ap.map(lambda index: getattr(t, f"output{index}"), range(out_num)),
+        )
+
+    def constraint(self, o, t):
+        program = ir_tools.copy_fused_ops_to_program(
+            o.trivial_op, tensor_match_ctx=t
+        )
+        print("before-umprime: ", program)
+        # umprime passes
+        pass_manager = ir_tools.create_pass_manager()
+        pass_manager.add_pass(ir_tools.create_access_topo_drr_pass("umprime"))
+        pass_manager.add_pass(ir_tools.create_dce_pass())
+        pass_manager.run(program)
+        print("before-access_topo_pass", program)
+        init_pass_manager = ir_tools.create_pass_manager()
+        init_down_spider = topo_drr_pass.InitDownSpiderAccessTopoPass("mm_out")
+        init_pass_manager.add_pass(
+            ir_tools.create_access_topo_drr_one_step_pass(init_down_spider)
+        )
+        outputs_name_list = ap.map(
+            lambda i: f"output{i}", range(self.number_of_outputs())
+        )
+        inputs_name_list = (
+            ap.map(
+                lambda i: f"input{i + 2}", range(self.number_of_inputs() - 2)
+            )
+            if self.number_of_inputs() > 2
+            else []
+        )
+        print('inputs_name_list: ', ', '.join(inputs_name_list))
+        init_fake_data_for_yield_input = (
+            topo_drr_pass.FakeDataForYieldAccessTopoPass(outputs_name_list)
+        )
+        init_pass_manager.add_pass(
+            ir_tools.create_access_topo_drr_one_step_pass(
+                init_fake_data_for_yield_input
+            )
+        )
+        init_pass_manager.run(program)
+        print("after-init-access_topo_pass", program)
+        pass_manager = ir_tools.create_pass_manager()
+        pass_manager.add_pass(ir_tools.create_access_topo_drr_pass("default"))
+        pass_manager.add_pass(ir_tools.create_dce_pass())
+        pass_manager.run(program)
+        print("after-apply-access_topo_pass", program)
+        pass_manager = ir_tools.create_pass_manager()
+        ap.map(
+            lambda dst_name: pass_manager.add_pass(
+                ir_tools.create_access_topo_drr_one_step_pass(
+                    matmul_epilogue_pass.RemoveDataOpPairPass(
+                        src_data_op_name="mm_out", dst_data_op_name=dst_name
+                    )
+                )
+            ),
+            inputs_name_list,
+        )
+        ap.map(
+            lambda dst_name: pass_manager.add_pass(
+                ir_tools.create_access_topo_drr_one_step_pass(
+                    matmul_epilogue_pass.RemoveDataOp2SumOp2DataOpPass(
+                        src_data_op_name="mm_out", dst_data_op_name=dst_name
+                    )
+                )
+            ),
+            inputs_name_list,
+        )
+
+        ap.map(
+            lambda dst_name: pass_manager.add_pass(
+                ir_tools.create_access_topo_drr_one_step_pass(
+                    matmul_epilogue_pass.RemoveDataOpPairPass(
+                        src_data_op_name="mm_out", dst_data_op_name=dst_name
+                    )
+                )
+            ),
+            outputs_name_list,
+        )
+        pass_manager.add_pass(ir_tools.create_dce_pass())
+        pass_manager.run(program)
+        print("after-remove-input-output-access_topo_pass", program)
+        return program.empty()
+
+    def _insert_load_from_global(self, program, input_names):
+        init_pass_manager = ir_tools.create_pass_manager()
+
+        def AddPass(input_name):
+            ir_pass = topo_drr_pass.InitNaiveLoadFromGlobalAccessTopoPass(
+                input_name
+            )
+            init_pass_manager.add_pass(
+                ir_tools.create_access_topo_drr_one_step_pass(ir_pass)
+            )
+
+        ap.map(AddPass, input_names)
+        init_pass_manager.run(program)
+
+    def _insert_store_to_global(self, program, output_names):
+        init_pass_manager = ir_tools.create_pass_manager()
+        ir_pass = topo_drr_pass.FakeDataStoreToGlobalForYieldAccessTopoPass(
+            output_names
+        )
+        init_pass_manager.add_pass(
+            ir_tools.create_access_topo_drr_one_step_pass(ir_pass)
+        )
+        init_pass_manager.run(program)
+
+    def _make_kernel_arg_translator(self):
+        return matmul_variadic_tpl.make_kernel_arg_translator()
+
+    def _apply_topo_access_passes(self, mut_program, anchor_data_op_name):
+        init_pass_manager = ir_tools.create_pass_manager()
+        init_down_spider = topo_drr_pass.InitDownSpiderAccessTopoPass(
+            anchor_data_op_name
+        )
+        init_pass_manager.add_pass(
+            ir_tools.create_access_topo_drr_one_step_pass(init_down_spider)
+        )
+        init_pass_manager.run(mut_program)
+        pass_manager = ir_tools.create_pass_manager()
+        pass_manager.add_pass(ir_tools.create_access_topo_drr_pass("default"))
+        pass_manager.add_pass(ir_tools.create_dce_pass())
+        pass_manager.run(mut_program)
+
+    def _simplify_index_program(self, mut_program):
+        pass_manager = ir_tools.create_pass_manager()
+        drr_pass = topo_drr_pass.ConvertUpSpiderStoreDataOpToYieldOpPass()
+        pass_manager.add_pass(
+            ir_tools.create_access_topo_drr_one_step_pass(drr_pass)
+        )
+        drr_pass = topo_drr_pass.ConvertDownSpiderStoreDataOpToYieldOpPass()
+        pass_manager.add_pass(
+            ir_tools.create_access_topo_drr_one_step_pass(drr_pass)
+        )
+        pass_manager.add_pass(ir_tools.create_dce_pass())
+        pass_manager.run(mut_program)
+        return mut_program
+
+    def _make_index_func_unique_id2index_program(
+        self, compute_program, anchor_data_op_name, input_names, output_names
+    ):
+        full_index_program = compute_program.clone()
+        self._apply_topo_access_passes(full_index_program, anchor_data_op_name)
+        print('full_index_program: ', full_index_program)
+
+        def MatchAndCopyInputIndex(dst_input_name):
+            pass_manager = ir_tools.create_pass_manager()
+            removed_programs = ap.MutableList()
+            rm_elementwise_drr_pass = (
+                matmul_epilogue_pass.RemoveElementInputIndexPass(
+                    src_data_op_name=anchor_data_op_name,
+                    dst_load_from_global_op_name=dst_input_name,
+                )
+            )
+            rm_elementwise_ir_pass = (
+                ir_tools.create_access_topo_drr_one_step_pass(
+                    rm_elementwise_drr_pass,
+                    matched_pattern_mut_list=removed_programs,
+                )
+            )
+            pass_manager.add_pass(rm_elementwise_ir_pass)
+            rm_broadcast_drr_pass = (
+                matmul_epilogue_pass.RemoveBroadcastInputIndexPass(
+                    src_data_op_name=anchor_data_op_name,
+                    dst_load_from_global_op_name=dst_input_name,
+                )
+            )
+            rm_broadcast_ir_pass = (
+                ir_tools.create_access_topo_drr_one_step_pass(
+                    rm_broadcast_drr_pass,
+                    matched_pattern_mut_list=removed_programs,
+                )
+            )
+            pass_manager.add_pass(rm_broadcast_ir_pass)
+            pass_manager.run(full_index_program)
+
+            def Converter(program):
+                return [dst_input_name, self._simplify_index_program(program)]
+
+            return ap.map(Converter, removed_programs)
+
+        input_and_index_programs = ap.flat_map(
+            MatchAndCopyInputIndex, input_names
+        )
+
+        def MatchAndCopyOutputIndex(dst_output_name):
+            print('full_index_program output: ', full_index_program)
+            pass_manager = ir_tools.create_pass_manager()
+            removed_programs = ap.MutableList()
+            drr_pass = matmul_epilogue_pass.RemoveOutputIndexPass(
+                src_data_op_name=anchor_data_op_name,
+                dst_store_to_global_op_name=dst_output_name,
+            )
+            ir_pass = ir_tools.create_access_topo_drr_one_step_pass(
+                drr_pass, matched_pattern_mut_list=removed_programs
+            )
+            pass_manager.add_pass(ir_pass)
+            pass_manager.run(full_index_program)
+
+            def Converter(program):
+                return [dst_output_name, self._simplify_index_program(program)]
+
+            print('len removed of output: ', len(removed_programs))
+            return ap.map(Converter, removed_programs)
+
+        output_and_index_programs = ap.flat_map(
+            MatchAndCopyOutputIndex, output_names
+        )
+        return ap.OrderedDict(
+            [*input_and_index_programs, *output_and_index_programs]
+        )
+
+    def _replace_with_load_from_register(
+        self, mut_program, load_ir_value_name, register_var_name
+    ):
+        pass_manager = ir_tools.create_pass_manager()
+        drr_pass = topo_drr_pass.ReplaceWithLoadFromRegisterPass(
+            name=load_ir_value_name, register_var_name=register_var_name
+        )
+        pass_manager.add_pass(
+            ir_tools.create_access_topo_drr_one_step_pass(drr_pass)
+        )
+        pass_manager.add_pass(ir_tools.create_dce_pass())
+        pass_manager.run(mut_program)
+        return mut_program
+
+    def _replace_with_store_to_register(
+        self, mut_program, store_ir_value_name, register_var_name
+    ):
+        pass_manager = ir_tools.create_pass_manager()
+        drr_pass = topo_drr_pass.ReplaceWithStoreToRegisterPass(
+            name=store_ir_value_name, register_var_name=register_var_name
+        )
+        pass_manager.add_pass(
+            ir_tools.create_access_topo_drr_one_step_pass(drr_pass)
+        )
+        pass_manager.add_pass(ir_tools.create_dce_pass())
+        pass_manager.run(mut_program)
+        return mut_program
+
+    def _get_program_translator(self, ctx, o, t):
+        outputs_name_list = ap.map(
+            lambda i: f"output{i}", range(self.number_of_outputs())
+        )
+        other_outputs_name_list = ap.map(
+            lambda i: f"output{i + 1}", range(self.number_of_outputs() - 1)
+        )
+        local_outputs_name_list = ap.map(
+            lambda i: f"out{i}", range(self.number_of_outputs())
+        )
+        inputs_name_list = (
+            ap.map(
+                lambda i: f"input{i + 2}", range(self.number_of_inputs() - 2)
+            )
+            if self.number_of_inputs() > 2
+            else []
+        )
+        mut_program = ir_tools.copy_fused_ops_to_program(
+            o.trivial_op, tensor_match_ctx=t
+        )
+        print("before-umprime: ", mut_program)
+        pass_manager = ir_tools.create_pass_manager()
+        pass_manager.add_pass(ir_tools.create_access_topo_drr_pass("umprime"))
+        pass_manager.add_pass(ir_tools.create_dce_pass())
+        pass_manager.run(mut_program)
+        self._insert_load_from_global(mut_program, input_names=["mm_out"])
+        self._insert_load_from_global(mut_program, input_names=inputs_name_list)
+        self._insert_store_to_global(
+            mut_program, output_names=outputs_name_list
+        )
+        kernel_arg_translator = self._make_kernel_arg_translator()
+        index_func_unique_id2index_program = (
+            self._make_index_func_unique_id2index_program(
+                mut_program,
+                anchor_data_op_name="mm_out",
+                input_names=inputs_name_list,
+                output_names=other_outputs_name_list,
+            )
+        )
+        print(
+            "index_func_unique_id2index_program:\n",
+            index_func_unique_id2index_program,
+        )
+        index_program_translator_map = index_program_translator_util.IndexProgramTranslatorMap(
+            index_func_unique_id2index_program=index_func_unique_id2index_program,
+            kernel_arg_translator=kernel_arg_translator,
+            anchor_iter_var_names=matmul_variadic_tpl.get_anchor_iter_var_names(),
+        )
+        self._replace_with_load_from_register(
+            mut_program, load_ir_value_name="mm_out", register_var_name="x"
+        )
+        self._replace_with_store_to_register(mut_program, "output0", "out")
+        print("mut_program:", mut_program)
+        op_compute_translator_maker = (
+            op_compute_translator_util.OpComputeTranslatorFactory()
+        )
+        program_translator = program_translator_util.ProgramTranslator(
+            program_property=mut_program.copy_to_const_program_data(),
+            kernel_arg_translator=kernel_arg_translator,
+            index_program_translator_map=index_program_translator_map,
+            op_translator_maker=op_compute_translator_maker,
+        )
+
+        return program_translator
+
+    def code_gen(self, ctx, o, t):
+        program_translator = self._get_program_translator(ctx, o, t)
+        mut_kernel_arg_id_registry = kernel_arg_id_util.KernelArgIdNameRegistry(
+            code_gen_ctx=ctx, tensor_match_ctx=t, name_prefix=""
+        )
+
+        template_module = matmul_variadic_tpl.MatmulVariadicTemplate(
+            program_translator=program_translator,
+            mut_kernel_arg_id_registry=mut_kernel_arg_id_registry,
+        )
+
+        def get_symbolic_shape_args_list(sym_dim):
+            return ctx.dim_expr_kernel_arg_id(sym_dim)
+
+        input0_shape_kargs = ap.map(
+            get_symbolic_shape_args_list, t.input0.symbolic_shape_to_list()
+        )
+        input1_shape_kargs = ap.map(
+            get_symbolic_shape_args_list, t.input1.symbolic_shape_to_list()
+        )
+        return template_module.compile(
+            input0_karg=ctx.in_tensor_data_ptr_kernel_arg_id(t.input0),
+            input1_karg=ctx.in_tensor_data_ptr_kernel_arg_id(t.input1),
+            output_karg=ctx.out_tensor_data_ptr_kernel_arg_id(t.output0),
+            input0_shape_kargs=input0_shape_kargs,
+            input1_shape_kargs=input1_shape_kargs,
+        )
+
+
+class NumberOfInputsTrait0:
+    def number_of_inputs(self):
+        return 0
+
+
+class NumberOfInputsTrait1:
+    def number_of_inputs(self):
+        return 1
+
+
+class NumberOfInputsTrait2:
+    def number_of_inputs(self):
+        return 2
+
+
+class NumberOfInputsTrait3:
+    def number_of_inputs(self):
+        return 3
+
+
+class NumberOfInputsTrait4:
+    def number_of_inputs(self):
+        return 4
+
+
+class NumberOfInputsTrait5:
+    def number_of_inputs(self):
+        return 5
+
+
+class NumberOfInputsTrait6:
+    def number_of_inputs(self):
+        return 6
+
+
+class NumberOfInputsTrait7:
+    def number_of_inputs(self):
+        return 7
+
+
+class NumberOfInputsTrait8:
+    def number_of_inputs(self):
+        return 8
+
+
+class NumberOfInputsTrait9:
+    def number_of_inputs(self):
+        return 9
+
+
+class NumberOfInputsTrait10:
+    def number_of_inputs(self):
+        return 10
+
+
+class NumberOfInputsTrait11:
+    def number_of_inputs(self):
+        return 11
+
+
+class NumberOfInputsTrait12:
+    def number_of_inputs(self):
+        return 12
+
+
+class NumberOfInputsTrait13:
+    def number_of_inputs(self):
+        return 13
+
+
+class NumberOfInputsTrait14:
+    def number_of_inputs(self):
+        return 14
+
+
+class NumberOfInputsTrait15:
+    def number_of_inputs(self):
+        return 15
+
+
+class NumberOfInputsTrait16:
+    def number_of_inputs(self):
+        return 16
+
+
+class NumberOfInputsTrait17:
+    def number_of_inputs(self):
+        return 17
+
+
+class NumberOfOutputsTrait0:
+    def number_of_outputs(self):
+        return 0
+
+
+class NumberOfOutputsTrait1:
+    def number_of_outputs(self):
+        return 1
+
+
+class NumberOfOutputsTrait2:
+    def number_of_outputs(self):
+        return 2
+
+
+class NumberOfOutputsTrait3:
+    def number_of_outputs(self):
+        return 3
+
+
+class NumberOfOutputsTrait4:
+    def number_of_outputs(self):
+        return 4
+
+
+class NumberOfOutputsTrait5:
+    def number_of_outputs(self):
+        return 5
+
+
+class NumberOfOutputsTrait6:
+    def number_of_outputs(self):
+        return 6
+
+
+class NumberOfOutputsTrait7:
+    def number_of_outputs(self):
+        return 7
+
+
+class NumberOfOutputsTrait8:
+    def number_of_outputs(self):
+        return 8
+
+
+class NumberOfOutputsTrait9:
+    def number_of_outputs(self):
+        return 9
+
+
+class NumberOfOutputsTrait10:
+    def number_of_outputs(self):
+        return 10
+
+
+class NumberOfOutputsTrait11:
+    def number_of_outputs(self):
+        return 11
+
+
+class NumberOfOutputsTrait12:
+    def number_of_outputs(self):
+        return 12
+
+
+class NumberOfOutputsTrait13:
+    def number_of_outputs(self):
+        return 13
+
+
+class NumberOfOutputsTrait14:
+    def number_of_outputs(self):
+        return 14
+
+
+class NumberOfOutputsTrait15:
+    def number_of_outputs(self):
+        return 15
+
+
+class NumberOfOutputsTrait16:
+    def number_of_outputs(self):
+        return 16
+
+
+class NumberOfOutputsTrait17:
+    def number_of_outputs(self):
+        return 17
+
+
+class NumberOfOutputsTrait18:
+    def number_of_outputs(self):
+        return 18
+
+
+class NumberOfOutputsTrait19:
+    def number_of_outputs(self):
+        return 19
+
+
+class NumberOfOutputsTrait20:
+    def number_of_outputs(self):
+        return 20
+
+
+class NumberOfOutputsTrait21:
+    def number_of_outputs(self):
+        return 21
+
+
+class NumberOfOutputsTrait22:
+    def number_of_outputs(self):
+        return 22
+
+
+def get_mixin_class(base_class, number_of_inputs, number_of_outputs):
+    num_inputs_to_input_trait_class = [
+        None,
+        NumberOfInputsTrait1,
+        NumberOfInputsTrait2,
+        NumberOfInputsTrait3,
+        NumberOfInputsTrait3,
+        NumberOfInputsTrait4,
+        NumberOfInputsTrait5,
+        NumberOfInputsTrait6,
+        NumberOfInputsTrait7,
+        NumberOfInputsTrait8,
+        NumberOfInputsTrait9,
+        NumberOfInputsTrait10,
+        NumberOfInputsTrait11,
+        NumberOfInputsTrait12,
+        NumberOfInputsTrait13,
+        NumberOfInputsTrait14,
+        NumberOfInputsTrait15,
+        NumberOfInputsTrait16,
+        NumberOfInputsTrait17,
+    ]
+    num_outputs_to_output_trait_class = [
+        None,
+        NumberOfOutputsTrait1,
+        NumberOfOutputsTrait2,
+        NumberOfOutputsTrait3,
+        NumberOfOutputsTrait4,
+        NumberOfOutputsTrait5,
+        NumberOfOutputsTrait6,
+        NumberOfOutputsTrait7,
+        NumberOfOutputsTrait8,
+        NumberOfOutputsTrait9,
+        NumberOfOutputsTrait10,
+        NumberOfOutputsTrait11,
+        NumberOfOutputsTrait12,
+        NumberOfOutputsTrait13,
+        NumberOfOutputsTrait14,
+        NumberOfOutputsTrait15,
+        NumberOfOutputsTrait16,
+        NumberOfOutputsTrait17,
+        NumberOfOutputsTrait18,
+        NumberOfOutputsTrait19,
+        NumberOfOutputsTrait20,
+        NumberOfOutputsTrait21,
+        NumberOfOutputsTrait22,
+    ]
+    return type(
+        f"MatmulEpilogueFusion{number_of_inputs}_{number_of_outputs}",
+        [
+            base_class,
+            num_inputs_to_input_trait_class[number_of_inputs],
+            num_outputs_to_output_trait_class[number_of_outputs],
+        ],
+        ap.SerializableAttrMap(),
+    )
+
+
+# abstract_drr.register_drr_pass("matmul_binary_outs_fusion", nice=0)(get_mixin_class(MatmulEpilogueFusion, 3, 2))
+
+
+def register_class(base_class, max_num_inputs, max_num_outputs):
+    def register_drr_class(num_inputs, num_outputs):
+        abstract_drr.register_drr_pass(
+            f"matmul_binary_in{num_inputs}_out{num_outputs}_fusion", nice=0
+        )(get_mixin_class(base_class, num_inputs, num_outputs))
+
+    def register_num_inputs_drr_classes(num_inputs):
+        def register_num_outputs_drr_classes(num_outputs):
+            return register_drr_class(num_inputs + 2, num_outputs + 1)
+
+        ap.map(register_num_outputs_drr_classes, range(max_num_outputs))
+
+    ap.map(register_num_inputs_drr_classes, range(max_num_inputs))
+
+
+register_class(
+    base_class=MatmulEpilogueFusion, max_num_inputs=10, max_num_outputs=10
+)

python/paddle/apy/matmul_pass/matmul_variadic_tpl.py

@@ -0,0 +1,363 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import ap
+import compile_command_util
+import kernel_arg_translator_util
+import low_level_ir_code_gen_ctx_util
+
+
+def make_kernel_arg_translator():
+    return kernel_arg_translator_util.KernelArgTranslator(
+        param_struct_name="args"
+    )
+
+
+def get_anchor_iter_var_names():
+    return ["coord.batch", "coord.row", "coord.column"]
+
+
+class MatmulVariadicTemplate:
+    def __init__(
+        self,
+        program_translator,
+        mut_kernel_arg_id_registry,
+    ):
+        self.program_translator = program_translator
+        self.mut_kernel_arg_id_registry = mut_kernel_arg_id_registry
+        self.kernel_arg_translator = make_kernel_arg_translator()
+        self.dtype2type_name = ap.OrderedDict(
+            [
+                [ap.PointerType.const_float_ptr, "const float*"],
+                [ap.PointerType.const_float16_ptr, "const half*"],
+                [ap.PointerType.const_bfloat16_ptr, "const bfloat16*"],
+                [ap.PointerType.float_ptr, "float*"],
+                [ap.PointerType.float16_ptr, "half*"],
+                [ap.PointerType.bfloat16_ptr, "bfloat16*"],
+                [ap.DataType.float, "float"],
+                [ap.DataType.float16, "half"],
+                [ap.DataType.bfloat16, "bfloat16"],
+                [ap.DataType.int64_t, "int64_t"],
+            ]
+        )
+        self.input_dim_karg_to_shape_access = ap.MutableOrderedDict()
+        self.kernel_name = "MatmulVariadicKernel"
+        self.library_name = "matmul_variadic_kernel"
+
+    def _register_name(self, pair):
+        registry = self.mut_kernel_arg_id_registry
+        registry.get_or_create_kernel_arg_id_manul_var_name(
+            kernel_arg_id=pair[0], cpp_var_name=pair[1]
+        )
+
+    def compile(
+        self,
+        input0_karg,
+        input1_karg,
+        output_karg,
+        input0_shape_kargs,
+        input1_shape_kargs,
+    ):
+        kargs_name_pair_list = [
+            [input0_karg, "input0"],
+            [input1_karg, "input1"],
+            [output_karg, "output"],
+            *ap.map(
+                lambda i: [input0_shape_kargs[i], f"input0_dim{i}"],
+                range(len(input0_shape_kargs)),
+            ),
+            *ap.map(
+                lambda i: [input1_shape_kargs[i], f"input1_dim{i}"],
+                range(len(input1_shape_kargs)),
+            ),
+        ]
+
+        ap.map(self._register_name, kargs_name_pair_list)
+        mut_lir_code_gen_ctx = (
+            low_level_ir_code_gen_ctx_util.CudaLikeIrCodeGenCtx(
+                compute_dtype=ap.DataType.float
+            )
+        )
+        self.program_translator.translate(
+            mut_kernel_arg_id_registry=self.mut_kernel_arg_id_registry,
+            mut_lir_code_gen_ctx=mut_lir_code_gen_ctx,
+        )
+        trivial_code_str = mut_lir_code_gen_ctx.get_stmts_joined_str(
+            indent="    "
+        )
+
+        project_module = self.make_project(
+            trivial_code_str,
+            input0_karg,
+            input1_karg,
+            output_karg,
+            input0_shape_kargs,
+            input1_shape_kargs,
+        )
+        return CodeGenResult(  # noqa: F821
+            module=project_module,
+            kernel_dispatch_func=KernelDispatch,
+            kernel_dispatch_const_data=ap.SerializableAttrMap(
+                kernel_args_getters=self.get_kernel_arg_runtime_getters()
+            ),
+        )
+
+    def get_kernel_arg_runtime_getters(self):
+        all_kernel_arg_id_and_unique_names = self.mut_kernel_arg_id_registry.all_kernel_arg_id2unique_name.items()
+        return ap.map(
+            lambda pair: pair[0].runtime_getter,
+            all_kernel_arg_id_and_unique_names,
+        )
+
+    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()
+        return ap.map(
+            lambda pair: pair[0].type, all_kernel_arg_id_and_unique_names
+        )
+
+    def get_kernel_arg_id_var_name(self, kernel_arg_id):
+        all_kernel_arg_id2unique_name = (
+            self.mut_kernel_arg_id_registry.all_kernel_arg_id2unique_name
+        )
+        return all_kernel_arg_id2unique_name[kernel_arg_id]
+
+    def get_kernel_arg_list_str(self, for_declare):
+        def declare_epilogue_arguments_field(pair):
+            kernel_arg_id = pair[0]
+            var_name = pair[1]
+            field_name = self.kernel_arg_translator.get_param_struct_field_name(
+                var_name
+            )
+            dtype = kernel_arg_id.type
+            type_name = self.dtype2type_name[dtype]
+            return (
+                f"{type_name} {field_name}" if for_declare else f"{field_name}"
+            )
+
+        all_kernel_arg_id_and_names = self.mut_kernel_arg_id_registry.all_kernel_arg_id2unique_name.items()
+        return ", ".join(
+            ap.map(
+                declare_epilogue_arguments_field, all_kernel_arg_id_and_names
+            )
+        )
+
+    def get_epilogue_arguments_fields_str(self, indent):
+        def declare_epilogue_arguments_field(pair):
+            kernel_arg_id = pair[0]
+            var_name = pair[1]
+            field_name = self.kernel_arg_translator.get_param_struct_field_name(
+                var_name
+            )
+            dtype = kernel_arg_id.type
+            type_name = self.dtype2type_name[dtype]
+            return f"{type_name} {field_name};"
+
+        generated_kernel_arg_id_and_names = self.mut_kernel_arg_id_registry.generated_kernel_arg_id2unique_name.items()
+        return f"\n{indent}".join(
+            ap.map(
+                declare_epilogue_arguments_field,
+                generated_kernel_arg_id_and_names,
+            )
+        )
+
+    def get_epilogue_arguments_init_str(self, param_obj_name, indent):
+        def declare_epilogue_arguments_assign(pair):
+            kernel_arg_id = pair[0]
+            var_name = pair[1]
+            field_name = self.kernel_arg_translator.get_param_struct_field_name(
+                var_name
+            )
+            return f"{param_obj_name}.{field_name} = {var_name};"
+
+        generated_kernel_arg_id_and_names = self.mut_kernel_arg_id_registry.generated_kernel_arg_id2unique_name.items()
+        return f"\n{indent}".join(
+            ap.map(
+                declare_epilogue_arguments_assign,
+                generated_kernel_arg_id_and_names,
+            )
+        )
+
+    def get_params_input_shape_init_str(
+        self, input_name, input_shape_kargs, indent
+    ):
+        def init_input_shape_with_args(i):
+            def get_creator():
+                return f"{input_name}_shape[{i}]"
+
+            karg_var_name = self.get_kernel_arg_id_var_name(
+                input_shape_kargs[i]
+            )
+            self.input_dim_karg_to_shape_access.get_or_create(
+                karg_var_name, get_creator
+            )
+            return f"{indent}{input_name}_shape[{i}] = {karg_var_name};"
+
+        shape_vector_init_str = (
+            f"{input_name}_shape.resize({len(input_shape_kargs)});\n"
+        )
+        return shape_vector_init_str + "\n".join(
+            ap.map(init_input_shape_with_args, range(len(input_shape_kargs)))
+        )
+
+    def make_project(
+        self,
+        trivial_code_str,
+        input0_karg,
+        input1_karg,
+        output_karg,
+        input0_shape_kargs,
+        input1_shape_kargs,
+    ):
+        code_template = """
+// auto generated codes
+#include "matmul.h"
+#include <vector>
+
+namespace ap {
+
+template <typename T>
+struct VariadicEpilogueFunctor {
+  struct Arguments {
+    ${AP_EPILOGUE_ARGUMENTS_FIELDS}
+  };
+
+  // Note: need to support vectorized operation
+  __forceinline__ __host__ __device__
+  T operator()(T x, const Arguments& args, const MatrixCoord& coord) const {
+    T out;
+    ${AP_EPILOGUE_COMPUTATION_STATEMENTS}
+    return out;
+  }
+};
+
+template <int TuningConfigId>
+static void RunMatmulWithVariadicKernel(const GemmEpilogueParams &params, ${AP_KERNEL_ARGS_DECLARE}) {
+  using ElementT = ${output_dtype};
+  using ElementComputeT = float;
+
+  typename VariadicEpilogueFunctor<ElementComputeT>::Arguments epilogue_args;
+
+  ${AP_EPILOGUE_ARGUMENTS_INIT}
+
+  constexpr int AlignA = Alignment<ElementT, ${k_value}>::kValue;
+  constexpr int AlignB = Alignment<ElementT, ${n_value}>::kValue;
+
+  MatmulAddVariadic<ElementT, ElementComputeT, VariadicEpilogueFunctor,
+                           AlignA, AlignB, TuningConfigId>(params, epilogue_args);
+}
+
+} // namespace ap
+
+extern "C" {
+
+void ${kernel_name}(void* stream_ptr, ${AP_KERNEL_ARGS_DECLARE}) {
+  std::vector<int64_t> ${input0}_shape;
+  ${AP_PARAMS_INPUT0_SHAPE_INIT}
+
+  std::vector<int64_t> ${input1}_shape;
+  ${AP_PARAMS_INPUT1_SHAPE_INIT}
+
+  ap::GemmEpilogueParams params(
+      stream_ptr, ${input0}, ${input1}, nullptr, ${output}, ${input0}_shape, ${input1}_shape, std::vector<int64_t>{});
+
+#if AP_ENABLE_AUTOTUNE
+  AP_AUTOTUNE_${output_dtype}(ap::RunMatmulWithVariadicKernel, stream_ptr, params, ${AP_KERNEL_ARGS_CALL});
+#else
+  ap::RunMatmulWithVariadicKernel<ap::DefaultConfig::kConfigId>(params, ${AP_KERNEL_ARGS_CALL});
+#endif
+}
+}
+  """
+
+        output_dtype = self.dtype2type_name[output_karg.type.data_type]
+        code = (
+            code_template.replace(
+                "${AP_EPILOGUE_COMPUTATION_STATEMENTS}", trivial_code_str
+            )
+            .replace(
+                "${AP_KERNEL_ARGS_DECLARE}",
+                self.get_kernel_arg_list_str(for_declare=True),
+            )
+            .replace(
+                "${AP_KERNEL_ARGS_CALL}",
+                self.get_kernel_arg_list_str(for_declare=False),
+            )
+            .replace(
+                "${AP_PARAMS_INPUT0_SHAPE_INIT}",
+                self.get_params_input_shape_init_str(
+                    "${input0}", input0_shape_kargs, indent="  "
+                ),
+            )
+            .replace(
+                "${AP_PARAMS_INPUT1_SHAPE_INIT}",
+                self.get_params_input_shape_init_str(
+                    "${input1}", input1_shape_kargs, indent="  "
+                ),
+            )
+            .replace(
+                "${AP_EPILOGUE_ARGUMENTS_FIELDS}",
+                self.get_epilogue_arguments_fields_str(indent="    "),
+            )
+            .replace(
+                "${AP_EPILOGUE_ARGUMENTS_INIT}",
+                self.get_epilogue_arguments_init_str(
+                    "epilogue_args", indent="  "
+                ),
+            )
+            .replace("${kernel_name}", self.kernel_name)
+            .replace("${input0}", self.get_kernel_arg_id_var_name(input0_karg))
+            .replace("${input1}", self.get_kernel_arg_id_var_name(input1_karg))
+            .replace("${output}", self.get_kernel_arg_id_var_name(output_karg))
+            .replace("${output_dtype}", output_dtype)
+            .replace("${k_value}", f"{input0_shape_kargs[-1].value}")
+            .replace("${n_value}", f"{input1_shape_kargs[-1].value}")
+        )
+
+        dir_name = ap.dirname(__file__)
+        compile_command_generator = (
+            compile_command_util.CompileCommandGenerator()
+        )
+        compile_cmd = compile_command_generator(
+            "matmul", dir_name, self.library_name
+        )
+        file_ext = compile_command_generator.file_ext
+
+        return CodeModule(  # noqa: F821
+            FuncDeclare(  # noqa: F821
+                ap.DataType.void,
+                self.kernel_name,
+                [ap.PointerType.void_ptr, *self.get_kernel_arg_types()],
+            ),
+            Project(  # noqa: F821
+                nested_files=Project.Directory(  # noqa: F821
+                    [
+                        f"{self.library_name}.{file_ext}",
+                        Project.FileContent(code),  # noqa: F821
+                    ],
+                    ["make.sh", Project.FileContent(compile_cmd)],  # noqa: F821
+                ),
+                compile_cmd="sh make.sh",
+                so_relative_path=f"lib{self.library_name}.so",
+            ),
+        )
+
+
+def KernelDispatch(ctx):
+    import ap
+
+    so_func = ctx.get_so_function("MatmulVariadicKernel")
+    stream_ptr = ctx.device_ctx.get_stream_addr_as_void_ptr()
+    getters = ctx.kernel_dispatch_const_data.kernel_args_getters
+    args = [stream_ptr, *ap.map(lambda getter: getter(ctx), getters)]
+    ap.apply(so_func, args)

python/paddle/apy/matmul_pass/op_compute_translator_util.py

@@ -0,0 +1,836 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import ap
+import code_gen_value_util
+
+
+class ApOpLoadFromRegisterCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        out = self.get_out_cg_val(0)
+        return [out]
+
+    def get_out_cg_val(self, i):
+        register_var_name_attr = self.op_property.attributes.register_var_name
+        register_var_name = register_var_name_attr.match(a_str=lambda x: x)
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type, register_var_name
+        )
+
+
+class ApOpLoadFromGlobalCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        index_func_unique_id_attr = (
+            self.op_property.attributes.index_func_unique_id
+        )
+        index_func_unique_id = index_func_unique_id_attr.match(
+            a_str=lambda x: x
+        )
+        offset_var_name = self.index_program_translator_map.get_offset_var_name(
+            index_func_unique_id=index_func_unique_id,
+            mut_kernel_arg_id_registry=mut_kernel_arg_id_registry,
+            mut_lir_code_gen_ctx=mut_lir_code_gen_ctx,
+        )
+        data_op_name = inputs[0].var_name
+        arg_name = mut_kernel_arg_id_registry.get_in_tensor_data_ptr_var_name(
+            data_op_name
+        )
+        ptr_var_name = self.kernel_arg_translator.get_use_name(arg_name)
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"{ptr_var_name}[{offset_var_name}]")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class ApOpStoreToRegisterCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        mut_lir_code_gen_ctx.stmts.append(
+            f"{self.get_out_var_name()} = {inputs[0].var_name};"
+        )
+        return []
+
+    def get_out_var_name(self):
+        register_var_name_attr = self.op_property.attributes.register_var_name
+        return register_var_name_attr.match(a_str=lambda x: x)
+
+
+class ApOpStoreToGlobalCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+        self.ptr_type2data_type = ap.OrderedDict(
+            [
+                [ap.PointerType.const_float_ptr, "const float"],
+                [ap.PointerType.const_float16_ptr, "const half"],
+                [ap.PointerType.float_ptr, "float"],
+                [ap.PointerType.float16_ptr, "half"],
+            ]
+        )
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        index_func_unique_id_attr = (
+            self.op_property.attributes.index_func_unique_id
+        )
+        index_func_unique_id = index_func_unique_id_attr.match(
+            a_str=lambda x: x
+        )
+        offset_var_name = self.index_program_translator_map.get_offset_var_name(
+            index_func_unique_id=index_func_unique_id,
+            mut_kernel_arg_id_registry=mut_kernel_arg_id_registry,
+            mut_lir_code_gen_ctx=mut_lir_code_gen_ctx,
+        )
+        arg_name = mut_kernel_arg_id_registry.get_out_tensor_data_ptr_var_name(
+            index_func_unique_id
+        )
+        glb_data_type = self.get_glb_type(
+            mut_kernel_arg_id_registry, index_func_unique_id
+        )
+        ptr_var_name = self.kernel_arg_translator.get_use_name(arg_name)
+        mut_lir_code_gen_ctx.store(
+            glb_data_type, ptr_var_name, offset_var_name, inputs[1].var_name
+        )
+        return []
+
+    def get_glb_type(self, mut_kernel_arg_id_registry, index_func_unique_id):
+        arg_name = mut_kernel_arg_id_registry.get_out_tensor_data_ptr_var_name(
+            index_func_unique_id
+        )
+        kernel_arg_name2ids = ap.OrderedDict(
+            ap.map(
+                lambda item: [item[1], item[0]],
+                mut_kernel_arg_id_registry.generated_kernel_arg_id2unique_name.items(),
+            )
+        )
+        kernel_arg_id = kernel_arg_name2ids[arg_name]
+        return self.ptr_type2data_type[kernel_arg_id.type]
+
+    def get_out_var_name(self):
+        register_var_name_attr = self.op_property.attributes.register_var_name
+        return register_var_name_attr.match(a_str=lambda x: x)
+
+
+class PdOpDataCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        out = self.get_out_cg_val(0)
+        return [out]
+
+    def get_out_cg_val(self, i):
+        name = self.op_property.attributes.name.match(a_str=lambda x: x)
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type, name
+        )
+
+
+class PdOpFullCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        value = self.op_property.attributes.value.match(a_f64=lambda x: x)
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"{value}")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpCastCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+        self.dtype2type_name = ap.OrderedDict(
+            [
+                [ap.DataType.float, "float"],
+                [ap.DataType.float16, "half"],
+                [ap.DataType.bfloat16, "nv_bfloat16"],
+                [ap.DataType.int32, "int"],
+                [ap.DataType.int64, "int64_t"],
+            ]
+        )
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        dtype = self.op_property.attributes.dtype.match(a_dtype=lambda x: x)
+        dtype_name = self.dtype2type_name[dtype]
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(
+            out, f"static_cast<{dtype_name}>({inputs[0].var_name})"
+        )
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpExpCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"expf({inputs[0].var_name})")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpReluCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(
+            out, f"({inputs[0].var_name} > 0 ? {inputs[0].var_name} : 0) "
+        )
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpErfCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        var_name = inputs[0].var_name
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"erf({var_name})")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpElementwisePowCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        exponent = inputs[1].var_name
+        var_name = inputs[0].var_name
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"ComputePow({var_name}, {exponent})")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpSinCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        var_name = inputs[0].var_name
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"sin({var_name})")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpTanhCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        var_name = inputs[0].var_name
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"tanh({var_name})")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpFloorCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        var_name = inputs[0].var_name
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"floor({var_name})")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class CinnOpScaleCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        scale = self.op_property.attributes.scale.match(a_f32=lambda x: x)
+        bias = self.op_property.attributes.bias.match(a_f32=lambda x: x)
+        bias_after_scale = self.op_property.attributes.bias_after_scale.match(
+            a_bool=lambda x: x
+        )
+        in_name = inputs[0].var_name
+        true_str = f"{scale} * {in_name} + {bias}"
+        false_str = f"{scale} * ({in_name} + {bias})"
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(
+            out, true_str if bias_after_scale else false_str
+        )
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpSubtractCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        a = inputs[0]
+        b = inputs[1]
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"({a.var_name} - {b.var_name})")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpAddCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        a = inputs[0]
+        b = inputs[1]
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"{a.var_name} + {b.var_name}")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpMultiplyCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        a = inputs[0]
+        b = inputs[1]
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"{a.var_name} * {b.var_name}")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpDivideCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        a = inputs[0]
+        b = inputs[1]
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(out, f"{a.var_name} / {b.var_name}")
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class PdOpMaximumCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        a = inputs[0]
+        b = inputs[1]
+        out = self.get_out_cg_val(0)
+        mut_lir_code_gen_ctx.let(
+            out,
+            f"(({a.var_name} >= {b.var_name}) ? ({a.var_name}) : ({b.var_name}))",
+        )
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class CinnOpYieldStoreCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        return inputs
+
+
+class CinnOpBroadcastCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        return inputs
+
+
+class CinnOpExpandCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        return [inputs[0]]
+
+
+class CinnOpGenerateShapeCodeGen:
+    def __init__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        out = self.get_out_cg_val(0)
+        return [out]
+
+    def get_out_cg_val(self, i):
+        return code_gen_value_util.CodeGenValue(
+            self.output_properties[i].type,
+            f"op{self.op_property.op_index}_out{i}",
+        )
+
+
+class OpComputeTranslatorFactory:
+    def __init__(self):
+        self.op_name2class = ap.OrderedDict(
+            [
+                ["ap_op.load_from_register", ApOpLoadFromRegisterCodeGen],
+                ["ap_op.load_from_global", ApOpLoadFromGlobalCodeGen],
+                ["ap_op.store_to_register", ApOpStoreToRegisterCodeGen],
+                ["ap_op.store_to_global", ApOpStoreToGlobalCodeGen],
+                ["pd_op.data", PdOpDataCodeGen],
+                ["pd_op.full", PdOpFullCodeGen],
+                ["pd_op.cast", PdOpCastCodeGen],
+                ["pd_op.exp", PdOpExpCodeGen],
+                ["pd_op.relu", PdOpReluCodeGen],
+                ["pd_op.sin", PdOpSinCodeGen],
+                ["pd_op.tanh", PdOpTanhCodeGen],
+                ["pd_op.floor", PdOpFloorCodeGen],
+                ["pd_op.erf", PdOpErfCodeGen],
+                ["pd_op.elementwise_pow", PdOpElementwisePowCodeGen],
+                ["cinn_op.scale", CinnOpScaleCodeGen],
+                ["pd_op.subtract", PdOpSubtractCodeGen],
+                ["pd_op.add", PdOpAddCodeGen],
+                ["pd_op.multiply", PdOpMultiplyCodeGen],
+                ["pd_op.divide", PdOpDivideCodeGen],
+                ["pd_op.maximum", PdOpMaximumCodeGen],
+                ["cinn_op.yield_store", CinnOpYieldStoreCodeGen],
+                ["cinn_op.broadcast", CinnOpBroadcastCodeGen],
+                ["pd_op.expand", CinnOpExpandCodeGen],
+                ["cinn_op.generate_shape", CinnOpGenerateShapeCodeGen],
+            ]
+        )
+
+    def __call__(
+        self,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        index_program_translator_map,
+    ):
+        cls = self._get_class(op_property.op_name)
+        return cls(
+            op_property=op_property,
+            input_properties=input_properties,
+            output_properties=output_properties,
+            kernel_arg_translator=kernel_arg_translator,
+            index_program_translator_map=index_program_translator_map,
+        )
+
+    def _get_class(self, op_name):
+        return self.op_name2class[op_name]

python/paddle/apy/matmul_pass/op_conversion_drr_pass.py

@@ -0,0 +1,222 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import access_topo_drr
+
+
+@access_topo_drr.register_drr_pass("pd_op_cast", tag="default")
+class PdOpCastAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.cast_op = o.ap_native_op("pd_op.cast")
+        o.cast_op([t.input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_tanh", tag="default")
+class PdOpTanhAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.tanh_op = o.ap_native_op("pd_op.tanh")
+        o.tanh_op([t.input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_floor", tag="default")
+class PdOpFloorAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.floor_op = o.ap_native_op("pd_op.floor")
+        o.floor_op([t.input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_erf", tag="default")
+class PdOpErfAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.erf_op = o.ap_native_op("pd_op.erf")
+        o.erf_op([t.input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_elementwise_pow", tag="default")
+class PdOpElementwisePowAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.source_op = o.ap_native_op("pd_op.elementwise_pow")
+        o.source_op([t.input0, t.input1], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.add")
+        o.result_op([t.input0, t.input1], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_exp", tag="default")
+class PdOpExpAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.exp_op = o.ap_native_op("pd_op.exp")
+        o.exp_op([t.input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("cinn_op_scale", tag="default")
+class CinnOpScaleAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.scale_op = o.ap_native_op("cinn_op.scale")
+        o.scale_op([t.input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_sin", tag="default")
+class PdOpSinAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.sin_op = o.ap_native_op("pd_op.sin")
+        o.sin_op([t.input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("cinn_op_yield_store", tag="default")
+class CinnOpYieldStoreAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.yield_op = o.ap_native_op("cinn_op.yield_store")
+        o.yield_op([t.input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_subtract", tag="default")
+class PdOpSubtractAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.source_op = o.ap_native_op("pd_op.subtract")
+        o.source_op([t.input0, t.input1], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.add")
+        o.result_op([t.input0, t.input1], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_divide", tag="default")
+class PdOpDivideAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.source_op = o.ap_native_op("pd_op.divide")
+        o.source_op([t.input0, t.input1], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.add")
+        o.result_op([t.input0, t.input1], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_multiply", tag="default")
+class PdOpMultiplyAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.source_op = o.ap_native_op("pd_op.multiply")
+        o.source_op([t.input0, t.input1], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.add")
+        o.result_op([t.input0, t.input1], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_maximum", tag="default")
+class PdOpMaximumAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.source_op = o.ap_native_op("pd_op.maximum")
+        o.source_op([t.input0, t.input1], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.add")
+        o.result_op([t.input0, t.input1], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_left_full_add", tag="default")
+class PdOpLeftFullAddAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.full_op = o.ap_native_op("pd_op.full")
+        o.full_op([], [t.intermediate])
+        o.source_op = o.ap_native_op("pd_op.add")
+        o.source_op([t.intermediate, t.input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_right_full_add", tag="default")
+class PdOpRightFullAddAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.full_op = o.ap_native_op("pd_op.full")
+        o.full_op([], [t.intermediate])
+        o.source_op = o.ap_native_op("pd_op.add")
+        o.source_op([t.input, t.intermediate], [t.output])
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass(
+    "full_generate_shape_expand_left_add", tag="default"
+)
+class FullGenerateShapeExpandLeftAddAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.full = o.ap_native_op("pd_op.full")
+        o.full([], [t.intermediate0])
+        o.generate_shape = o.ap_native_op("cinn_op.generate_shape")
+        o.generate_shape([t.input0], [t.intermediate1])
+        o.expand = o.ap_native_op("pd_op.expand")
+        o.expand([t.intermediate0, t.intermediate1], [t.expanded_input])
+        o.add = o.ap_native_op("pd_op.add")
+        o.add([t.expanded_input, t.input0], [t.output])
+
+    def result_pattern(self, o, t):
+        o.relu = o.ap_native_op("pd_op.relu")
+        o.relu([t.input0], [t.output])
+
+
+@access_topo_drr.register_drr_pass(
+    "full_generate_shape_expand_right_add", tag="default"
+)
+class FullGenerateShapeExpandRightAddAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.full = o.ap_native_op("pd_op.full")
+        o.full([], [t.intermediate0])
+        o.generate_shape = o.ap_native_op("cinn_op.generate_shape")
+        o.generate_shape([t.input0], [t.intermediate1])
+        o.expand = o.ap_native_op("pd_op.expand")
+        o.expand([t.intermediate0, t.intermediate1], [t.expanded_input])
+        o.add = o.ap_native_op("pd_op.add")
+        o.add([t.input0, t.expanded_input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.relu = o.ap_native_op("pd_op.relu")
+        o.relu([t.input0], [t.output])

python/paddle/apy/matmul_pass/op_index_translator_util.py

@@ -0,0 +1,226 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import ap
+import index_code_gen_value_util
+
+
+class PdOpDataCodeGen:
+    def __init__(
+        self,
+        index_program_id,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        anchor_iter_var_names,
+    ):
+        self.index_program_id = index_program_id
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.anchor_iter_var_names = anchor_iter_var_names
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        return [
+            index_code_gen_value_util.IndexCodeGenValue(
+                self.anchor_iter_var_names
+            )
+        ]
+
+
+class PdOpFullIntArrayCodeGen:
+    def __init__(
+        self,
+        index_program_id,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        anchor_iter_var_names,
+    ):
+        self.index_program_id = index_program_id
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.anchor_iter_var_names = anchor_iter_var_names
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        out = index_code_gen_value_util.IndexCodeGenValue(None)
+
+        def get_int64(attr):
+            return attr.match(a_i64=lambda x: x)
+
+        def convert_list(lst):
+            return ap.map(get_int64, lst)
+
+        out.const_data = self.op_property.attributes.value.match(
+            a_array=convert_list
+        )
+        return [out]
+
+
+class PdOpSumCodeGen:
+    def __init__(
+        self,
+        index_program_id,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        anchor_iter_var_names,
+    ):
+        self.index_program_id = index_program_id
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.anchor_iter_var_names = anchor_iter_var_names
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        input_iter_var_names = inputs[0].iter_var_names
+        reduced_axes_set = ap.OrderedDict(
+            ap.map(lambda x: [int(x), True], inputs[1].const_data)
+        )
+        non_reduced_axes = ap.filter(
+            lambda x: reduced_axes_set.contains(x) == False,  # noqa: E712
+            range(len(input_iter_var_names)),
+        )
+        output_iter_var_names = ap.map(
+            lambda i: input_iter_var_names[i], non_reduced_axes
+        )
+        return [
+            index_code_gen_value_util.IndexCodeGenValue(output_iter_var_names)
+        ]
+
+
+class CinnOpReshapeCodeGen:
+    def __init__(
+        self,
+        index_program_id,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        anchor_iter_var_names,
+    ):
+        self.index_program_id = index_program_id
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.anchor_iter_var_names = anchor_iter_var_names
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        symbolic_shape = self.input_properties[0].symbolic_shape
+
+        def get_or_create_dim_var_name(dim_expr):
+            arg_var_name = mut_kernel_arg_id_registry.get_dim_expr_var_name(
+                dim_expr
+            )
+            return self.kernel_arg_translator.get_use_name(arg_var_name)
+
+        def get_dim_var_name(i):
+            dim_expr = symbolic_shape[i]
+            return get_or_create_dim_var_name(dim_expr)
+
+        rank = len(symbolic_shape)
+        stride_dims_list = ap.map(
+            lambda num_dims: ap.map(
+                lambda i: get_dim_var_name(num_dims + i + 1),
+                range(rank - 1 - num_dims),
+            ),
+            range(rank),
+        )
+        var_name_and_dims_list = ap.map(
+            lambda pair: [pair[0], *pair[1]],
+            zip(inputs[0].iter_var_names, stride_dims_list),
+        )
+        offset_expr = " + ".join(
+            ap.map(lambda elts: " * ".join(elts), var_name_and_dims_list)
+        )
+        assert len(self.output_properties[0].symbolic_shape) == 1, (
+            "len(self.output_properties[0]) should be 1"
+        )
+        return [
+            index_code_gen_value_util.IndexCodeGenValue([f"({offset_expr})"])
+        ]
+
+
+class CfYieldCodeGen:
+    def __init__(
+        self,
+        index_program_id,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        anchor_iter_var_names,
+    ):
+        self.index_program_id = index_program_id
+        self.op_property = op_property
+        self.input_properties = input_properties
+        self.output_properties = output_properties
+        self.kernel_arg_translator = kernel_arg_translator
+        self.anchor_iter_var_names = anchor_iter_var_names
+
+    def __call__(
+        self, inputs, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        return []
+
+
+class OpIndexTranslatorFactory:
+    def __init__(self):
+        self.op_name2class = ap.OrderedDict(
+            [
+                ["pd_op.data", PdOpDataCodeGen],
+                ["pd_op.full_int_array", PdOpFullIntArrayCodeGen],
+                ["pd_op.sum", PdOpSumCodeGen],
+                ["cinn_op.reshape", CinnOpReshapeCodeGen],
+                ["cf.yield", CfYieldCodeGen],
+            ]
+        )
+
+    def __call__(
+        self,
+        index_program_id,
+        op_property,
+        input_properties,
+        output_properties,
+        kernel_arg_translator,
+        anchor_iter_var_names,
+    ):
+        cls = self._get_class(op_property.op_name)
+        return cls(
+            index_program_id=index_program_id,
+            op_property=op_property,
+            input_properties=input_properties,
+            output_properties=output_properties,
+            kernel_arg_translator=kernel_arg_translator,
+            anchor_iter_var_names=anchor_iter_var_names,
+        )
+
+    def _get_class(self, op_name):
+        return self.op_name2class[op_name]

python/paddle/apy/matmul_pass/program_translator_util.py

@@ -0,0 +1,81 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import ap
+
+
+class ProgramTranslator:
+    def __init__(
+        self,
+        program_property,
+        kernel_arg_translator,
+        index_program_translator_map,
+        op_translator_maker,
+    ):
+        self.program_property = program_property
+        self.kernel_arg_translator = kernel_arg_translator
+        self.index_program_translator_map = index_program_translator_map
+        self.op_translator_maker = op_translator_maker
+        self.ir_value_index2translated_value = ap.MutableList()
+
+        def PushNone(x):
+            self.ir_value_index2translated_value.append(None)
+
+        map(PushNone, self.program_property.values)
+
+    # mut_kernel_arg_id_registry: mutable KernelArgIdLazyContext
+    # mut_lir_code_gen_ctx: mutable low level ir code generation context
+    def translate(self, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx):
+        def TranslateOp(op_property):
+            self._translate_op(
+                op_property, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+            )
+
+        map(TranslateOp, self.program_property.ops)
+
+    def _translate_op(
+        self, op_property, mut_kernel_arg_id_registry, mut_lir_code_gen_ctx
+    ):
+        op_translator = self.op_translator_maker(
+            op_property=op_property,
+            input_properties=map(
+                self._get_value_property, op_property.input_value_indexes
+            ),
+            output_properties=map(
+                self._get_value_property, op_property.output_value_indexes
+            ),
+            kernel_arg_translator=self.kernel_arg_translator,
+            index_program_translator_map=self.index_program_translator_map,
+        )
+        inputs = map(
+            self._get_translated_value, op_property.input_value_indexes
+        )
+        outputs = op_translator(
+            inputs,
+            mut_kernel_arg_id_registry=mut_kernel_arg_id_registry,
+            mut_lir_code_gen_ctx=mut_lir_code_gen_ctx,
+        )
+        map(
+            self._set_translated_value,
+            zip(op_property.output_value_indexes, outputs),
+        )
+
+    def _get_value_property(self, i):
+        return self.program_property.values[i]
+
+    def _get_translated_value(self, i):
+        return self.ir_value_index2translated_value[i]
+
+    def _set_translated_value(self, pair):
+        self.ir_value_index2translated_value[pair[0]] = pair[1]

python/paddle/apy/matmul_pass/topo_drr_pass.py

@@ -0,0 +1,489 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import access_topo_drr
+import ap
+import pir
+
+
+class FakeDataForYieldAccessTopoPass(access_topo_drr.DrrPass):
+    def __init__(self, fake_data_names):
+        self.num_outputs = len(fake_data_names)
+        self.fake_data_names = fake_data_names
+        self.undefined_place = pir.a_place(pir.UndefinedPlace())
+
+    def source_pattern(self, o, t):
+        o.yield_op = o.ap_native_op("cf.yield")
+
+        def get_yield_input(i):
+            return getattr(t, f"output{i}")
+
+        o.yield_op(ap.map(get_yield_input, range(self.num_outputs)), [])
+
+    def result_pattern(self, o, t):
+        self.result_pattern_data_op(o, t)
+        self.result_pattern_up_spider(o, t)
+
+    def result_pattern_data_op(self, o, t):
+        ap.map(
+            lambda i: self.data_op_for_output(o, t, i), range(self.num_outputs)
+        )
+
+    def data_op_for_output(self, o, t, i):
+        t.declare_internal_native_ir_value(f"data_out{i}")
+        data_op_unique_name = f"data_op_for_output{i}"
+        setattr(o, data_op_unique_name, o.ap_native_op("pd_op.data"))
+        data_op = getattr(o, data_op_unique_name)
+        data_op.name = lambda o, t: self.get_name(o, t, i)
+        data_op.shape = lambda o, t: self.get_shape(o, t, i)
+        data_op.dtype = lambda o, t: self.get_dtype(o, t, i)
+        data_op.place = lambda o, t: self.get_place(o, t, i)
+        data_op([], [getattr(t, f"data_out{i}")])
+
+    def get_name(self, o, t, i):
+        return pir.a_str(self.fake_data_names[i])
+
+    def get_shape(self, o, t, i):
+        ir_tensor = getattr(t, f"output{i}")
+
+        def GetDims(dtype, dims, data_layout):
+            return dims
+
+        return pir.a_intarray(ir_tensor.type.match(t_dtensor=GetDims))
+
+    def get_dtype(self, o, t, i):
+        ir_tensor = getattr(t, f"output{i}")
+
+        def GetDtype(dtype, dims, data_layout):
+            return dtype
+
+        return pir.a_dtype(ir_tensor.type.match(t_dtensor=GetDtype))
+
+    def get_place(self, o, t, i):
+        return self.undefined_place
+
+    def result_pattern_up_spider(self, o, t):
+        ap.map(
+            lambda i: self.up_spider_for_output(o, t, i),
+            range(self.num_outputs),
+        )
+
+    def up_spider_for_output(self, o, t, i):
+        t.declare_internal_native_ir_value(f"add_out{i}")
+        up_spider_op_name = f"up_spider_op{i}"
+        setattr(o, up_spider_op_name, o.ap_native_op("ap_op.up_spider"))
+        getattr(o, up_spider_op_name)(
+            [getattr(t, f"output{i}"), getattr(t, f"data_out{i}")], []
+        )
+
+
+class FakeDataStoreToGlobalForYieldAccessTopoPass(access_topo_drr.DrrPass):
+    def __init__(self, fake_data_names):
+        self.num_outputs = len(fake_data_names)
+        self.fake_data_names = fake_data_names
+        self.undefined_place = pir.a_place(pir.UndefinedPlace())
+
+    def source_pattern(self, o, t):
+        o.yield_op = o.ap_native_op("cf.yield")
+
+        def get_yield_input(i):
+            return getattr(t, f"output{i}")
+
+        o.yield_op(ap.map(get_yield_input, range(self.num_outputs)), [])
+
+    def result_pattern(self, o, t):
+        self.result_pattern_data_op(o, t)
+        self.result_pattern_store_to_global_op(o, t)
+
+    def result_pattern_data_op(self, o, t):
+        ap.map(
+            lambda i: self.data_op_for_output(o, t, i), range(self.num_outputs)
+        )
+
+    def data_op_for_output(self, o, t, i):
+        t.declare_internal_native_ir_value(f"data_out{i}")
+        data_op_unique_name = f"data_op_for_output{i}"
+        setattr(o, data_op_unique_name, o.ap_native_op("pd_op.data"))
+        data_op = getattr(o, data_op_unique_name)
+        data_op.name = lambda o, t: self.get_name(o, t, i)
+        data_op.shape = lambda o, t: self.get_shape(o, t, i)
+        data_op.dtype = lambda o, t: self.get_dtype(o, t, i)
+        data_op.place = lambda o, t: self.get_place(o, t, i)
+        data_op([], [getattr(t, f"data_out{i}")])
+
+    def get_name(self, o, t, i):
+        return pir.a_str(self.fake_data_names[i])
+
+    def get_shape(self, o, t, i):
+        ir_tensor = getattr(t, f"output{i}")
+
+        def GetDims(dtype, dims, data_layout):
+            return dims
+
+        return pir.a_intarray(ir_tensor.type.match(t_dtensor=GetDims))
+
+    def get_dtype(self, o, t, i):
+        ir_tensor = getattr(t, f"output{i}")
+
+        def GetDtype(dtype, dims, data_layout):
+            return dtype
+
+        return pir.a_dtype(ir_tensor.type.match(t_dtensor=GetDtype))
+
+    def get_place(self, o, t, i):
+        return self.undefined_place
+
+    def result_pattern_store_to_global_op(self, o, t):
+        ap.map(
+            lambda i: self.store_to_global_op_for_output(o, t, i),
+            range(self.num_outputs),
+        )
+
+    def store_to_global_op_for_output(self, o, t, i):
+        store_to_global_op_name = f"store_to_global_op{i}"
+        setattr(
+            o, store_to_global_op_name, o.ap_native_op("ap_op.store_to_global")
+        )
+        store_to_global_op = getattr(o, store_to_global_op_name)
+        store_to_global_op.index_func_unique_id = lambda o, t: pir.a_str(
+            self.fake_data_names[i]
+        )
+        store_to_global_op(
+            [getattr(t, f"data_out{i}"), getattr(t, f"output{i}")], []
+        )
+
+
+class ConvertUpSpiderStoreDataOpToYieldOpPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.data_op = o.ap_native_op("pd_op.data")
+        o.data_op([], [t.input1])
+        o.load_from_global_op = o.ap_native_op("ap_op.load_from_global")
+        o.load_from_global_op([t.input1], [t.tmp1])
+        o.up_spider_op = o.ap_native_op("ap_op.up_spider")
+        o.up_spider_op([t.input0, t.tmp1], [])
+
+    def result_pattern(self, o, t):
+        o.yield_op = o.ap_native_op("cf.yield")
+        o.yield_op([t.input0], [])
+
+
+class ConvertDownSpiderStoreDataOpToYieldOpPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.data_mm_op = o.ap_native_op("pd_op.data")
+        o.data_mm_op([], [t.input1])
+        o.down_spider_op = o.ap_native_op("ap_op.down_spider")
+        o.down_spider_op([t.input1], [t.tmp1])
+        o.store_to_global = o.ap_native_op("ap_op.store_to_global")
+        o.store_to_global([t.input0, t.tmp1], [])
+
+    def result_pattern(self, o, t):
+        o.yield_op = o.ap_native_op("cf.yield")
+        o.yield_op([t.input0], [])
+
+
+class InitDownSpiderAccessTopoPass(access_topo_drr.DrrPass):
+    def __init__(self, data_input_name):
+        self.data_input_name_attr = pir.a_str(data_input_name)
+
+    def source_pattern(self, o, t):
+        o.data_op = o.ap_native_op("pd_op.data")
+        o.data_op([], [t.output])
+
+    def constraint(self, o, t):
+        return o.data_op.name == self.data_input_name_attr
+
+    def result_pattern(self, o, t):
+        t.declare_internal_native_ir_value("input")
+        o.new_data_op = o.ap_native_op("pd_op.data")
+        o.new_data_op.name = lambda o, t: o.data_op.name
+        o.new_data_op.shape = lambda o, t: o.data_op.shape
+        o.new_data_op.dtype = lambda o, t: o.data_op.dtype
+        o.new_data_op.place = lambda o, t: o.data_op.place
+        o.new_data_op([], [t.input])
+        o.down_spider = o.ap_native_op("ap_op.down_spider")
+        o.down_spider([t.input], [t.output])
+
+
+class InitNaiveLoadFromGlobalAccessTopoPass(access_topo_drr.DrrPass):
+    def __init__(self, data_input_name):
+        self.data_input_name_attr = pir.a_str(data_input_name)
+
+    def source_pattern(self, o, t):
+        o.data_op = o.ap_native_op("pd_op.data")
+        o.data_op([], [t.output])
+
+    def constraint(self, o, t):
+        return o.data_op.name == self.data_input_name_attr
+
+    def result_pattern(self, o, t):
+        t.declare_internal_native_ir_value("input")
+        o.new_data_op = o.ap_native_op("pd_op.data")
+        o.new_data_op.name = lambda o, t: o.data_op.name
+        o.new_data_op.shape = lambda o, t: o.data_op.shape
+        o.new_data_op.dtype = lambda o, t: o.data_op.dtype
+        o.new_data_op.place = lambda o, t: o.data_op.place
+        o.new_data_op([], [t.input])
+        o.load_from_global = o.ap_native_op("ap_op.load_from_global")
+        o.load_from_global.index_func_unique_id = lambda o, t: (
+            self.data_input_name_attr
+        )
+        o.load_from_global([t.input], [t.output])
+
+
+class ReplaceWithLoadFromRegisterPass(access_topo_drr.DrrPass):
+    def __init__(self, name, register_var_name):
+        self.name = pir.a_str(name)
+        self.register_var_name = pir.a_str(register_var_name)
+
+    def source_pattern(self, o, t):
+        o.data_op = o.ap_native_op("pd_op.data")
+        o.data_op([], [t.input])
+        o.load_from_global = o.ap_native_op("ap_op.load_from_global")
+        o.load_from_global.index_func_unique_id = self.name
+        o.load_from_global([t.input], [t.output])
+
+    def result_pattern(self, o, t):
+        o.load_from_register = o.ap_native_op("ap_op.load_from_register")
+        o.load_from_register.name = lambda o, t: self.name
+        o.load_from_register.register_var_name = lambda o, t: (
+            self.register_var_name
+        )
+        o.load_from_register.type = lambda o, t: pir.a_type(t.output.type)
+        o.load_from_register.symbolic_shape_or_data = lambda o, t: pir.a_symbol(
+            t.output.get_symbolic_shape_or_data()
+        )
+        o.load_from_register([], [t.output])
+
+
+class ReplaceWithStoreToRegisterPass(access_topo_drr.DrrPass):
+    def __init__(self, name, register_var_name):
+        self.name = pir.a_str(name)
+        self.register_var_name = pir.a_str(register_var_name)
+
+    def source_pattern(self, o, t):
+        o.data_op = o.ap_native_op("pd_op.data")
+        o.data_op([], [t.output])
+        o.store_to_global_op = o.ap_native_op("ap_op.store_to_global")
+        o.store_to_global_op.index_func_unique_id = self.name
+        o.store_to_global_op([t.output, t.output_val], [])
+
+    def result_pattern(self, o, t):
+        o.store_to_register_op = o.ap_native_op("ap_op.store_to_register")
+        o.store_to_register_op.name = lambda o, t: self.name
+        o.store_to_register_op.register_var_name = lambda o, t: (
+            self.register_var_name
+        )
+        o.store_to_register_op([t.output_val], [])
+
+
+@access_topo_drr.register_drr_pass("down_spider_relu", tag="default")
+class DownSpiderReluAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.spider0 = o.ap_native_op("ap_op.down_spider")
+        o.spider0([t.input], [t.tmp])
+        o.relu1 = o.ap_native_op("pd_op.relu")
+        o.relu1([t.tmp], [t.output])
+
+    def result_pattern(self, o, t):
+        o.fustion_op = o.ap_native_op("ap_op.down_spider")
+        o.fustion_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass(
+    "down_spider_load_from_global", tag="default"
+)
+class DownSpiderLoadFromGlobalAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.spider0 = o.ap_native_op("ap_op.down_spider")
+        o.spider0([t.input], [t.tmp])
+        o.load_from_global_op = o.ap_native_op("ap_op.load_from_global")
+        o.load_from_global_op([t.tmp], [t.output])
+
+    def result_pattern(self, o, t):
+        o.fustion_op = o.ap_native_op("ap_op.down_spider")
+        o.fustion_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("down_spider_up_spider", tag="default")
+class DownSpiderUpSpiderAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.down_spider_op = o.ap_native_op("ap_op.down_spider")
+        o.down_spider_op([t.input], [t.tmp0])
+        o.up_spider_op = o.ap_native_op("ap_op.up_spider")
+        o.up_spider_op([t.tmp0, t.input], [])
+
+    def result_pattern(self, o, t):
+        pass
+
+
+@access_topo_drr.register_drr_pass("left_down_spider_add", tag="default")
+class LeftDownSpiderAddAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.spider = o.ap_native_op("ap_op.down_spider")
+        o.spider([t.input0], [t.tmp0])
+        o.add = o.ap_native_op("pd_op.add")
+        o.add([t.tmp0, t.input1], [t.output])
+
+    def result_pattern(self, o, t):
+        o.down_spider = o.ap_native_op("ap_op.down_spider")
+        o.down_spider([t.input0], [t.output])
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.input0, t.input1], [])
+
+
+@access_topo_drr.register_drr_pass("right_down_spider_add", tag="default")
+class RightDownSpiderAddAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.spider = o.ap_native_op("ap_op.down_spider")
+        o.spider([t.input0], [t.tmp0])
+        o.add = o.ap_native_op("pd_op.add")
+        o.add([t.input1, t.tmp0], [t.output])
+
+    def result_pattern(self, o, t):
+        o.down_spider = o.ap_native_op("ap_op.down_spider")
+        o.down_spider([t.input0], [t.output])
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.input0, t.input1], [])
+
+
+@access_topo_drr.register_drr_pass("expand_up_spider", tag="default")
+class ExpandUpSpiderAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.expand = o.ap_native_op("pd_op.expand")
+        o.expand([t.input1, t.input2], [t.expanded_input])
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.input0, t.expanded_input], [])
+
+    def constraint(self, o, t):
+        input_shape = t.input1.symbolic_shape_to_list()
+        output_shape = t.expanded_input.symbolic_shape_to_list()
+        rank_diff = len(output_shape) - len(input_shape)
+        return rank_diff > 0
+
+        # TODO: Get Inner Expanded Axes
+        def GetInnerExpanded_axes(i):
+            if input_shape[i] == output_shape[i + rank_diff]:
+                return []
+            else:
+                return [i]
+
+        input_rank = len(t.input1.symbolic_shape_to_list())
+        inner_expanded_axes = ap.flat_map(
+            GetInnerExpanded_axes, range(input_rank)
+        )
+        return rank_diff > 0 and len(inner_expanded_axes) == 0
+
+    def result_pattern(self, o, t):
+        t.declare_internal_native_ir_value("reduced_input")
+        o.sum = o.ap_native_op("pd_op.sum")
+        o.sum.axis = self.get_axis
+        o.sum.keepdim = self.get_keepdim
+        o.sum([t.input0], [t.reduced_input])
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.reduced_input, t.input1], [])
+
+    def get_keepdim(self, o, t):
+        return pir.a_bool(False)
+
+    def get_axis(self, o, t):
+        input_rank = len(t.input1.symbolic_shape_to_list())
+        output_rank = len(t.expanded_input.symbolic_shape_to_list())
+        axes = range(output_rank - input_rank)
+        return pir.a_intarray(axes)
+
+
+@access_topo_drr.register_drr_pass("cinn_broadcast_up_spider", tag="default")
+class CinnBroadcastUpSpiderAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.broadcast_op = o.ap_native_op("cinn_op.broadcast")
+        o.broadcast_op([t.input1], [t.expanded_input])
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.input0, t.expanded_input], [])
+
+    def constraint(self, o, t):
+        input_shape = t.input1.symbolic_shape_to_list()
+        output_shape = t.expanded_input.symbolic_shape_to_list()
+        rank_diff = len(output_shape) - len(input_shape)
+        return rank_diff > 0
+
+        # TODO: Get Inner Expanded Axes
+        def GetInnerExpanded_axes(i):
+            if input_shape[i] == output_shape[i + rank_diff]:
+                return []
+            else:
+                return [i]
+
+        input_rank = len(t.input1.symbolic_shape_to_list())
+        inner_expanded_axes = ap.flat_map(
+            GetInnerExpanded_axes, range(input_rank)
+        )
+        return rank_diff > 0 and len(inner_expanded_axes) == 0
+
+    def result_pattern(self, o, t):
+        t.declare_internal_native_ir_value("reduced_input")
+        o.sum = o.ap_native_op("pd_op.sum")
+        o.sum.axis = self.get_axis
+        o.sum.keepdim = self.get_keepdim
+        o.sum([t.input0], [t.reduced_input])
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.reduced_input, t.input1], [])
+
+    def get_keepdim(self, o, t):
+        return pir.a_bool(False)
+
+    def get_axis(self, o, t):
+        input_rank = len(t.input1.symbolic_shape_to_list())
+        output_rank = len(t.expanded_input.symbolic_shape_to_list())
+        axes = range(output_rank - input_rank)
+        return pir.a_intarray(axes)
+
+
+@access_topo_drr.register_drr_pass("right_down_spider_up_spider", tag="default")
+class RightDownSpiderUpSpiderAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.expand = o.ap_native_op("ap_op.down_spider")
+        o.expand([t.input1], [t.output1])
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.input0, t.output1], [])
+
+    def result_pattern(self, o, t):
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.input0, t.input1], [])
+
+
+@access_topo_drr.register_drr_pass("left_down_spider_up_spider", tag="default")
+class LeftDownSpiderUpSpiderAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.expand = o.ap_native_op("ap_op.down_spider")
+        o.expand([t.input0], [t.output0])
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.output0, t.input1], [])
+
+    def result_pattern(self, o, t):
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.input0, t.input1], [])
+
+
+@access_topo_drr.register_drr_pass(
+    "triangle_left_down_spider_up_spider", tag="default"
+)
+class TriangleLeftDownSpiderUpSpiderAccessTopoPass(access_topo_drr.DrrPass):
+    def source_pattern(self, o, t):
+        o.expand = o.ap_native_op("ap_op.down_spider")
+        o.expand([t.input0], [t.output0])
+        o.up_spider = o.ap_native_op("ap_op.up_spider")
+        o.up_spider([t.input0, t.output0], [])
+
+    def result_pattern(self, o, t):
+        pass

python/paddle/apy/matmul_pass/umprime.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import access_topo_drr
+import ap
+import pir
+
+
+@access_topo_drr.register_drr_pass("pd_op_static_relu", tag="umprime")
+class PdOpReluAccessTopoPass(access_topo_drr.DrrPass):
+    def __init__(self):
+        self.zero = pir.a_f64(ap.DataValue.float64("0"))
+
+    def source_pattern(self, o, t):
+        o.full_op = o.ap_native_op("pd_op.full")
+        o.full_op([], [t.intermediate])
+        o.maximum_op = o.ap_native_op("pd_op.maximum")
+        o.maximum_op([t.input, t.intermediate], [t.output])
+
+    def constraint(self, o, t):
+        return o.full_op.value == self.zero
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input], [t.output])
+
+
+@access_topo_drr.register_drr_pass("pd_op_dynamic_relu", tag="umprime")
+class PdOpDynReluAccessTopoPass(access_topo_drr.DrrPass):
+    def __init__(self):
+        self.zero = pir.a_f64(ap.DataValue.float64("0"))
+
+    def source_pattern(self, o, t):
+        o.full_op = o.ap_native_op("pd_op.full")
+        o.full_op([], [t.intermediate0])
+        o.generate_shape_op = o.ap_native_op("cinn_op.generate_shape")
+        o.generate_shape_op([t.input0], [t.intermediate1])
+        o.expand_op = o.ap_native_op("pd_op.expand")
+        o.expand_op([t.intermediate0, t.intermediate1], [t.intermediate2])
+        o.maximum_op = o.ap_native_op("pd_op.maximum")
+        o.maximum_op([t.input1, t.intermediate2], [t.output])
+
+    def constraint(self, o, t):
+        return o.full_op.value == self.zero
+
+    def result_pattern(self, o, t):
+        o.result_op = o.ap_native_op("pd_op.relu")
+        o.result_op([t.input1], [t.output])

python/paddle/apy/sys/__builtin__.py

@@ -0,0 +1,40 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+DataType = __builtin__DataType  # noqa: F821
+DataValue = __builtin__DataValue  # noqa: F821
+PointerType = __builtin__PointerType  # noqa: F821
+PointerValue = __builtin__PointerValue  # noqa: F821
+MutableList = __builtin__MutableList  # noqa: F821
+OrderedDict = __builtin__OrderedDict  # noqa: F821
+MutableOrderedDict = __builtin__MutableOrderedDict  # noqa: F821
+AttrMap = __builtin__AttrMap  # noqa: F821
+SerializableAttrMap = __builtin__BuiltinSerializableAttrMap  # noqa: F821
+
+_raise = __builtin__raise  # noqa: F821
+
+foreach = __builtin__foreach  # noqa: F821
+range = __builtin__range  # noqa: F821
+map = __builtin__map  # noqa: F821
+reduce = __builtin__reduce  # noqa: F821
+filter = __builtin__filter  # noqa: F821
+zip = __builtin__zip  # noqa: F821
+flat_map = __builtin__flat_map  # noqa: F821
+apply = __builtin__apply  # noqa: F821
+replace_or_trim_left_comma = __builtin__replace_or_trim_left_comma  # noqa: F821
+
+sorted = __builtin__sorted  # noqa: F821
+
+dirname = __builtin__dirname  # noqa: F821
+basename = __builtin__basename  # noqa: F821

python/paddle/apy/sys/__builtin_registry_item__.py

@@ -0,0 +1,62 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import __builtin__
+
+
+class RegistryEntry:
+    def __init__(self):
+        self.__tag_name__ = None
+        self.__nice__ = None
+        self.__values__ = __builtin__.MutableList()
+        self.__child_register_item_name2value__ = (
+            __builtin__.MutableOrderedDict()
+        )
+
+    # tag_name: str
+    # nice: int
+    def __getattr__(self, attrname):
+        def contains():
+            return self.__child_register_item_name2value__.contains(attrname)
+
+        def find():
+            return self.__child_register_item_name2value__[attrname]
+
+        def create():
+            register_entry = RegistryEntry()
+            self.__child_register_item_name2value__[attrname] = register_entry
+            return register_entry
+
+        return find() if contains() else create()
+
+    def __call__(self, tag_name, nice):
+        registry_obj = RegistryObject(tag_name, nice)
+        self.__values__.append(registry_obj)
+        return RegisterItemDecorator(registry_obj)
+
+
+class RegistryObject:
+    def __init__(self, tag_name, nice):
+        self.tag_name = tag_name
+        self.nice = nice
+        self.value = None
+
+
+class RegisterItemDecorator:
+    def __init__(self, register_obj):
+        self.register_obj = register_obj
+
+    def __call__(self, value):
+        self.register_obj.value = value
+        return value

python/paddle/apy/sys/__builtin_trivial_op_names__.py

@@ -0,0 +1,40 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+def GetGroupedTrivialOpNames():
+    return [
+        "pd_op.sin",
+        "pd_op.add",
+        "pd_op.relu",
+        "pd_op.data",
+        "pd_op.full",
+        "pd_op.cast",
+        "pd_op.exp",
+        "pd_op.relu",
+        "pd_op.tanh",
+        "pd_op.floor",
+        "pd_op.erf",
+        "pd_op.elementwise_pow",
+        "cinn_op.scale",
+        "pd_op.subtract",
+        "pd_op.add",
+        "pd_op.multiply",
+        "pd_op.divide",
+        "pd_op.maximum",
+        "cinn_op.yield_store",
+        "cinn_op.broadcast",
+        "pd_op.expand",
+        "cinn_op.generate_shape",
+    ]

python/paddle/apy/sys/ap.py

@@ -0,0 +1,44 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import __builtin__
+
+DataType = __builtin__.DataType
+DataValue = __builtin__.DataValue
+PointerType = __builtin__.PointerType
+PointerValue = __builtin__.PointerValue
+MutableList = __builtin__.MutableList
+OrderedDict = __builtin__.OrderedDict
+MutableOrderedDict = __builtin__.MutableOrderedDict
+AttrMap = __builtin__.AttrMap
+SerializableAttrMap = __builtin__.SerializableAttrMap
+
+_raise = __builtin__._raise
+
+foreach = __builtin__.foreach
+range = __builtin__.range
+map = __builtin__.map
+reduce = __builtin__.reduce
+filter = __builtin__.filter
+zip = __builtin__.zip
+flat_map = __builtin__.flat_map
+apply = __builtin__.apply
+replace_or_trim_left_comma = __builtin__.replace_or_trim_left_comma
+
+registry = __builtin__registry  # noqa: F821
+
+sorted = __builtin__.sorted
+
+dirname = __builtin__.dirname
+basename = __builtin__.basename