Skip to content

Commit

Permalink
[PIR]Add Fused dropout add pass (PaddlePaddle#58272)
Browse files Browse the repository at this point in the history
* tmp

* fix flashattn compile bug

* modify unsqueeze

* modify part complex bug

* modify

* modify cpu

* [PIR]Migrate maximum into pir

* Polish code

* add ir_grad of static_gradient

* add test

* modify bug

* modify

* test_with_pir

* close one test

* add_math_op_patch

* modify

* modify

* add mean fill_constant test

* modify cpu int32 test

* get_shape_tensor

* delete

* add default place

* add dropout unittest

* Update test/legacy_test/test_elementwise_div_op.py

* modify review comment

* modify add test

* modify logic and sub

* modify

* add fused_gemm_epilogue pass

* polish code

* udate cmake

* fused_dropout_add_pass

* modify cmake

* modify

* Update test/cpp/pir/pattern_rewrite/CMakeLists.txt

* Update paddle/fluid/pir/transforms/fusion/fused_dropout_add_pass.cc

Co-authored-by: Yuanle Liu <yuanlehome@163.com>

* Update paddle/fluid/pir/transforms/fusion/fused_dropout_add_pass.cc

Co-authored-by: Yuanle Liu <yuanlehome@163.com>

* Update paddle/fluid/pybind/pir.cc

* Update test/ir/new_ir/fused_pass/test_fused_dropout_add.py

* Update test/ir/new_ir/fused_pass/test_fused_dropout_add.py

* Update paddle/fluid/pir/transforms/fusion/fused_dropout_add_pass.cc

* add ci converage

---------

Co-authored-by: zyfncg <zhangyunfei07@baidu.com>
Co-authored-by: 0x45f <wangzhen45@baidu.com>
Co-authored-by: Yuanle Liu <yuanlehome@163.com>
  • Loading branch information
4 people authored and jiahy0825 committed Oct 26, 2023
1 parent e1c882a commit 734cd84
Show file tree
Hide file tree
Showing 11 changed files with 320 additions and 13 deletions.
15 changes: 14 additions & 1 deletion paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,20 @@
OPS_API_TEMPLATE = """
{{"{name}", (PyCFunction)(void (*)(void)){name}, METH_VARARGS | METH_KEYWORDS, "C++ interface function for {name}."}},"""

NEED_GEN_STATIC_ONLY_APIS = ['fetch']
NEED_GEN_STATIC_ONLY_APIS = [
'fetch',
'fused_embedding_eltwise_layernorm',
'fused_fc_elementwise_layernorm',
'fused_multi_transformer_xpu',
'fused_scale_bias_relu_conv_bnstats',
'fusion_transpose_flatten_concat',
'generate_sequence_xpu',
'layer_norm_act_xpu',
'multi_encoder_xpu',
'multihead_matmul',
'squeeze_excitation_block',
'yolo_box_xpu',
]

NO_NEED_GEN_STATIC_ONLY_APIS = [
'add_n_',
Expand Down
15 changes: 7 additions & 8 deletions paddle/fluid/pir/dialect/operator/ir/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -98,7 +98,7 @@ set(api_source_file_tmp ${api_source_file}.tmp)
add_custom_command(
OUTPUT ${api_header_file} ${api_source_file}
COMMAND
${PYTHON_EXECUTABLE} ${api_gen_file} --op_yaml_files ${api_gen_yaml_files}
${PYTHON_EXECUTABLE} ${api_gen_file} --op_yaml_files ${op_yaml_files}
--op_compat_yaml_file ${op_compat_yaml_file} --namespaces ${op_namespace}
--api_def_h_file ${api_header_file_tmp} --api_def_cc_file
${api_source_file_tmp}
Expand Down Expand Up @@ -129,10 +129,9 @@ set(python_c_source_file_tmp ${python_c_source_file}.tmp)
add_custom_command(
OUTPUT ${python_c_header_file} ${python_c_source_file}
COMMAND
${PYTHON_EXECUTABLE} ${python_c_gen_file} --op_yaml_files
${api_gen_yaml_files} --op_compat_yaml_file ${op_compat_yaml_file}
--namespaces "paddle,pybind" --python_c_def_h_file
${python_c_header_file_tmp} --python_c_def_cc_file
${PYTHON_EXECUTABLE} ${python_c_gen_file} --op_yaml_files ${op_yaml_files}
--op_compat_yaml_file ${op_compat_yaml_file} --namespaces "paddle,pybind"
--python_c_def_h_file ${python_c_header_file_tmp} --python_c_def_cc_file
${python_c_source_file_tmp}
COMMAND ${CMAKE_COMMAND} -E copy_if_different ${python_c_header_file_tmp}
${python_c_header_file}
Expand Down Expand Up @@ -160,9 +159,9 @@ set(ops_api_source_file_tmp ${ops_api_source_file}.tmp)
add_custom_command(
OUTPUT ${ops_api_source_file}
COMMAND
${PYTHON_EXECUTABLE} ${ops_api_gen_file} --op_yaml_files
${api_gen_yaml_files} --op_compat_yaml_file ${op_compat_yaml_file}
--namespaces "paddle,pybind" --ops_api_file ${ops_api_source_file_tmp}
${PYTHON_EXECUTABLE} ${ops_api_gen_file} --op_yaml_files ${op_yaml_files}
--op_compat_yaml_file ${op_compat_yaml_file} --namespaces "paddle,pybind"
--ops_api_file ${ops_api_source_file_tmp}
COMMAND ${CMAKE_COMMAND} -E copy_if_different ${ops_api_source_file_tmp}
${ops_api_source_file}
COMMENT "copy_if_different ${ops_api_source_file}"
Expand Down
146 changes: 146 additions & 0 deletions paddle/fluid/pir/transforms/fusion/fused_dropout_add_pass.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,146 @@
// Copyright (c) 2023 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.

#include "paddle/fluid/pir/transforms/fusion/fused_dropout_add_pass.h"

#include "paddle/fluid/pir/drr/api/drr_pattern_base.h"
#include "paddle/pir/pass/pass.h"
#include "paddle/pir/pass/pass_registry.h"
#include "paddle/pir/pattern_rewrite/pattern_rewrite_driver.h"

namespace {

class FusedDropoutAddPattern
: public pir::drr::DrrPatternBase<FusedDropoutAddPattern> {
public:
void operator()(pir::drr::DrrPatternContext *ctx) const override {
pir::drr::SourcePattern pat = ctx->SourcePattern();
const auto &dropout = pat.Op("pd_op.dropout",
{{"p", pat.Attr("p")},
{"is_test", pat.Attr("is_test")},
{"mode", pat.Attr("mod")},
{"seed", pat.Attr("seed")},
{"fix_seed", pat.Attr("fix_seed")}});
const auto &add = pat.Op("pd_op.add");

dropout({&pat.Tensor("x"), &pat.Tensor("seed_tensor")},
{&pat.Tensor("dropout_out"), &pat.Tensor("mask")});
pat.Tensor("add_out") = add(pat.Tensor("dropout_out"), pat.Tensor("y"));

pir::drr::ResultPattern res = pat.ResultPattern();
const auto &fused_dropout_add =
res.Op("pd_op.fused_dropout_add",
{{{"p", pat.Attr("p")},
{"is_test", pat.Attr("is_test")},
{"mode", pat.Attr("mod")},
{"seed", pat.Attr("seed")},
{"fix_seed", pat.Attr("fix_seed")}}});
fused_dropout_add(
{&res.Tensor("x"), &res.Tensor("y"), &res.Tensor("seed_tensor")},
{&res.Tensor("add_out"), &res.Tensor("mask")});
}
};

class FusedDropoutGradAddGradPattern
: public pir::drr::DrrPatternBase<FusedDropoutAddPattern> {
public:
void operator()(pir::drr::DrrPatternContext *ctx) const override {
pir::drr::SourcePattern pat = ctx->SourcePattern();
const auto &dropout = pat.Op("pd_op.dropout",
{{"p", pat.Attr("p")},
{"is_test", pat.Attr("is_test")},
{"mode", pat.Attr("mod")},
{"seed", pat.Attr("seed")},
{"fix_seed", pat.Attr("fix_seed")}});
const auto &add = pat.Op("pd_op.add");

const auto &add_grad = pat.Op("pd_op.add_grad");
const auto &dropout_grad = pat.Op("pd_op.dropout_grad",
{{"p", pat.Attr("p")},
{"is_test", pat.Attr("is_test")},
{"mode", pat.Attr("mod")}});

dropout({&pat.Tensor("x"), &pat.Tensor("seed_tensor")},
{&pat.Tensor("dropout_out"), &pat.Tensor("mask")});
pat.Tensor("add_out") = add(pat.Tensor("dropout_out"), pat.Tensor("y"));
add_grad({&pat.Tensor("dropout_out"),
&pat.Tensor("y"),
&pat.Tensor("add_out_grad")},
{&pat.Tensor("dropout_out_grad"), &pat.Tensor("y_grad")});
dropout_grad({&pat.Tensor("mask"), &pat.Tensor("dropout_out_grad")},
{&pat.Tensor("x_grad")});

pir::drr::ResultPattern res = pat.ResultPattern();
const auto &fused_dropout_add =
res.Op("pd_op.fused_dropout_add",
{{{"p", pat.Attr("p")},
{"is_test", pat.Attr("is_test")},
{"mode", pat.Attr("mod")},
{"seed", pat.Attr("seed")},
{"fix_seed", pat.Attr("fix_seed")}}});

const auto &fused_dropout_add_grad =
res.Op("pd_op.fused_dropout_add_grad",
{{{"p", pat.Attr("p")},
{"is_test", pat.Attr("is_test")},
{"mode", pat.Attr("mod")},
{"fix_seed", pat.Attr("fix_seed")}}});

fused_dropout_add(
{&res.Tensor("x"), &res.Tensor("y"), &res.Tensor("seed_tensor")},
{&res.Tensor("add_out"), &res.Tensor("mask")});
fused_dropout_add_grad({&res.Tensor("mask"), &res.Tensor("add_out_grad")},
{&res.Tensor("x_grad"), &res.Tensor("y_grad")});
}
};

class FusedDropoutAddPass : public pir::Pass {
public:
FusedDropoutAddPass() : pir::Pass("fused_dropout_add_pass", 1) {}

bool Initialize(pir::IrContext *context) override {
pir::RewritePatternSet ps(context);

ps.Add(FusedDropoutAddPattern().Build(context));
ps.Add(FusedDropoutGradAddGradPattern().Build(context));
patterns_ = pir::FrozenRewritePatternSet(std::move(ps));
return true;
}

void Run(pir::Operation *op) override {
pir::GreedyRewriteConfig cfg;
cfg.use_top_down_traversal = true;
cfg.max_iterations = 10;
pir::ApplyPatternsGreedily(op->region(0), patterns_, cfg);
}

bool CanApplyOn(pir::Operation *op) const override {
return op->isa<::pir::ModuleOp>() && op->num_regions() > 0;
}

private:
pir::FrozenRewritePatternSet patterns_;
};

} // namespace

namespace pir {

std::unique_ptr<Pass> CreateFusedDropoutAddPass() {
return std::make_unique<FusedDropoutAddPass>();
}

} // namespace pir

REGISTER_IR_PASS(fused_dropout_add_pass, FusedDropoutAddPass);
26 changes: 26 additions & 0 deletions paddle/fluid/pir/transforms/fusion/fused_dropout_add_pass.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
// Copyright (c) 2023 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 <memory>
#include "paddle/pir/core/dll_decl.h"

namespace pir {

class Pass;

IR_API std::unique_ptr<Pass> CreateFusedDropoutAddPass();

} // namespace pir
1 change: 1 addition & 0 deletions paddle/fluid/pybind/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ set(PYBIND_DEPS
pd_op_dialect
program_translator
pd_inplace_pass
fusion_passes
pir
new_profiler
jit_layer
Expand Down
2 changes: 2 additions & 0 deletions paddle/fluid/pybind/pir.cc
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
#include "paddle/fluid/pir/dialect/operator/ir/pd_api.h"
#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
#include "paddle/fluid/pir/transforms/fusion/fused_dropout_add_pass.h"
#include "paddle/fluid/pir/transforms/inplace_pass.h"
#include "paddle/phi/core/enforce.h"
#include "paddle/pir/core/block.h"
Expand Down Expand Up @@ -66,6 +67,7 @@ using pir::Value;
using pybind11::return_value_policy;

USE_PASS(dead_code_elimination_pass);
USE_PASS(fused_dropout_add_pass);
USE_PASS(inplace_pass);

PHI_DECLARE_bool(print_ir);
Expand Down
4 changes: 2 additions & 2 deletions python/paddle/autograd/ir_backward.py
Original file line number Diff line number Diff line change
Expand Up @@ -105,12 +105,12 @@ def prepare_grad_outputs(grad_outputs, outputs, state):
if output.shape != grad.shape:
raise ValueError(
"The shape of grad_output[%d] %s should be the same as the shape of output[%d] %s"
% (i, str(output.shape), i, str(grad.shape))
% (i, str(grad.shape), i, str(output.shape))
)
if output.dtype != grad.dtype:
raise ValueError(
"The dtype of grad_output[%d] %s should be the same as the dtype of output[%d] %s"
% (i, str(output.dtype), i, str(grad.dtype))
% (i, str(grad.dtype), i, str(output.dtype))
)
feedop = grad.get_defining_op()
update_bwdop_structure(
Expand Down
4 changes: 2 additions & 2 deletions python/paddle/incubate/nn/functional/fused_dropout_add.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
from paddle import _C_ops
from paddle.base import core
from paddle.common_ops_import import default_main_program
from paddle.framework import LayerHelper, in_dynamic_mode
from paddle.framework import LayerHelper, in_dynamic_or_pir_mode


def fused_dropout_add(
Expand Down Expand Up @@ -84,7 +84,7 @@ def fused_dropout_add(
"mode argument should be 'downscale_in_infer' or 'upscale_in_train'"
)
seed = None
if in_dynamic_mode():
if in_dynamic_or_pir_mode():
if default_main_program().random_seed != 0:
seed = default_main_program().random_seed
out, seed_offset = _C_ops.fused_dropout_add(
Expand Down
2 changes: 2 additions & 0 deletions test/ir/new_ir/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -19,3 +19,5 @@ foreach(target ${TEST_IR_SYSTEM_CASES})
endforeach()

set_tests_properties(test_pd_inplace_pass PROPERTIES TIMEOUT 60)

add_subdirectory(fused_pass)
5 changes: 5 additions & 0 deletions test/ir/new_ir/fused_pass/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
file(
GLOB TEST_INTERP_CASES
RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
"test_*.py")
string(REPLACE ".py" "" TEST_INTERP_CASES "${TEST_INTERP_CASES}")
Loading

0 comments on commit 734cd84

Please sign in to comment.