diff --git a/paddle/fluid/ir/dialect/pd_op.yaml b/paddle/fluid/ir/dialect/pd_op.yaml index 790d6f923480bacf5d42a724059c6e203dfc1de6..c81e1cbcccbaf54df52ec26b97bf7cb9a5567a54 100644 --- a/paddle/fluid/ir/dialect/pd_op.yaml +++ b/paddle/fluid/ir/dialect/pd_op.yaml @@ -18,6 +18,7 @@ data_transform: {} attrs: - {typename: str, name: name} + - {typename: int, name: col} outputs: - {typename: Tensor, name: out, optional: false, intermediate: false} no_need_buffer: null diff --git a/paddle/fluid/ir/phi_kernel_adaptor/phi_kernel_util.cc b/paddle/fluid/ir/phi_kernel_adaptor/phi_kernel_util.cc index fcf74c2efabcebd8525fb6794d9a73c785a6e87c..0da1fdc1462952faa9470116ad1cfab62e4167fa 100644 --- a/paddle/fluid/ir/phi_kernel_adaptor/phi_kernel_util.cc +++ b/paddle/fluid/ir/phi_kernel_adaptor/phi_kernel_util.cc @@ -62,8 +62,9 @@ void BuildScope(ir::Block* block, for (size_t i = 0; i < input_num; ++i) { auto var = scope->Var("fetch"); auto fetch_list = var->GetMutable(); - // for now only support one fetch - fetch_list->resize(1); + int index = + (*it)->attributes().at("col").dyn_cast().data(); + fetch_list->resize(index + 1); } continue; } @@ -148,7 +149,11 @@ void BuildScope(ir::Block* block, } auto var = scope->Var(name); // Only support DenseTensor or Vector - if (ptr.type().isa()) { + + if (!ptr.type()) { + var->GetMutable(); + } else if (ptr.type() + .isa()) { var->GetMutable(); } else if (ptr.type().isa()) { auto tensor_array = diff --git a/paddle/fluid/ir/phi_kernel_adaptor/phi_kernel_util.h b/paddle/fluid/ir/phi_kernel_adaptor/phi_kernel_util.h index dc3938d8dd9c8f3e75776342d3ccbd21700cc5a6..bb12c7b9c895fe13ddbdd644ca4c21f65aa254ee 100644 --- a/paddle/fluid/ir/phi_kernel_adaptor/phi_kernel_util.h +++ b/paddle/fluid/ir/phi_kernel_adaptor/phi_kernel_util.h @@ -146,6 +146,19 @@ void BuildPhiContext( ctx->EmplaceBackAttr(attr_map[t].dyn_cast().data()); } else if (attr_type_name == "ir::BoolAttribute") { ctx->EmplaceBackAttr(attr_map[t].dyn_cast().data()); + } else if (attr_type_name == "ir::ArrayAttribute") { + auto array_list = attr_map[t].dyn_cast().data(); + if (array_list[0].isa()) { + std::vector vec_res; + for (size_t i = 0; i < array_list.size(); ++i) { + vec_res.push_back( + array_list[0].dyn_cast().data()); + } + ctx->EmplaceBackAttr(vec_res); + } else { + PADDLE_THROW(phi::errors::Unimplemented("attr type not support [%s] ", + attr_type_name)); + } } else if (attr_type_name == "paddle::dialect::PlaceAttribute") { ctx->EmplaceBackAttr( attr_map[t].dyn_cast().data()); @@ -166,14 +179,22 @@ void BuildPhiContext( // process fetch op auto fetch_var = scope->Var("fetch"); auto* fetch_list = fetch_var->GetMutable(); - auto* out_tensor = &(PADDLE_GET(phi::DenseTensor, fetch_list->at(0))); + int index = + op->attributes().at("col").dyn_cast().data(); + auto* out_tensor = &(PADDLE_GET(phi::DenseTensor, fetch_list->at(index))); ctx->EmplaceBackOutput(out_tensor); } else { for (size_t i = 0; i < op->num_results(); ++i) { ir::Value out_ptr = op->result(i); auto name = name_map.at(out_ptr); - ctx->EmplaceBackOutput(OutType(const_cast( - &(scope->Var(name)->Get())))); + if (out_ptr.type()) { + ctx->EmplaceBackOutput(OutType(const_cast( + &(scope->Var(name)->Get())))); + } else { + phi::DenseTensor* ptr = nullptr; + OutType out_ptr(ptr); + ctx->EmplaceBackOutput(out_ptr); + } if (output_map != nullptr) { // only deal with single input for now, [todo] need support multi input diff --git a/paddle/fluid/ir/transforms/pd_op_to_kernel_pass.cc b/paddle/fluid/ir/transforms/pd_op_to_kernel_pass.cc index c60abfd012fbbf87ac872688ea5970068a5b53bf..4353b03fb9ca8cd39fa62cb7fd87ac0eb6ae5e01 100644 --- a/paddle/fluid/ir/transforms/pd_op_to_kernel_pass.cc +++ b/paddle/fluid/ir/transforms/pd_op_to_kernel_pass.cc @@ -218,7 +218,9 @@ std::unique_ptr PdOpLowerToKernelPass(ir::Program* prog) { if ((*it)->num_results() > 0) { for (size_t i = 0; i < (*it)->num_results(); ++i) { auto result_type = (*it)->result(i).type(); - if (result_type.isa()) { + if (!result_type) { + op_output_types.push_back(result_type); + } else if (result_type.isa()) { auto allocated_dense_tensor_dtype = paddle::dialect::AllocatedDenseTensorType::get( ctx, diff --git a/paddle/fluid/ir_adaptor/translator/op_translator.cc b/paddle/fluid/ir_adaptor/translator/op_translator.cc index e1bf7ecfe2e8444bd638cb6cdd6e010732d24d23..b12b067e72b5d9035e82cd0b55d41d240d3029b3 100644 --- a/paddle/fluid/ir_adaptor/translator/op_translator.cc +++ b/paddle/fluid/ir_adaptor/translator/op_translator.cc @@ -916,6 +916,8 @@ struct FetchOpTranscriber : public OpTranscriber { OpOutputTypeList op_output_types; ir::AttributeMap attribute_map = { {"name", ir::StrAttribute::get(ctx, op_desc.InputArgumentNames()[0])}, + {"col", + ir::Int32Attribute::get(ctx, op_desc.GetAttrIfExists("col"))}, }; op_output_types.push_back(op_inputs[0].type()); diff --git a/paddle/phi/kernels/cpu/fetch_kernel.cc b/paddle/phi/kernels/cpu/fetch_kernel.cc index 5880b7687a9f17807da4a87f8aedf21255b5a2dc..6d8a5aec2855d818433312090e146225ecaa0c47 100644 --- a/paddle/phi/kernels/cpu/fetch_kernel.cc +++ b/paddle/phi/kernels/cpu/fetch_kernel.cc @@ -34,6 +34,11 @@ PD_REGISTER_KERNEL(fetch, double, int, int64_t, + uint8_t, + int8_t, + int16_t, + phi::float16, + phi::bfloat16, phi::dtype::complex, phi::dtype::complex, bool) {} diff --git a/test/cpp/ir/pattern_rewrite/pattern_rewrite_test.cc b/test/cpp/ir/pattern_rewrite/pattern_rewrite_test.cc index 8a8a73b09345965d651db2a43952607111a3f5ce..c61836928a401272694a77e6589813f1017602c0 100644 --- a/test/cpp/ir/pattern_rewrite/pattern_rewrite_test.cc +++ b/test/cpp/ir/pattern_rewrite/pattern_rewrite_test.cc @@ -414,7 +414,7 @@ void BuildProgram(ir::Builder &builder) { // NOLINT auto transpose2_op = builder.Build( transpose1_op.out(), std::vector{0, 3, 1, 2}); - builder.Build(transpose2_op.out(), "out"); + builder.Build(transpose2_op.out(), "out", 0); } // TODO(wilber): Add a normal test. diff --git a/test/ir/new_ir/test_standalone_new_ir.py b/test/ir/new_ir/test_standalone_new_ir.py index f5aeb6ce075e85916eae04a8082e4ae33a8160f1..ca924467d4e97bc8b55624b6e527c48b9d061c6f 100644 --- a/test/ir/new_ir/test_standalone_new_ir.py +++ b/test/ir/new_ir/test_standalone_new_ir.py @@ -89,5 +89,35 @@ class TestFeedOp(unittest.TestCase): np.testing.assert_array_equal(out[0], gold_res) +class TestAddGradOp(unittest.TestCase): + def test_with_new_ir(self): + place = paddle.CPUPlace() + exe = paddle.static.Executor(place) + + main_program = paddle.static.Program() + new_scope = paddle.static.Scope() + with paddle.static.scope_guard(new_scope): + with paddle.static.program_guard(main_program): + x = paddle.static.data("x", [2, 2], dtype="float32") + y = paddle.static.data("y", [2, 2], dtype="float32") + x.stop_gradient = False + + z = x * y + + paddle.static.gradients(z, x) + + np_a = np.random.rand(2, 2).astype("float32") + np_b = np.random.rand(2, 2).astype("float32") + out = exe.run( + main_program, + feed={"x": np_a, "y": np_b}, + fetch_list=[z.name], + ) + + gold_res = np_a * np_b + + np.testing.assert_array_equal(out[0], gold_res) + + if __name__ == "__main__": unittest.main()