[AP] Add paddle.cc.ap.FacadeOp: a custom op machanism for ap pass only

paddle/ap/CMakeLists.txt

@@ -52,6 +52,13 @@ cc_library(
   SRCS ${ap_pir_srcs}
   DEPS ${AP_COMMON_DEPS} ${ap_pir_deps})
 
+file(GLOB_RECURSE ap_hlir_srcs "src/paddle/hlir/*.cc")
+set(ap_hlir_deps axpr ap_drr ap_pir)
+cc_library(
+  ap_hlir
+  SRCS ${ap_hlir_srcs}
+  DEPS ${AP_COMMON_DEPS} ${ap_hlir_deps})
+
 file(GLOB_RECURSE ap_reified_drr_srcs "src/reified_drr/*.cc")
 set(ap_reified_drr_deps axpr ap_drr ap_code_module ap_code_gen)
 cc_library(
@@ -60,7 +67,7 @@ cc_library(
   DEPS ${AP_COMMON_DEPS} ${ap_reified_drr_deps})
 
 file(GLOB_RECURSE ap_pass_srcs "src/paddle/pass/*.cc")
-set(ap_pass_deps axpr ap_pir ap_drr ap_code_module ap_code_gen ap_reified_drr)
+set(ap_pass_deps axpr ap_hlir ap_drr ap_code_module ap_code_gen ap_reified_drr)
 cc_library(
   ap_pass
   SRCS ${ap_pass_srcs}

paddle/ap/include/axpr/attr_map_method_class.h

@@ -50,12 +50,35 @@ struct AttrMapMethodClass {
   }
 };
 
+template <typename ValueT>
+struct TypeImplBuiltinAttrMapMethodClass {
+  using This = TypeImplBuiltinAttrMapMethodClass;
+  using Self = TypeImpl<AttrMap<ValueT>>;
+
+  adt::Result<ValueT> Call(const Self&) { return &This::StaticConstruct; }
+
+  static adt::Result<ValueT> StaticConstruct(const ValueT&,
+                                             const std::vector<ValueT>& args) {
+    return This{}.Construct(args);
+  }
+
+  adt::Result<ValueT> Construct(const std::vector<ValueT>& args) {
+    const auto& packed_args = CastToPackedArgs(args);
+    const auto& [pos_args, kwargs] = *packed_args;
+    ADT_CHECK(pos_args->empty())
+        << adt::errors::TypeError{std::string() +
+                                  "the construct of AttrMap "
+                                  "takes no positional arguments."};
+    return kwargs;
+  }
+};
+
 template <typename ValueT>
 struct MethodClassImpl<ValueT, AttrMap<ValueT>>
     : public AttrMapMethodClass<ValueT> {};
 
 template <typename ValueT>
 struct MethodClassImpl<ValueT, TypeImpl<AttrMap<ValueT>>>
-    : public EmptyMethodClass<ValueT> {};
+    : public TypeImplBuiltinAttrMapMethodClass<ValueT> {};
 
 }  // namespace ap::axpr

paddle/ap/include/axpr/binary_func.h

@@ -23,6 +23,7 @@ namespace ap::axpr {
   _(Sub, -)                         \
   _(Mul, *)                         \
   _(Div, /)                         \
+  _(FloorDiv, /)                    \
   _(Mod, %)                         \
   _(EQ, ==)                         \
   _(NE, !=)                         \

paddle/ap/include/axpr/builtin_class_instance_method_class.h

@@ -114,6 +114,58 @@ struct MethodClassImpl<ValueT, BuiltinClassInstance<ValueT>> {
     return class_ops->Equals(self, rhs_val);
   }
 
+  adt::Result<ValueT> Add(InterpreterBase<ValueT>* interpreter,
+                          const Self& self,
+                          const ValueT& rhs_val) {
+    const auto& opt_func = GetClassAttr(self, "__add__");
+    const auto& class_attrs = self.type.class_attrs();
+    ADT_CHECK(opt_func.has_value())
+        << adt::errors::AttributeError{std::string() + class_attrs->class_name +
+                                       " class has no attribute '__add__'"};
+    std::vector<ValueT> args{rhs_val};
+    ADT_LET_CONST_REF(ret, interpreter->InterpretCall(opt_func.value(), args));
+    return ret;
+  }
+
+  adt::Result<ValueT> Sub(InterpreterBase<ValueT>* interpreter,
+                          const Self& self,
+                          const ValueT& rhs_val) {
+    const auto& opt_func = GetClassAttr(self, "__sub__");
+    const auto& class_attrs = self.type.class_attrs();
+    ADT_CHECK(opt_func.has_value())
+        << adt::errors::AttributeError{std::string() + class_attrs->class_name +
+                                       " class has no attribute '__sub__'"};
+    std::vector<ValueT> args{rhs_val};
+    ADT_LET_CONST_REF(ret, interpreter->InterpretCall(opt_func.value(), args));
+    return ret;
+  }
+
+  adt::Result<ValueT> Mul(InterpreterBase<ValueT>* interpreter,
+                          const Self& self,
+                          const ValueT& rhs_val) {
+    const auto& opt_func = GetClassAttr(self, "__mul__");
+    const auto& class_attrs = self.type.class_attrs();
+    ADT_CHECK(opt_func.has_value())
+        << adt::errors::AttributeError{std::string() + class_attrs->class_name +
+                                       " class has no attribute '__mul__'"};
+    std::vector<ValueT> args{rhs_val};
+    ADT_LET_CONST_REF(ret, interpreter->InterpretCall(opt_func.value(), args));
+    return ret;
+  }
+
+  adt::Result<ValueT> FloorDiv(InterpreterBase<ValueT>* interpreter,
+                               const Self& self,
+                               const ValueT& rhs_val) {
+    const auto& opt_func = GetClassAttr(self, "__floordiv__");
+    const auto& class_attrs = self.type.class_attrs();
+    ADT_CHECK(opt_func.has_value()) << adt::errors::AttributeError{
+        std::string() + class_attrs->class_name +
+        " class has no attribute '__floordiv__'"};
+    std::vector<ValueT> args{rhs_val};
+    ADT_LET_CONST_REF(ret, interpreter->InterpretCall(opt_func.value(), args));
+    return ret;
+  }
+
   adt::Result<ValueT> GetItem(InterpreterBase<ValueT>* interpreter,
                               const Self& self,
                               const ValueT& idx_val) {

paddle/ap/include/axpr/builtin_frame_util.h

@@ -41,6 +41,7 @@ void VisitEachBuiltinFrameAttr(const YieldT& Yield) {
   Yield("__builtin_not__", &BuiltinNot);
 
   Yield("__builtin__foreach", &ForEach);
+
   auto YieldTwice = [&](const auto& name, const auto& value) {
     Yield(name, value);
     Yield(std::string("__builtin__") + name, value);

paddle/ap/include/axpr/dim_expr_method_class.h

@@ -22,6 +22,8 @@
 #include "paddle/pir/include/dialect/shape/utils/dim_expr.h"
 
 namespace ap::axpr {
+template <typename ValueT>
+axpr::TypeImpl<axpr::BuiltinClassInstance<ValueT>> GetDimExprClass();
 
 template <typename ValueT>
 struct DimExprMethodClass {
@@ -41,6 +43,38 @@ struct DimExprMethodClass {
     return hash_value;
   }
 
+  static adt::Result<ValueT> Add(const ValueT& self_val,
+                                 const std::vector<ValueT>& args) {
+    ADT_LET_CONST_REF(lhs, self_val.template CastTo<Self>());
+    ADT_CHECK(args.size() == 1);
+    ADT_LET_CONST_REF(rhs, args.at(0).template CastTo<Self>());
+    return GetDimExprClass<ValueT>().New(lhs + rhs);
+  }
+
+  static adt::Result<ValueT> Sub(const ValueT& self_val,
+                                 const std::vector<ValueT>& args) {
+    ADT_LET_CONST_REF(lhs, self_val.template CastTo<Self>());
+    ADT_CHECK(args.size() == 1);
+    ADT_LET_CONST_REF(rhs, args.at(0).template CastTo<Self>());
+    return GetDimExprClass<ValueT>().New(lhs - rhs);
+  }
+
+  static adt::Result<ValueT> Mul(const ValueT& self_val,
+                                 const std::vector<ValueT>& args) {
+    ADT_LET_CONST_REF(lhs, self_val.template CastTo<Self>());
+    ADT_CHECK(args.size() == 1);
+    ADT_LET_CONST_REF(rhs, args.at(0).template CastTo<Self>());
+    return GetDimExprClass<ValueT>().New(lhs * rhs);
+  }
+
+  static adt::Result<ValueT> FloorDiv(const ValueT& self_val,
+                                      const std::vector<ValueT>& args) {
+    ADT_LET_CONST_REF(lhs, self_val.template CastTo<Self>());
+    ADT_CHECK(args.size() == 1);
+    ADT_LET_CONST_REF(rhs, args.at(0).template CastTo<Self>());
+    return GetDimExprClass<ValueT>().New(lhs / rhs);
+  }
+
   static adt::Result<ValueT> Match(axpr::InterpreterBase<ValueT>* interpreter,
                                    const ValueT& self_val,
                                    const std::vector<ValueT>& packed_args_val) {
@@ -93,6 +127,10 @@ axpr::TypeImpl<axpr::BuiltinClassInstance<ValueT>> GetDimExprClass() {
   static auto cls(
       axpr::MakeBuiltinClass<ValueT>("DimExpr", [&](const auto& Define) {
         Define("__str__", &Impl::ToString);
+        Define("__add__", &Impl::Add);
+        Define("__sub__", &Impl::Sub);
+        Define("__mul__", &Impl::Mul);
+        Define("__floordiv__", &Impl::FloorDiv);
         Define("__hash__", &Impl::Hash);
         Define("match", &Impl::Match);
       }));

paddle/ap/include/axpr/type_util.h

@@ -98,6 +98,7 @@ AttrMap<ValueT> GetObjectTypeName2Type() {
                                               OrderedDict<ValueT>,
                                               MutableOrderedDict<ValueT>,
                                               AttrMap<axpr::SerializableValue>,
+                                              AttrMap<ValueT>,
                                               ValueImplTypes...>::Call(&object);
   return object;
 }

paddle/ap/include/drr/builtin_frame_util.h

@@ -27,9 +27,7 @@ void VisitEachBuiltinFrameClass(const DoEachT& DoEach) {
   DoEach(drr::Type<DrrCtx>{}.GetClass());
 }
 
-template <typename VisitorT>
-ap::axpr::AttrMap<axpr::Value> MakeBuiltinFrameAttrMap(
-    const VisitorT& Visitor) {
+inline ap::axpr::AttrMap<axpr::Value> MakeBuiltinFrameAttrMap() {
   ap::axpr::AttrMap<axpr::Value> attr_map;
   ap::axpr::VisitEachBuiltinFrameAttr<axpr::Value>(
       [&](const std::string& k, const axpr::Value& v) { attr_map->Set(k, v); });
@@ -38,7 +36,6 @@ ap::axpr::AttrMap<axpr::Value> MakeBuiltinFrameAttrMap(
     attr_map->Set(std::string("__builtin__") + cls.Name(), cls);
   };
   VisitEachBuiltinFrameClass(Insert);
-  Visitor(Insert);
   return attr_map;
 }
 

paddle/ap/include/drr/drr_interpreter.h

@@ -23,11 +23,8 @@ namespace ap::drr {
 
 class DrrInterpreter {
  public:
-  explicit DrrInterpreter(
-      const axpr::TypeImpl<axpr::BuiltinClassInstance<axpr::Value>>&
-          backend_ir_ctx,
-      const std::weak_ptr<ap::memory::CirclableRefListBase>&
-          circlable_ref_list);
+  explicit DrrInterpreter(const std::weak_ptr<ap::memory::CirclableRefListBase>&
+                              circlable_ref_list);
 
   using Function = ap::axpr::Value;
 

paddle/ap/include/paddle/pir/manual_op.h → paddle/ap/include/paddle/hlir/manual_op.h

@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include "paddle/ap/include/axpr/attr_map.h"
 #include "paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/infer_symbolic_shape.h"
 #include "paddle/phi/core/infermeta_utils.h"
 #include "paddle/pir/include/core/builder.h"
@@ -25,6 +26,22 @@
 
 namespace ap::dialect {
 
+class IR_API FacadeOp
+    : public pir::Op<FacadeOp, ::paddle::dialect::InferSymbolicShapeInterface> {
+ public:
+  using Op::Op;
+  static const char *name() { return "ap_op.facade"; }
+  static constexpr uint32_t attributes_num = 3;
+  static const char *attributes_name[attributes_num];
+  static void Build(pir::Builder &builder,             // NOLINT
+                    pir::OperationArgument &argument,  // NOLINT
+                    const std::vector<pir::Value> &inputs,
+                    const pir::AttributeMap &attributes,
+                    const std::vector<pir::Type> &output_types);
+  void VerifySig() const {}
+  bool InferSymbolicShape(pir::InferSymbolicShapeContext *infer_context);
+};
+
 class IR_API UpSpiderOp
     : public pir::Op<UpSpiderOp,
                      pir::SideEffectTrait,
@@ -134,6 +151,7 @@ class IR_API StoreToGlobalOp
 
 }  // namespace ap::dialect
 
+IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(ap::dialect::FacadeOp);
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(ap::dialect::UpSpiderOp);
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(ap::dialect::DownSpiderOp);
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(ap::dialect::LoadFromRegisterOp);

paddle/ap/include/paddle/meta_tensor_ptr_method_class.h

@@ -103,10 +103,11 @@ struct MetaTensorPtrMethodClass {
 
   adt::Result<axpr::Value> SetDims(const Self& self,
                                    const axpr::Value& dims_val) {
+    if (dims_val.CastableTo<DDim>()) {
+      ADT_LET_CONST_REF(ddim, dims_val.CastTo<DDim>());
+      return SetDimsByDDim(self, ddim);
+    }
     return dims_val.Match(
-        [&](const DDim& ddims) -> adt::Result<axpr::Value> {
-          return SetDimsByDDim(self, ddims);
-        },
         [&](const adt::List<axpr::Value>& list) -> adt::Result<axpr::Value> {
           return SetDimsByIntList(self, list);
         },

paddle/ap/include/paddle/pass/convert_pd_facade_to_ap_facade.h

@@ -0,0 +1,41 @@
+// Copyright (c) 2025 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.
+
+#pragma once
+
+#include <memory>
+#include <optional>
+#include "paddle/pir/include/pass/pass.h"
+
+namespace ap::memory {
+
+class CirclableRefListBase;
+
+}
+
+namespace ap::axpr {
+
+struct Value;
+
+}
+
+namespace cinn {
+namespace dialect {
+namespace ir {
+
+std::unique_ptr<::pir::Pass> CreateConvertPdFacadeToApFacadePass();
+
+}  // namespace ir
+}  // namespace dialect
+}  // namespace cinn

paddle/ap/include/paddle/pass/ir_helper_method_class.h

@@ -18,10 +18,10 @@
 #include "paddle/ap/include/axpr/callable_helper.h"
 #include "paddle/ap/include/axpr/lambda_expr_builder.h"
 #include "paddle/ap/include/drr/drr_value_helper.h"
+#include "paddle/ap/include/paddle/hlir/op_dialect.h"
 #include "paddle/ap/include/paddle/pass/ap_drr_helper.h"
 #include "paddle/ap/include/paddle/pass/ap_generic_drr_pass.h"
 #include "paddle/ap/include/paddle/pass/ir_helper.h"
-#include "paddle/ap/include/paddle/pir/op_dialect.h"
 #include "paddle/ap/include/paddle/pir/packed_ir_op_inner_source_pattern_helper.h"
 #include "paddle/ap/include/paddle/pir/pass_manager_method_class.h"
 #include "paddle/ap/include/paddle/pir/pass_method_class.h"

paddle/ap/include/paddle/phi/ap_infer_meta_helper.h

@@ -15,8 +15,11 @@
 #pragma once
 
 #include "paddle/ap/include/adt/adt.h"
+#include "paddle/ap/include/axpr/attr_map.h"
 #include "paddle/ap/include/axpr/core_expr.h"
+#include "paddle/ap/include/axpr/value.h"
 #include "paddle/phi/core/meta_tensor.h"
+#include "paddle/pir/include/core/operation_utils.h"
 
 namespace phi {
 
@@ -29,6 +32,12 @@ struct ApInferMetaHelper {
   adt::Result<adt::Ok> InferMeta(const std::string& lambda,
                                  const std::vector<const MetaTensor*>* inputs,
                                  std::vector<MetaTensor*>* outputs);
+
+  adt::Result<adt::Ok> InferMetaByAxprHook(
+      const ::paddle::optional<std::vector<const MetaTensor*>>& inputs,
+      const std::string& infer_meta_func_name,
+      const std::string& serialized_attributes,
+      const std::vector<MetaTensor*>& outputs);
 };
 
 }  // namespace phi