wehub/paddlepaddle--paddle
1
0
Fork 0
Code Issues Pull Requests Actions 2 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
ff97fb6869c2eeecb7b4c09af41ce69f43f2c574
paddlepaddle--paddle/python/paddle/apy/matmul_pass/matmul_variadic_ptn.py
T
wehub-resource-sync e25996e7db chore: import upstream snapshot with attribution
2026-07-13 12:40:42 +08:00

674 lines
21 KiB
Python
Raw Blame History

# 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
)
Reference in New Issue View Git Blame Copy Permalink