[PIR] Add op_callstack to Pir (#62139)

xingmingyyj · SigureMo · web-flow · commit 376aba57d037 · 2024-03-06T14:14:28.000+08:00
---------

Co-authored-by: SigureMo &lt;sigure.qaq@gmail.com&gt;
diff --git a/paddle/cinn/hlir/framework/pir/utils.cc b/paddle/cinn/hlir/framework/pir/utils.cc
@@ -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>()) {
diff --git a/paddle/fluid/pir/dialect/op_generator/python_c_gen.py b/paddle/fluid/pir/dialect/op_generator/python_c_gen.py
@@ -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}
@@ -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());
@@ -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());
@@ -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);
 
 
diff --git a/paddle/fluid/pybind/CMakeLists.txt b/paddle/fluid/pybind/CMakeLists.txt
@@ -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
diff --git a/paddle/fluid/pybind/manual_static_op_function.h b/paddle/fluid/pybind/manual_static_op_function.h
@@ -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"
@@ -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());
@@ -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());
@@ -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());
@@ -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;
@@ -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 (...) {
@@ -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 (...) {
@@ -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());
@@ -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 (...) {
@@ -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 (...) {
@@ -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 (...) {
@@ -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 (...) {
@@ -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 (...) {
@@ -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 (...) {
@@ -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 (...) {
@@ -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(
@@ -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);
 }
 
@@ -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());
@@ -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());
diff --git a/paddle/fluid/pybind/op_callstack_utils.cc b/paddle/fluid/pybind/op_callstack_utils.cc
@@ -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
+//
+//     https://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);
+  }
+}
diff --git a/paddle/fluid/pybind/op_callstack_utils.h b/paddle/fluid/pybind/op_callstack_utils.h
