[CustomPass] add register_subgraph_pass api

[ronny1996]

PR types

Others

PR changes

Others

Description

#55446

import paddle.fluid as fluid
import paddle
import os
import numpy as np

paddle.enable_static()

def func(x, y, z):
    return x + y + z

@paddle.jit.to_static(
    input_spec=[
        paddle.static.InputSpec(([1, 1]), "int64", "x"),
        paddle.static.InputSpec([1, 1], "int64", "y"),
        paddle.static.InputSpec([1, 1], "int64", "z"),
        paddle.static.InputSpec([1, 1], "int64", "limit"),
    ]
)


def test_while(x, y, z, limit):
    res = func(x, y, z) # Pre-while call for inference
    while (paddle.squeeze(res, 1) < paddle.squeeze(limit, 1)):
        res = func(x, y, res)

    return res

MODLE_FILE = "./test_while"

def setUp():
    for lib in os.listdir(os.getenv("CUSTOM_DEVICE_ROOT")):
        if lib.endswith(".so"):
            paddle.utils.cpp_extension.extension_utils.load_op_meta_info_and_register_op(
                lib
            )
    
    paddle.jit.save(test_while, MODLE_FILE)

@paddle.incubate.passes.ir.RegisterPass
def generate_add_n():
    def pattern(x, y, z):
        return paddle.add(paddle.add(x, y), z)

    def replace(x, y, z):
        return paddle.incubate.passes.ir.PassDesc.OP.my_add_n(X=x, Y=y, Z=z)

    return pattern, replace

paddle.fluid.core.register_subgraph_pass("generate_add_n")

def main():
    setUp()
    config = paddle.inference.Config()
    config.set_prog_file(MODLE_FILE + ".pdmodel")
    config.enable_memory_optim()
    config.enable_custom_device("npu")
    pass_builder = config.pass_builder()
    pass_builder.append_pass("generate_add_n")
    predictor = paddle.inference.create_predictor(config)

    np_inputs = [
        np.ones((1, 1)).astype("int64"), # x
        np.ones((1, 1)).astype("int64"), # y
        np.ones((1, 1)).astype("int64"), # z
        20 * np.ones((1, 1)).astype("int64"), # limit
    ]
    input_names = predictor.get_input_names()
    for i, name in enumerate(input_names):
        input_tensor = predictor.get_input_handle(name)
        input_tensor.copy_from_cpu(np_inputs[i])
    print("################start predictor run")
    predictor.run()
    results = []
    output_names = predictor.get_output_names()
    for i, name in enumerate(output_names):
        output_tensor = predictor.get_output_handle(name)
        output_data = output_tensor.copy_to_cpu()
        results.append(output_data)
    print(results)

if __name__ == '__main__':
    main() # 替换Pass，并predictor.run

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[qili93]

LGTM