Skip to content

Commit

Permalink
[PIR] Add op_callstack to Pir (#62139)
Browse files Browse the repository at this point in the history
---------

Co-authored-by: SigureMo <sigure.qaq@gmail.com>
  • Loading branch information
xingmingyyj and SigureMo authored Mar 6, 2024
1 parent c7b3acf commit 376aba5
Show file tree
Hide file tree
Showing 6 changed files with 210 additions and 17 deletions.
11 changes: 9 additions & 2 deletions paddle/cinn/hlir/framework/pir/utils.cc
Original file line number Diff line number Diff line change
Expand Up @@ -465,9 +465,16 @@ static utils::Attribute ConvertArrayAttribute(
CASE_ATTRIBUTE(float, FloatAttribute)
} else if (attr_vec[0].isa<::pir::DoubleAttribute>()) {
CASE_ATTRIBUTE(double, DoubleAttribute)
} else if (attr_vec[0].isa<::pir::StrAttribute>()) {
std::vector<std::string> dst_attr;
for (auto element : attr_vec) {
dst_attr.push_back(
element.dyn_cast<::pir::StrAttribute>().AsString());
}
} else {
LOG(FATAL) << "only support bool/int32/int64/float/double attribute in "
"ArrayAttribute";
LOG(FATAL)
<< "only support bool/int32/int64/float/double/string attribute in "
"ArrayAttribute";
}
}
} else if (src_attr.isa<::pir::shape::SymbolAttribute>()) {
Expand Down
12 changes: 10 additions & 2 deletions paddle/fluid/pir/dialect/op_generator/python_c_gen.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@
#include "paddle/fluid/pybind/op_function_common.h"
#include "paddle/phi/common/int_array.h"
#include "paddle/phi/core/enforce.h"
#include "paddle/fluid/pybind/op_callstack_utils.h"
{body}
Expand All @@ -71,8 +72,10 @@
{attrs}
// Call ir static api
CallStackRecorder callstack_recoder("{api_name}");
callstack_recoder.Record();
auto static_api_out = paddle::dialect::{api_name}({args});
callstack_recoder.AttachToOps();
return ToPyObject(static_api_out);
}} catch (...) {{
ThrowExceptionToPython(std::current_exception());
Expand All @@ -94,8 +97,10 @@
{attrs}
// Call ir static api
CallStackRecorder callstack_recoder("{api_name}");
callstack_recoder.Record();
paddle::dialect::{api_name}({args});
callstack_recoder.AttachToOps();
return nullptr;
}} catch (...) {{
ThrowExceptionToPython(std::current_exception());
Expand Down Expand Up @@ -129,7 +134,10 @@
{cast_attrs}
// Call ir static api
CallStackRecorder callstack_recoder("{api_name}");
callstack_recoder.Record();
auto static_api_out = paddle::dialect::{api_name}({args_with_mutable_attrs});
callstack_recoder.AttachToOps();
return ToPyObject(static_api_out);
Expand Down
3 changes: 2 additions & 1 deletion paddle/fluid/pybind/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -151,7 +151,8 @@ set(PYBIND_SRCS
auto_parallel_py.cc
eval_frame_tools.cc
cpython_internals.c
eval_frame.c)
eval_frame.c
op_callstack_utils.cc)

if(NOT WITH_SHARED_IR)
# Note: We want to compile pir source into paddle.so directly, because
Expand Down
66 changes: 54 additions & 12 deletions paddle/fluid/pybind/manual_static_op_function.h
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
#include "paddle/fluid/pybind/eager_utils.h"
#include "paddle/fluid/pybind/exception.h"
#include "paddle/fluid/pybind/op_callstack_utils.h"
#include "paddle/fluid/pybind/op_function_common.h"
#include "paddle/phi/common/int_array.h"
#include "paddle/phi/core/enforce.h"
Expand All @@ -43,8 +44,10 @@ static PyObject *static_api_parameter(PyObject *self,
PyObject *name_obj = PyTuple_GET_ITEM(args, 0);
std::string name = CastPyArg2String(name_obj, "name", 0);
// Call ir static api
CallStackRecorder callstack_recoder("parameter");
callstack_recoder.Record();
auto static_api_out = paddle::dialect::parameter(name);

callstack_recoder.AttachToOps();
return ToPyObject(static_api_out);
} catch (...) {
ThrowExceptionToPython(std::current_exception());
Expand All @@ -67,8 +70,10 @@ static PyObject *static_api_set_parameter(PyObject *self,
PyObject *name_obj = PyTuple_GET_ITEM(args, 1);
std::string name = CastPyArg2String(name_obj, "name", 1);
// Call ir static api
CallStackRecorder callstack_recoder("set_parameter");
callstack_recoder.Record();
paddle::dialect::set_parameter(parameter, name);

callstack_recoder.AttachToOps();
Py_RETURN_NONE;
} catch (...) {
ThrowExceptionToPython(std::current_exception());
Expand All @@ -91,8 +96,10 @@ static PyObject *static_api_set_persistable_value(PyObject *self,
PyObject *name_obj = PyTuple_GET_ITEM(args, 1);
std::string name = CastPyArg2String(name_obj, "name", 1);
// Call ir static api
CallStackRecorder callstack_recoder("shadow_output");
callstack_recoder.Record();
paddle::dialect::shadow_output(persist_value, name);

callstack_recoder.AttachToOps();
Py_RETURN_NONE;
} catch (...) {
ThrowExceptionToPython(std::current_exception());
Expand All @@ -119,7 +126,10 @@ PyObject *static_api_full(PyObject *self, PyObject *args, PyObject *kwargs) {
!PyObject_CheckIRValue(value_obj)) {
std::vector<int64_t> shape = CastPyArg2Longs(shape_obj, "full", 0);
float value = CastPyArg2Float(value_obj, "full", 1);
CallStackRecorder callstack_recoder("full");
callstack_recoder.Record();
auto static_api_out = paddle::dialect::full(shape, value, dtype, place);
callstack_recoder.AttachToOps();
return ToPyObject(static_api_out);
} else {
pir::Value shape, value;
Expand All @@ -146,8 +156,12 @@ PyObject *static_api_full(PyObject *self, PyObject *args, PyObject *kwargs) {
phi::CPUPlace());
}

CallStackRecorder callstack_recoder("full_with_tensor");
callstack_recoder.Record();
auto static_api_out =
paddle::dialect::full_with_tensor(shape, value, dtype);
callstack_recoder.AttachToOps();

return ToPyObject(static_api_out);
}
} catch (...) {
Expand All @@ -169,7 +183,10 @@ static PyObject *static_api_create_array(PyObject *self,
CastPyArg2DataTypeDirectly(dtype_obj, "create_array", 0);

// Call ir static api
CallStackRecorder callstack_recoder("create_array");
callstack_recoder.Record();
auto static_api_out = paddle::dialect::create_array(dtype);
callstack_recoder.AttachToOps();

return ToPyObject(static_api_out);
} catch (...) {
Expand All @@ -194,8 +211,10 @@ static PyObject *static_api_create_array_like(PyObject *self,
float value = CastPyArg2Float(value_obj, "create_array_like", 1);

// Call ir static api
CallStackRecorder callstack_recoder("create_array_like");
callstack_recoder.Record();
auto static_api_out = paddle::dialect::create_array_like(input, value);

callstack_recoder.AttachToOps();
return ToPyObject(static_api_out);
} catch (...) {
ThrowExceptionToPython(std::current_exception());
Expand All @@ -215,7 +234,10 @@ static PyObject *static_api_array_length(PyObject *self,
auto x = CastPyArg2Value(x_obj, "array_length", 0);

// Call ir static api
CallStackRecorder callstack_recoder("array_length");
callstack_recoder.Record();
auto static_api_out = paddle::dialect::array_length(x);
callstack_recoder.AttachToOps();

return ToPyObject(static_api_out);
} catch (...) {
Expand Down Expand Up @@ -248,7 +270,10 @@ static PyObject *static_api_array_read(PyObject *self,
}

// Call ir static api
CallStackRecorder callstack_recoder("array_read");
callstack_recoder.Record();
auto static_api_out = paddle::dialect::array_read(array, i);
callstack_recoder.AttachToOps();

return ToPyObject(static_api_out);
} catch (...) {
Expand Down Expand Up @@ -282,7 +307,10 @@ static PyObject *static_api_array_write_(PyObject *self,
}

// Call ir static api
CallStackRecorder callstack_recoder("array_write_");
callstack_recoder.Record();
auto static_api_out = paddle::dialect::array_write_(array, x, i);
callstack_recoder.AttachToOps();

return ToPyObject(static_api_out);
} catch (...) {
Expand Down Expand Up @@ -321,7 +349,10 @@ static PyObject *static_api_array_to_tensor(PyObject *self,
auto use_stack = CastPyArg2Boolean(use_stack_obj, "array_to_tensor", 2);

// Call ir static api
CallStackRecorder callstack_recoder("array_to_tensor");
callstack_recoder.Record();
auto static_api_out = paddle::dialect::array_to_tensor(x, axis, use_stack);
callstack_recoder.AttachToOps();

return ToPyObject(static_api_out);
} catch (...) {
Expand All @@ -341,10 +372,10 @@ PyObject *static_api_add_n_array(PyObject *self,
PyObject *inputs_obj = PyTuple_GET_ITEM(args, 0);
auto inputs = CastPyArg2VectorOfValue(inputs_obj, "add_n", 0);

// Parse Attributes

// Call ir static api
CallStackRecorder callstack_recoder("add_n_array");
callstack_recoder.Record();
auto static_api_out = paddle::dialect::add_n_array(inputs);
callstack_recoder.AttachToOps();

return ToPyObject(static_api_out);
} catch (...) {
Expand Down Expand Up @@ -395,7 +426,10 @@ static PyObject *static_api_slice_array(PyObject *self,
}

// Call ir static api
CallStackRecorder callstack_recoder("slice_array");
callstack_recoder.Record();
auto static_api_out = paddle::dialect::slice_array(input, starts, ends);
callstack_recoder.AttachToOps();

return ToPyObject(static_api_out);
} catch (...) {
Expand Down Expand Up @@ -430,9 +464,11 @@ static PyObject *static_api_slice_array_dense(PyObject *self,
starts = paddle::dialect::full_int_array(
starts_tmp, phi::DataType::INT64, phi::CPUPlace());
}

// Call ir static api
CallStackRecorder callstack_recoder("slice_array_dense");
callstack_recoder.Record();
auto static_api_out = paddle::dialect::slice_array_dense(input, starts);
callstack_recoder.AttachToOps();

return ToPyObject(static_api_out);
} catch (...) {
Expand Down Expand Up @@ -754,7 +790,8 @@ static PyObject *static_api_run_custom_op(PyObject *self,

argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
::pir::PassStopGradientsDefaultly(argument);

CallStackRecorder callstack_recoder("run_custom_op");
callstack_recoder.Record();
std::vector<pir::Value> op_results;
pir::Operation *op =
paddle::dialect::ApiBuilder::Instance().GetBuilder()->Build(
Expand All @@ -772,7 +809,7 @@ static PyObject *static_api_run_custom_op(PyObject *self,
op_results.push_back(op->result(i));
}
}

callstack_recoder.AttachToOps();
return ToPyObject(op_results);
}

Expand Down Expand Up @@ -811,10 +848,13 @@ static PyObject *static_api_fused_gemm_epilogue(PyObject *self,
PyObject *activation_obj = PyTuple_GET_ITEM(args, 5);
std::string activation =
CastPyArg2String(activation_obj, "fused_gemm_epilogue", 5);

// Call ir static api
CallStackRecorder callstack_recoder("fused_gemm_epilogue");
callstack_recoder.Record();
auto out = paddle::dialect::fused_gemm_epilogue(
x, y, bias, trans_x, trans_y, activation);
callstack_recoder.AttachToOps();

return ToPyObject(out);
} catch (...) {
ThrowExceptionToPython(std::current_exception());
Expand All @@ -836,8 +876,10 @@ static PyObject *static_api_array_pop(PyObject *self,
auto index = CastPyArg2Int(index_obj, "array_pop", 1);

// Call ir static api
CallStackRecorder callstack_recoder("array_pop");
callstack_recoder.Record();
auto static_api_out = paddle::dialect::array_pop(input, index);

callstack_recoder.AttachToOps();
return ToPyObject(static_api_out);
} catch (...) {
ThrowExceptionToPython(std::current_exception());
Expand Down
104 changes: 104 additions & 0 deletions paddle/fluid/pybind/op_callstack_utils.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,104 @@
// Copyright (c) 2024 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 <Python.h>
#include <frameobject.h>

#include "paddle/fluid/framework/op_proto_maker.h"
#include "paddle/fluid/pir/dialect/operator/ir/api_builder.h"
#include "paddle/fluid/platform/enforce.h"
#include "paddle/fluid/pybind/op_callstack_utils.h"

pir::Attribute CallStackRecorder::GetOpCallstackInfo() {
PyObject* traceback_str = PyUnicode_FromString("traceback");
PyObject* traceback_module = PyImport_Import(traceback_str);

if (NULL == traceback_module) {
Py_DECREF(traceback_str);
Py_DECREF(traceback_module);
PADDLE_THROW(paddle::platform::errors::PreconditionNotMet(
"Failed to import traceback module while getting callstack information "
"for %s.",
api_name_));
}
PyObject* tb = PyObject_GetAttrString(traceback_module, "extract_stack");
PyObject* stack = PyObject_CallObject(tb, NULL);
if (NULL == stack) {
Py_DECREF(tb);
Py_DECREF(traceback_str);
Py_DECREF(traceback_module);
PADDLE_THROW(paddle::platform::errors::PreconditionNotMet(
"Failed to get callstack object while getting callstack information "
"for "
"%s.",
api_name_));
}
Py_ssize_t stack_size = PyList_Size(stack);
std::vector<pir::Attribute> op_callstack_infos;
for (Py_ssize_t i = 0; i < stack_size; ++i) {
PyObject* frame_summary = PyList_GetItem(stack, i);
PyObject* filename = PyObject_GetAttrString(frame_summary, "filename");
PyObject* lineno = PyObject_GetAttrString(frame_summary, "lineno");
PyObject* name = PyObject_GetAttrString(frame_summary, "name");
PyObject* line = PyObject_GetAttrString(frame_summary, "line");
PyObject* callstack_info = PyUnicode_FromFormat(
" File \"%S\", line %S, in %S", filename, lineno, name);
PyObject* callstack_source_line = PyUnicode_FromFormat(" %S", line);
op_callstack_infos.push_back(
pir::StrAttribute::get(pir::IrContext::Instance(),
std::string(PyUnicode_AsUTF8(callstack_info))));
op_callstack_infos.push_back(pir::StrAttribute::get(
pir::IrContext::Instance(),
std::string(PyUnicode_AsUTF8(callstack_source_line))));
Py_DECREF(callstack_info);
Py_DECREF(callstack_source_line);
Py_DECREF(filename);
Py_DECREF(lineno);
Py_DECREF(name);
Py_DECREF(line);
}
Py_DECREF(tb);
Py_DECREF(traceback_str);
Py_DECREF(traceback_module);
return pir::ArrayAttribute::get(pir::IrContext::Instance(),
op_callstack_infos);
}

void CallStackRecorder::Record() {
auto before_insertion_point =
paddle::dialect::ApiBuilder::Instance().GetCurrentInsertionPoint();
before_insertion_iterator_ = (--before_insertion_point.second);
before_insertion_block_ = before_insertion_point.first;
}

void CallStackRecorder::AttachToOps() {
before_insertion_iterator_++;
pir::Attribute callstack_info_attr = GetOpCallstackInfo();
pir::InsertionPoint after_insertion_point =
paddle::dialect::ApiBuilder::Instance().GetCurrentInsertionPoint();
PADDLE_ENFORCE_EQ(before_insertion_block_,
after_insertion_point.first,
paddle::platform::errors::PreconditionNotMet(
"The block obtained before and after calling the "
"static API %s is inconsistent.",
api_name_));
auto after_insertion_iterator = after_insertion_point.second;
for (auto block_iterator = before_insertion_iterator_;
block_iterator != after_insertion_iterator;
block_iterator++) {
block_iterator->set_attribute(paddle::framework::OpProtoAndCheckerMaker::
OpCreationCallstackAttrName(),
callstack_info_attr);
}
}
Loading

0 comments on commit 376aba5

Please sign in to comment.