[AutoParallel] reshard record event for time line (#59159)

wanghuancoder · web-flow · commit 9c15e019847f · 2023-11-21T10:50:29.000+08:00
* reshard record event for time line
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/nd_mesh_reshard_function.h b/paddle/phi/core/distributed/auto_parallel/reshard/nd_mesh_reshard_function.h
@@ -28,6 +28,8 @@ class SameNdMeshReshardFunction final : public ReshardFunction {
             const DistTensor& in,
             const TensorDistAttr& out_dist_attr,
             DistTensor* out) override;
+
+  std::string Name() override { return "SameNdMeshReshard"; }
 };
 
 }  // namespace distributed
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/p_to_r_reshard_function.h b/paddle/phi/core/distributed/auto_parallel/reshard/p_to_r_reshard_function.h
@@ -31,6 +31,8 @@ class PToRReshardFunction final : public ReshardFunction {
             const DistTensor& in,
             const TensorDistAttr& out_dist_attr,
             DistTensor* out) override;
+
+  std::string Name() override { return "PToRReshard"; }
 };
 
 }  // namespace distributed
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/p_to_s_reshard_function.h b/paddle/phi/core/distributed/auto_parallel/reshard/p_to_s_reshard_function.h
@@ -27,6 +27,8 @@ class PToSReshardFunction final : public ReshardFunction {
             const DistTensor& in,
             const TensorDistAttr& out_dist_attr,
             DistTensor* out) override;
+
+  std::string Name() override { return "PToSReshard"; }
 };
 
 }  // namespace distributed
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/r_to_p_reshard_function.h b/paddle/phi/core/distributed/auto_parallel/reshard/r_to_p_reshard_function.h
@@ -28,6 +28,8 @@ class RToPReshardFunction final : public ReshardFunction {
             const DistTensor& in,
             const TensorDistAttr& out_dist_attr,
             DistTensor* out) override;
+
+  std::string Name() override { return "RToPReshard"; }
 };
 
 }  // namespace distributed
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/r_to_s_reshard_function.h b/paddle/phi/core/distributed/auto_parallel/reshard/r_to_s_reshard_function.h
@@ -28,6 +28,8 @@ class RToSReshardFunction final : public ReshardFunction {
             const DistTensor& in,
             const TensorDistAttr& out_dist_attr,
             DistTensor* out) override;
+
+  std::string Name() override { return "RToSReshard"; }
 };
 
 class RToSReshardFunctionCrossMesh final : public ReshardFunction {
@@ -39,6 +41,8 @@ class RToSReshardFunctionCrossMesh final : public ReshardFunction {
             const DistTensor& in,
             const TensorDistAttr& out_dist_attr,
             DistTensor* out) override;
+
+  std::string Name() override { return "RToSReshardCrossMesh"; }
 };
 
 }  // namespace distributed
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.cc b/paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.cc
@@ -14,6 +14,7 @@
 
 #include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.h"
 
+#include "paddle/phi/api/profiler/event_tracing.h"
 #include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
 #include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
 
@@ -24,6 +25,8 @@ std::shared_ptr<DistTensor> ReshardFunction::Eval(
     DeviceContext* dev_ctx,
     const DistTensor& in,
     const TensorDistAttr& out_dist_attr) {
+  phi::RecordEvent reshard_record_event(
+      Name(), phi::TracerEventType::OperatorInner, 1);
   std::shared_ptr<DistTensor> out = std::make_shared<DistTensor>();
   Eval(dev_ctx, in, out_dist_attr, out.get());
   return out;
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.h b/paddle/phi/core/distributed/auto_parallel/reshard/reshard_function.h
@@ -43,6 +43,8 @@ class ReshardFunction {
                     const TensorDistAttr& out_dist_attr,
                     DistTensor* out) = 0;
 
+  virtual std::string Name() { return "ReshardBase"; }
+
  protected:
   void SetValue(DistTensor* tensor, const DenseTensor& value);
   void SetDistProps(DistTensor* tensor,
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/s_to_r_reshard_function.h b/paddle/phi/core/distributed/auto_parallel/reshard/s_to_r_reshard_function.h
@@ -31,6 +31,8 @@ class SToRReshardFunction final : public ReshardFunction {
             const DistTensor& in,
             const TensorDistAttr& out_dist_attr,
             DistTensor* out) override;
+
+  std::string Name() override { return "SToRReshard"; }
 };
 
 class SToRReshardFunctionCrossMesh final : public ReshardFunction {
@@ -42,6 +44,8 @@ class SToRReshardFunctionCrossMesh final : public ReshardFunction {
             const DistTensor& in,
             const TensorDistAttr& out_dist_attr,
             DistTensor* out) override;
+
+  std::string Name() override { return "SToRReshardCrossMesh"; }
 };
 
 }  // namespace distributed
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/s_to_s_reshard_function.h b/paddle/phi/core/distributed/auto_parallel/reshard/s_to_s_reshard_function.h
@@ -31,6 +31,8 @@ class SToSReshardFunction final : public ReshardFunction {
             const DistTensor& in,
             const TensorDistAttr& out_dist_attr,
             DistTensor* out) override;
+
+  std::string Name() override { return "SToSReshard"; }
 };
 
 }  // namespace distributed
diff --git a/paddle/phi/core/distributed/auto_parallel/reshard/same_status_reshard_function.h b/paddle/phi/core/distributed/auto_parallel/reshard/same_status_reshard_function.h
@@ -28,6 +28,8 @@ class SameStatusReshardFunction final : public ReshardFunction {
             const DistTensor& in,
             const TensorDistAttr& out_dist_attr,
             DistTensor* out) override;
+
+  std::string Name() override { return "SameStatusReshard"; }
 };
 
 }  // namespace distributed
