add immutable_layout_trait (#63349)

kangguangli · web-flow · commit 30b8732ff828 · 2024-04-11T15:01:17.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml
@@ -698,6 +698,7 @@
   args : (str name, int col)
   output : Tensor(out)
   interfaces : paddle::dialect::InferSymbolicShapeInterface
+  traits: pir::ImmutableLayoutTrait
 
 - op : fetch
   args : (Tensor x, str name, int col)
@@ -708,7 +709,7 @@
   kernel :
     func : fetch
     param : [x]
-  traits : pir::SideEffectTrait
+  traits : pir::SideEffectTrait, pir::ImmutableLayoutTrait
   interfaces : paddle::dialect::InferSymbolicShapeInterface
 
 - op : floor_divide
diff --git a/paddle/pir/include/core/builtin_op.h b/paddle/pir/include/core/builtin_op.h
@@ -87,7 +87,8 @@ class IR_API SetParameterOp : public pir::Op<SetParameterOp, SideEffectTrait> {
 /// \brief ShdowOutputOp: ShdowOutputOp(OpOperand, {StrAttribute,
 /// StrAttribute})
 ///
-class IR_API ShadowOutputOp : public pir::Op<ShadowOutputOp, SideEffectTrait> {
+class IR_API ShadowOutputOp
+    : public pir::Op<ShadowOutputOp, SideEffectTrait, ImmutableLayoutTrait> {
  public:
   using Op::Op;
   static const char *name() { return "builtin.shadow_output"; }
diff --git a/paddle/pir/include/core/op_trait.h b/paddle/pir/include/core/op_trait.h
@@ -121,6 +121,14 @@ class IR_API SideEffectTrait : public OpTraitBase<SideEffectTrait> {
   using Base::Base;
 };
 
+///
+/// \brief This trait marks the op's layout can't be modified.
+///
+class IR_API ImmutableLayoutTrait : public OpTraitBase<ImmutableLayoutTrait> {
+ public:
+  using Base::Base;
+};
+
 }  // namespace pir
 
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(pir::SameOperandsShapeTrait)
@@ -131,3 +139,4 @@ IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(pir::SameOperandsAndResultTypeTrait)
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(pir::SameTypeOperandsTrait)
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(pir::OneResultTrait)
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(pir::SideEffectTrait)
+IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(pir::ImmutableLayoutTrait)
diff --git a/paddle/pir/src/core/op_trait.cc b/paddle/pir/src/core/op_trait.cc
@@ -288,3 +288,4 @@ IR_DEFINE_EXPLICIT_TYPE_ID(pir::SameOperandsAndResultTypeTrait)
 IR_DEFINE_EXPLICIT_TYPE_ID(pir::SameTypeOperandsTrait)
 IR_DEFINE_EXPLICIT_TYPE_ID(pir::OneResultTrait)
 IR_DEFINE_EXPLICIT_TYPE_ID(pir::SideEffectTrait)
+IR_DEFINE_EXPLICIT_TYPE_ID(pir::ImmutableLayoutTrait)
diff --git a/test/cpp/pir/operator/layout_transformation_interface_test.cc b/test/cpp/pir/operator/layout_transformation_interface_test.cc
@@ -21,6 +21,7 @@
 #include "paddle/pir/include/core/builtin_dialect.h"
 #include "paddle/pir/include/core/builtin_type.h"
 #include "paddle/pir/include/core/dialect.h"
+#include "paddle/pir/include/core/op_trait.h"
 #include "paddle/pir/include/core/type.h"
 
 TEST(layout_transformation_interface_test, operator) {
@@ -55,3 +56,24 @@ TEST(layout_transformation_interface_test, operator) {
   EXPECT_EQ(layout_transformation_iface.RelevantOutputs(fused_conv).size(),
             fused_conv->results().size());
 }
+
+TEST(immutable_layout_trait_test, operator) {
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
+
+  pir::Program program(ctx);
+  pir::Builder builder(ctx, program.block());
+
+  auto build_input_value = [&](std::vector<int64_t> shape = {2, 2}) {
+    auto uniform = builder.Build<paddle::dialect::UniformOp>(
+        shape, phi::DataType::FLOAT32, 0.0, 1.0, 2, phi::CPUPlace());
+    return uniform;
+  };
+
+  auto out = builder.Build<pir::ShadowOutputOp>(
+      build_input_value(std::vector<int64_t>{2, 2, 2, 2}).out(), "test");
+  EXPECT_TRUE(out->HasTrait<pir::ImmutableLayoutTrait>());
+
+  auto immutable_layout_trait = out->dyn_cast<pir::ImmutableLayoutTrait>();
+  EXPECT_TRUE(immutable_layout_trait);
+}
