chore: import upstream snapshot with attribution

paddle/ap/include/code_gen/arg_source_ctx.h

@@ -0,0 +1,213 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/adt/adt.h"
+#include "paddle/pir/include/dialect/shape/utils/dim_expr.h"
+
+namespace ap::code_gen {
+
+struct NativeIrValueSource {
+  int native_ir_value_index;
+
+  bool operator==(const NativeIrValueSource& other) const {
+    return this->native_ir_value_index == other.native_ir_value_index;
+  }
+};
+
+struct PackedIrValueSource {
+  int packed_ir_value_index;
+  int tensor_member_index;
+
+  bool operator==(const PackedIrValueSource& other) const {
+    return this->packed_ir_value_index == other.packed_ir_value_index &&
+           this->tensor_member_index == other.tensor_member_index;
+  }
+};
+
+using TensorSourceImpl = std::variant<NativeIrValueSource, PackedIrValueSource>;
+
+struct TensorSource : public TensorSourceImpl {
+  using TensorSourceImpl::TensorSourceImpl;
+  ADT_DEFINE_VARIANT_METHODS(TensorSourceImpl);
+};
+
+struct InTensorSource {
+  TensorSource tensor_source;
+
+  bool operator==(const InTensorSource& other) const {
+    return this->tensor_source == other.tensor_source;
+  }
+};
+
+struct OutTensorSource {
+  TensorSource tensor_source;
+
+  bool operator==(const OutTensorSource& other) const {
+    return this->tensor_source == other.tensor_source;
+  }
+};
+
+using InOutTensorSourceImpl = std::variant<InTensorSource, OutTensorSource>;
+
+struct InOutTensorSource : public InOutTensorSourceImpl {
+  using InOutTensorSourceImpl::InOutTensorSourceImpl;
+  ADT_DEFINE_VARIANT_METHODS(InOutTensorSourceImpl);
+};
+
+struct ShapeDimSource {
+  InOutTensorSource tensor_source;
+  int dim_axis;
+
+  bool operator==(const ShapeDimSource& other) const {
+    return this->tensor_source == other.tensor_source &&
+           this->dim_axis == other.dim_axis;
+  }
+};
+
+struct DataDimSource {
+  InOutTensorSource tensor_source;
+  int dim_axis;
+
+  bool operator==(const DataDimSource& other) const {
+    return this->tensor_source == other.tensor_source &&
+           this->dim_axis == other.dim_axis;
+  }
+};
+
+using DimSourceImpl = std::variant<ShapeDimSource, DataDimSource>;
+struct DimSource : public DimSourceImpl {
+  using DimSourceImpl::DimSourceImpl;
+  ADT_DEFINE_VARIANT_METHODS(DimSourceImpl);
+};
+
+template <typename BirNode /* backend ir node*/>
+struct ArgSourceCtxImpl {
+  std::vector<std::pair<BirNode, InTensorSource>> input_and_tensor_source_pairs;
+  std::vector<std::pair<BirNode, OutTensorSource>>
+      output_and_tensor_source_pairs;
+  std::vector<std::pair<symbol::DimExpr, DimSource>>
+      dim_expr_and_dim_source_pairs;
+  std::unordered_map<symbol::DimExpr, DimSource> dim_expr2dim_source;
+
+  std::optional<const InTensorSource*> GetInputTensorSource(
+      const BirNode& node) const {
+    for (const auto& [k, v] : this->input_and_tensor_source_pairs) {
+      if (k == node) {
+        return &v;
+      }
+    }
+    return std::nullopt;
+  }
+
+  std::optional<const OutTensorSource*> GetOutputTensorSource(
+      const BirNode& node) const {
+    for (const auto& [k, v] : this->output_and_tensor_source_pairs) {
+      if (k == node) {
+        return &v;
+      }
+    }
+    return std::nullopt;
+  }
+
+  std::optional<const DimSource*> GetDimExprSource(
+      const symbol::DimExpr& dim_expr) const {
+    const auto& iter = this->dim_expr2dim_source.find(dim_expr);
+    if (iter == this->dim_expr2dim_source.end()) {
+      return std::nullopt;
+    }
+    return &iter->second;
+  }
+
+  bool HasDirectOrIndirectDimExprSource(const symbol::DimExpr& dim_expr) const {
+    if (GetDimExprSource(dim_expr).has_value()) {
+      return true;
+    }
+    return dim_expr.Match([&](const auto& impl) {
+      return HasDirectOrIndirectDimExprSourceImpl(impl);
+    });
+  }
+
+ private:
+  bool HasDirectOrIndirectDimExprSourceImpl(int64_t) const { return true; }
+
+  bool HasDirectOrIndirectDimExprSourceImpl(const std::string& dim_expr) const {
+    return GetDimExprSource(dim_expr).has_value();
+  }
+
+  using Negative = symbol::Negative<symbol::DimExpr>;
+  bool HasDirectOrIndirectDimExprSourceImpl(const Negative& dim_expr) const {
+    return HasDirectOrIndirectUnaryDimExprSource(dim_expr);
+  }
+
+  template <typename T>
+  bool HasDirectOrIndirectUnaryDimExprSource(const T& dim_expr) const {
+    const auto& [operand] = *dim_expr;
+    return HasDirectOrIndirectDimExprSource(operand);
+  }
+
+  using Add = symbol::Add<symbol::DimExpr>;
+  bool HasDirectOrIndirectDimExprSourceImpl(const Add& dim_expr) const {
+    return HasDirectOrIndirectVariadicDimExprSource(dim_expr);
+  }
+
+  using Mul = symbol::Mul<symbol::DimExpr>;
+  bool HasDirectOrIndirectDimExprSourceImpl(const Mul& dim_expr) const {
+    return HasDirectOrIndirectVariadicDimExprSource(dim_expr);
+  }
+
+  using Div = symbol::Div<symbol::DimExpr>;
+  bool HasDirectOrIndirectDimExprSourceImpl(const Div& dim_expr) const {
+    const auto& operand_lhs = (*dim_expr).lhs;
+    const auto& operand_rhs = (*dim_expr).rhs;
+    for (const auto& operand : {operand_lhs, operand_rhs}) {
+      if (!HasDirectOrIndirectDimExprSource(operand)) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  using Max = symbol::Max<symbol::DimExpr>;
+  bool HasDirectOrIndirectDimExprSourceImpl(const Max& dim_expr) const {
+    return HasDirectOrIndirectVariadicDimExprSource(dim_expr);
+  }
+
+  using Min = symbol::Min<symbol::DimExpr>;
+  bool HasDirectOrIndirectDimExprSourceImpl(const Min& dim_expr) const {
+    return HasDirectOrIndirectVariadicDimExprSource(dim_expr);
+  }
+
+  using Broadcast = symbol::Broadcast<symbol::DimExpr>;
+  bool HasDirectOrIndirectDimExprSourceImpl(const Broadcast& dim_expr) const {
+    return HasDirectOrIndirectVariadicDimExprSource(dim_expr);
+  }
+
+  template <typename T>
+  bool HasDirectOrIndirectVariadicDimExprSource(const T& dim_expr) const {
+    const auto& [operands] = dim_expr;
+    for (const auto& operand : *operands) {
+      if (!HasDirectOrIndirectDimExprSource(operand)) {
+        return false;
+      }
+    }
+    return true;
+  }
+};
+
+template <typename BirNode /* backend ir node*/>
+ADT_DEFINE_RC(ArgSourceCtx, ArgSourceCtxImpl<BirNode>);
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/arg_source_helper.h

@@ -0,0 +1,360 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/adt/adt.h"
+#include "paddle/ap/include/axpr/anf_expr_builder.h"
+#include "paddle/ap/include/axpr/core_expr.h"
+#include "paddle/ap/include/axpr/function.h"
+#include "paddle/ap/include/axpr/lambda_expr_builder.h"
+#include "paddle/ap/include/axpr/serializable_value.h"
+#include "paddle/ap/include/code_gen/arg_source_ctx.h"
+#include "paddle/ap/include/code_gen/kernel_arg_id.h"
+
+namespace ap::code_gen {
+
+template <typename BirNode>
+struct ArgSourceHelper {
+  const ArgSourceCtx<BirNode>& arg_source_ctx;
+
+  adt::Result<axpr::Function<axpr::SerializableValue>>
+  MakeRuntimeKerneArgsGetter(
+      const std::list<KernelArgId<BirNode>>& kernel_arg_ids) const {
+    auto GetBody =
+        [&](axpr::LetVar* dispatch_ctx) -> adt::Result<axpr::LetVar*> {
+      std::vector<axpr::LetVar> items;
+      items.reserve(kernel_arg_ids.size());
+      for (const auto& kernel_arg_id : kernel_arg_ids) {
+        ADT_LET_CONST_REF(elt, MakeRuntimeGetter(dispatch_ctx, kernel_arg_id));
+        items.emplace_back(*elt);
+      }
+      auto* ctx = dispatch_ctx->ctx();
+      auto* ret_ptr = &ctx->Var(ctx->NewTmpVarName());
+      *ret_ptr = ctx->Var(axpr::kBuiltinList()).Call(items);
+      return ret_ptr;
+    };
+    ADT_LET_CONST_REF(lambda, CreateLambda("ctx", GetBody));
+    return axpr::Function<axpr::SerializableValue>{lambda, std::nullopt};
+  }
+
+  adt::Result<axpr::Function<axpr::SerializableValue>>
+  MakeRuntimeKerneArgGetter(const KernelArgId<BirNode>& kernel_arg_id) const {
+    auto GetBody =
+        [&](axpr::LetVar* dispatch_ctx) -> adt::Result<axpr::LetVar*> {
+      return MakeRuntimeGetter(dispatch_ctx, kernel_arg_id);
+    };
+    ADT_LET_CONST_REF(lambda, CreateLambda("ctx", GetBody));
+    return axpr::Function<axpr::SerializableValue>{lambda, std::nullopt};
+  }
+
+ private:
+  adt::Result<axpr::LetVar*> MakeRuntimeGetter(
+      axpr::LetVar* dispatch_ctx,
+      const KernelArgId<BirNode>& kernel_arg_id) const {
+    return kernel_arg_id.Match([&](const auto& impl) {
+      return MakeRuntimeGetterImpl(dispatch_ctx, impl);
+    });
+  }
+
+  adt::Result<axpr::LetVar*> MakeRuntimeGetterImpl(
+      axpr::LetVar* dispatch_ctx,
+      const InTensorDataPtrKernelArgId<BirNode>& kernel_arg_id) const {
+    const auto& opt_in_tensor_source =
+        arg_source_ctx->GetInputTensorSource(kernel_arg_id->ir_value);
+    ADT_CHECK(opt_in_tensor_source.has_value());
+    const auto* in_tensor_source = opt_in_tensor_source.value();
+    ADT_LET_CONST_REF(
+        tensor_var_ptr,
+        MakeGetterAnfExprByInTensorSource(dispatch_ctx, *in_tensor_source));
+    return &tensor_var_ptr->Attr("data_ptr");
+  }
+
+  adt::Result<axpr::LetVar*> MakeRuntimeGetterImpl(
+      axpr::LetVar* dispatch_ctx,
+      const OutTensorDataPtrKernelArgId<BirNode>& kernel_arg_id) const {
+    const auto& opt_out_tensor_source =
+        arg_source_ctx->GetOutputTensorSource(kernel_arg_id->ir_value);
+    ADT_CHECK(opt_out_tensor_source.has_value());
+    const auto* out_tensor_source = opt_out_tensor_source.value();
+    ADT_LET_CONST_REF(
+        tensor_var_ptr,
+        MakeGetterAnfExprByOutTensorSource(dispatch_ctx, *out_tensor_source));
+    return &tensor_var_ptr->Attr("data_ptr");
+  }
+
+  template <typename GetVarPtrT>
+  adt::Result<axpr::Lambda<axpr::CoreExpr>> CreateLambda(
+      const std::string& dispatch_ctx_name, const GetVarPtrT& GetVarPtr) const {
+    axpr::LambdaExprBuilder lmbd;
+    auto GetBody =
+        [&](axpr::LetContext& ctx) -> adt::Result<axpr::AnfExpr> {  // NOLINT
+      ADT_LET_CONST_REF(var_ptr, GetVarPtr(&ctx.Var(dispatch_ctx_name)));
+      return static_cast<axpr::AnfExpr>(*var_ptr);
+    };
+    ADT_LET_CONST_REF(anf_expr, lmbd.TryLambda({dispatch_ctx_name}, GetBody));
+    const auto& core_expr = ConvertAnfExprToCoreExpr(anf_expr);
+    ADT_LET_CONST_REF(
+        atomic, core_expr.template TryGet<axpr::Atomic<axpr::CoreExpr>>());
+    ADT_LET_CONST_REF(lambda,
+                      atomic.template TryGet<axpr::Lambda<axpr::CoreExpr>>());
+    return lambda;
+  }
+
+  adt::Result<axpr::LetVar*> MakeGetterAnfExprByInTensorSource(
+      axpr::LetVar* dispatch_ctx,
+      const InTensorSource& in_tensor_source) const {
+    auto& inputs_var = dispatch_ctx->Attr("inputs");
+    return MakeGetterAnfExprByTensorSource(&inputs_var,
+                                           in_tensor_source.tensor_source);
+  }
+
+  adt::Result<axpr::LetVar*> MakeGetterAnfExprByOutTensorSource(
+      axpr::LetVar* dispatch_ctx,
+      const OutTensorSource& out_tensor_source) const {
+    auto& outputs_var = dispatch_ctx->Attr("outputs");
+    return MakeGetterAnfExprByTensorSource(&outputs_var,
+                                           out_tensor_source.tensor_source);
+  }
+
+  adt::Result<axpr::LetVar*> MakeGetterAnfExprByTensorSource(
+      axpr::LetVar* in_out_tensors, const TensorSource& tensor_source) const {
+    using RetT = adt::Result<axpr::LetVar*>;
+    return tensor_source.Match(
+        [&](const NativeIrValueSource& native) -> RetT {
+          return &in_out_tensors->At(native.native_ir_value_index);
+        },
+        [&](const PackedIrValueSource& packed) -> RetT {
+          return &in_out_tensors->At(packed.packed_ir_value_index)
+                      .At(packed.tensor_member_index);
+        });
+  }
+
+  adt::Result<axpr::LetVar*> MakeRuntimeGetterImpl(
+      axpr::LetVar* dispatch_ctx,
+      const DimExprKernelArgId<BirNode>& kernel_arg_id) const {
+    ADT_CHECK(arg_source_ctx->HasDirectOrIndirectDimExprSource(
+        kernel_arg_id->dim_expr));
+    return MakeGetterAnfExprByDimExpr(dispatch_ctx, kernel_arg_id->dim_expr);
+  }
+
+  adt::Result<axpr::LetVar*> MakeGetterAnfExprByDimSource(
+      axpr::LetVar* dispatch_ctx, const DimSource& dim_source) const {
+    using RetT = adt::Result<axpr::LetVar*>;
+    return dim_source.Match(
+        [&](const ShapeDimSource& shape_dim_source) -> RetT {
+          ADT_LET_CONST_REF(tensor_var_ptr,
+                            MakeGetterAnfExprByInOutTensorSource(
+                                dispatch_ctx, shape_dim_source.tensor_source));
+          auto* ctx = dispatch_ctx->ctx();
+          auto* dim_expr =
+              &tensor_var_ptr->Attr("shape").At(shape_dim_source.dim_axis);
+          auto* data_value = &ctx->Var("DataValue").Call(*dim_expr);
+          auto* ret = &ctx->Var(ctx->NewTmpVarName());
+          *ret = data_value->Attr("cast").Call(
+              ctx->Var("DataType").Attr("const_int64"));
+          return ret;
+        },
+        [&](const DataDimSource& data_dim_source) -> RetT {
+          return adt::errors::TypeError{"DataDimSource is not supported yet."};
+        });
+  }
+
+  adt::Result<axpr::LetVar*> MakeGetterAnfExprByInOutTensorSource(
+      axpr::LetVar* dispatch_ctx,
+      const InOutTensorSource& in_out_tensor_source) const {
+    using RetT = adt::Result<axpr::LetVar*>;
+    return in_out_tensor_source.Match(
+        [&](const InTensorSource& in_tensor_source) -> RetT {
+          return MakeGetterAnfExprByInTensorSource(dispatch_ctx,
+                                                   in_tensor_source);
+        },
+        [&](const OutTensorSource& out_tensor_source) -> RetT {
+          return MakeGetterAnfExprByOutTensorSource(dispatch_ctx,
+                                                    out_tensor_source);
+        });
+  }
+
+  adt::Result<axpr::LetVar*> MakeGetterAnfExprByDimExpr(
+      axpr::LetVar* dispatch_ctx, const symbol::DimExpr& dim_expr) const {
+    const auto& opt_dim_source = arg_source_ctx->GetDimExprSource(dim_expr);
+    if (opt_dim_source.has_value()) {
+      return MakeGetterAnfExprByDimSource(dispatch_ctx,
+                                          *opt_dim_source.value());
+    }
+    return dim_expr.Match([&](const auto& impl) {
+      return MakeGetterAnfExprByDimExprImpl(dispatch_ctx, impl);
+    });
+  }
+
+  adt::Result<axpr::LetVar*> MakeGetterAnfExprByDimExprImpl(
+      axpr::LetVar* dispatch_ctx, int64_t c) const {
+    auto* ctx = dispatch_ctx->ctx();
+    auto* ret_var = &ctx->Var(ctx->NewTmpVarName());
+    *ret_var = ctx->Int64(c);
+    return ret_var;
+  }
+
+  adt::Result<axpr::LetVar*> MakeGetterAnfExprByDimExprImpl(
+      axpr::LetVar* dispatch_ctx, const std::string& symbol) const {
+    return adt::errors::NotImplementedError{
+        "Dead code. Symbols have been handled in MakeGetterAnfExprByDimExpr()"};
+  }
+
+  adt::Result<axpr::LetVar*> MakeGetterAnfExprByDimExprImpl(
+      axpr::LetVar* dispatch_ctx,
+      const symbol::Negative<symbol::DimExpr>& dim_expr) const {
+    return adt::errors::NotImplementedError{
+        "Dead code. Negative dim_exprs have been handled in "
+        "MakeGetterAnfExprByDimExprImpl(dispatch_ctx, const "
+        "symbol::Add<symbol::DimExpr>&)"};
+  }
+
+  adt::Result<axpr::LetVar*> MakeGetterAnfExprByDimExprImpl(
+      axpr::LetVar* dispatch_ctx,
+      const symbol::Add<symbol::DimExpr>& dim_expr) const {
+    const auto& [operands] = dim_expr;
+    ADT_CHECK(operands->size() > 0);
+    auto* ctx = dispatch_ctx->ctx();
+    ADT_LET_CONST_REF(
+        init_var_ptr,
+        MakeGetterAnfExprByDimExpr(dispatch_ctx, operands->at(0)));
+    axpr::LetVar* ret_var_ptr = init_var_ptr;
+    for (int i = 1; i < operands->size(); ++i) {
+      const auto& operand = operands->at(i);
+      auto* tmp_var_ptr = &ctx->Var(ctx->NewTmpVarName());
+      if (operand.template Has<symbol::Negative<symbol::DimExpr>>()) {
+        const auto& [operand_operand] =
+            *operand.template Get<symbol::Negative<symbol::DimExpr>>();
+        ADT_LET_CONST_REF(
+            operand_operand_var_ptr,
+            MakeGetterAnfExprByDimExpr(dispatch_ctx, operand_operand));
+        *tmp_var_ptr = ctx->Call(
+            axpr::kBuiltinSub(), *ret_var_ptr, *operand_operand_var_ptr);
+      } else {
+        ADT_LET_CONST_REF(operand_var_ptr,
+                          MakeGetterAnfExprByDimExpr(dispatch_ctx, operand));
+        *tmp_var_ptr =
+            ctx->Call(axpr::kBuiltinAdd(), *ret_var_ptr, *operand_var_ptr);
+      }
+      ret_var_ptr = tmp_var_ptr;
+    }
+    return ret_var_ptr;
+  }
+
+  adt::Result<axpr::LetVar*> MakeGetterAnfExprByDimExprImpl(
+      axpr::LetVar* dispatch_ctx,
+      const symbol::Mul<symbol::DimExpr>& dim_expr) const {
+    const auto& [operands] = dim_expr;
+    ADT_CHECK(operands->size() > 0);
+    auto* ctx = dispatch_ctx->ctx();
+    ADT_LET_CONST_REF(
+        init_var_ptr,
+        MakeGetterAnfExprByDimExpr(dispatch_ctx, operands->at(0)));
+    axpr::LetVar* ret_var_ptr = init_var_ptr;
+    for (int i = 1; i < operands->size(); ++i) {
+      const auto& operand = operands->at(i);
+      auto* tmp_var_ptr = &ctx->Var(ctx->NewTmpVarName());
+      ADT_LET_CONST_REF(operand_var_ptr,
+                        MakeGetterAnfExprByDimExpr(dispatch_ctx, operand));
+      *tmp_var_ptr =
+          ctx->Call(axpr::kBuiltinMul(), *ret_var_ptr, *operand_var_ptr);
+      ret_var_ptr = tmp_var_ptr;
+    }
+    return ret_var_ptr;
+  }
+
+  adt::Result<axpr::LetVar*> MakeGetterAnfExprByDimExprImpl(
+      axpr::LetVar* dispatch_ctx,
+      const symbol::Div<symbol::DimExpr>& dim_expr) const {
+    const auto& operand_lhs = (*dim_expr).lhs;
+    const auto& operand_rhs = (*dim_expr).rhs;
+    auto* ctx = dispatch_ctx->ctx();
+    ADT_LET_CONST_REF(init_var_ptr,
+                      MakeGetterAnfExprByDimExpr(dispatch_ctx, operand_lhs));
+    axpr::LetVar* ret_var_ptr = init_var_ptr;
+    auto* tmp_var_ptr = &ctx->Var(ctx->NewTmpVarName());
+    ADT_LET_CONST_REF(operand_var_ptr,
+                      MakeGetterAnfExprByDimExpr(dispatch_ctx, operand_rhs));
+    *tmp_var_ptr =
+        ctx->Call(axpr::kBuiltinDiv(), *ret_var_ptr, *operand_var_ptr);
+    ret_var_ptr = tmp_var_ptr;
+    return ret_var_ptr;
+  }
+
+  adt::Result<axpr::LetVar*> MakeGetterAnfExprByDimExprImpl(
+      axpr::LetVar* dispatch_ctx,
+      const symbol::Max<symbol::DimExpr>& dim_expr) const {
+    const auto& [operands] = dim_expr;
+    ADT_CHECK(operands->size() > 0);
+    auto* ctx = dispatch_ctx->ctx();
+    ADT_LET_CONST_REF(
+        init_var_ptr,
+        MakeGetterAnfExprByDimExpr(dispatch_ctx, operands->at(0)));
+    axpr::LetVar* ret_var_ptr = init_var_ptr;
+    for (int i = 1; i < operands->size(); ++i) {
+      const auto& operand = operands->at(i);
+      auto* tmp_var_ptr = &ctx->Var(ctx->NewTmpVarName());
+      ADT_LET_CONST_REF(operand_var_ptr,
+                        MakeGetterAnfExprByDimExpr(dispatch_ctx, operand));
+      *tmp_var_ptr = ctx->Call("max", *ret_var_ptr, *operand_var_ptr);
+      ret_var_ptr = tmp_var_ptr;
+    }
+    return ret_var_ptr;
+  }
+
+  adt::Result<axpr::LetVar*> MakeGetterAnfExprByDimExprImpl(
+      axpr::LetVar* dispatch_ctx,
+      const symbol::Min<symbol::DimExpr>& dim_expr) const {
+    const auto& [operands] = dim_expr;
+    ADT_CHECK(operands->size() > 0);
+    auto* ctx = dispatch_ctx->ctx();
+    ADT_LET_CONST_REF(
+        init_var_ptr,
+        MakeGetterAnfExprByDimExpr(dispatch_ctx, operands->at(0)));
+    axpr::LetVar* ret_var_ptr = init_var_ptr;
+    for (int i = 1; i < operands->size(); ++i) {
+      const auto& operand = operands->at(i);
+      auto* tmp_var_ptr = &ctx->Var(ctx->NewTmpVarName());
+      ADT_LET_CONST_REF(operand_var_ptr,
+                        MakeGetterAnfExprByDimExpr(dispatch_ctx, operand));
+      *tmp_var_ptr = ctx->Call("min", *ret_var_ptr, *operand_var_ptr);
+      ret_var_ptr = tmp_var_ptr;
+    }
+    return ret_var_ptr;
+  }
+
+  adt::Result<axpr::LetVar*> MakeGetterAnfExprByDimExprImpl(
+      axpr::LetVar* dispatch_ctx,
+      const symbol::Broadcast<symbol::DimExpr>& dim_expr) const {
+    const auto& [operands] = dim_expr;
+    ADT_CHECK(operands->size() > 0);
+    auto* ctx = dispatch_ctx->ctx();
+    ADT_LET_CONST_REF(
+        init_var_ptr,
+        MakeGetterAnfExprByDimExpr(dispatch_ctx, operands->at(0)));
+    axpr::LetVar* ret_var_ptr = init_var_ptr;
+    for (int i = 1; i < operands->size(); ++i) {
+      const auto& operand = operands->at(i);
+      auto* tmp_var_ptr = &ctx->Var(ctx->NewTmpVarName());
+      ADT_LET_CONST_REF(operand_var_ptr,
+                        MakeGetterAnfExprByDimExpr(dispatch_ctx, operand));
+      *tmp_var_ptr = ctx->Call("broadcast", *ret_var_ptr, *operand_var_ptr);
+      ret_var_ptr = tmp_var_ptr;
+    }
+    return ret_var_ptr;
+  }
+};
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/arg_source_maker.h

@@ -0,0 +1,168 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/adt/adt.h"
+#include "paddle/ap/include/axpr/core_expr.h"
+#include "paddle/ap/include/code_gen/arg_source_ctx.h"
+#include "paddle/ap/include/code_gen/matched_result_pattern_helper.h"
+
+namespace ap::code_gen {
+
+template <typename BirNode>
+struct ArgSourceMaker {
+  const code_gen::MatchedResultPatternHelper<BirNode>& matched_res_ptn_helper;
+
+  using DrrValue = drr::Value;
+  using DrrNode = drr::Node;
+  using DrrPackedIrOp = drr::PackedIrOp<DrrNode>;
+
+  adt::Result<ArgSourceCtx<BirNode>> MakeArgSourceCtx(
+      const DrrPackedIrOp& res_ptn_ir_op) const {
+    ADT_LET_CONST_REF(input_and_tensor_source_pairs,
+                      MakeInputAndTensorSourcePairs(res_ptn_ir_op));
+    ADT_LET_CONST_REF(output_and_tensor_source_pairs,
+                      MakeOutputAndTensorSourcePairs(res_ptn_ir_op));
+    std::vector<std::pair<symbol::DimExpr, DimSource>>
+        dim_expr_and_dim_source_pairs;
+    ADT_RETURN_IF_ERR(CollectDimExprAndDimSourcePairs<InTensorSource>(
+        &dim_expr_and_dim_source_pairs, input_and_tensor_source_pairs));
+    ADT_RETURN_IF_ERR(CollectDimExprAndDimSourcePairs<OutTensorSource>(
+        &dim_expr_and_dim_source_pairs, output_and_tensor_source_pairs));
+    std::unordered_map<symbol::DimExpr, DimSource> dim_expr2dim_source{
+        dim_expr_and_dim_source_pairs.begin(),
+        dim_expr_and_dim_source_pairs.end()};
+    return ArgSourceCtx<BirNode>{input_and_tensor_source_pairs,
+                                 output_and_tensor_source_pairs,
+                                 dim_expr_and_dim_source_pairs,
+                                 dim_expr2dim_source};
+  }
+
+ private:
+  adt::Result<std::vector<std::pair<BirNode, InTensorSource>>>
+  MakeInputAndTensorSourcePairs(const DrrPackedIrOp& res_ptn_ir_op) const {
+    using Ok = adt::Result<adt::Ok>;
+    std::vector<BirNode> inputs;
+    {
+      auto CollectInput = [&](const BirNode& bir_node) -> Ok {
+        inputs.emplace_back(bir_node);
+        return adt::Ok{};
+      };
+      ADT_RETURN_IF_ERR(
+          matched_res_ptn_helper.VisitMatchedBirInputOfRestPtnPackedIrOp(
+              res_ptn_ir_op, CollectInput));
+    }
+    using Pair = std::pair<BirNode, InTensorSource>;
+    std::vector<Pair> ret;
+    ret.reserve(inputs.size());
+    {
+      std::size_t input_idx = 0;
+      auto DoEachIndex = [&](std::size_t index) -> Ok {
+        ADT_CHECK(index < inputs.size());
+        const auto& bir_node = inputs.at(index);
+        NativeIrValueSource native_ir_value_source{input_idx};
+        TensorSource tensor_source{native_ir_value_source};
+        InTensorSource in_tensor_source{tensor_source};
+        Pair pair{bir_node, in_tensor_source};
+        ret.emplace_back(pair);
+        ++input_idx;
+        return adt::Ok{};
+      };
+      auto DoEachSlice = [&](std::size_t start, std::size_t end) -> Ok {
+        for (std::size_t i = start; i < end; ++i) {
+          ADT_CHECK(i < inputs.size());
+          const auto& bir_node = inputs.at(i);
+          PackedIrValueSource packed_ir_value_source{input_idx, i};
+          TensorSource tensor_source{packed_ir_value_source};
+          InTensorSource in_tensor_source{tensor_source};
+          Pair pair{bir_node, in_tensor_source};
+          ret.emplace_back(pair);
+        }
+        ++input_idx;
+        return adt::Ok{};
+      };
+      ADT_RETURN_IF_ERR(matched_res_ptn_helper.VisitApKernelInputIndexOrSlice(
+          res_ptn_ir_op, DoEachIndex, DoEachSlice));
+    }
+    return ret;
+  }
+
+  adt::Result<std::vector<std::pair<BirNode, OutTensorSource>>>
+  MakeOutputAndTensorSourcePairs(const DrrPackedIrOp& res_ptn_ir_op) const {
+    using Ok = adt::Result<adt::Ok>;
+    std::vector<BirNode> outputs;
+    {
+      auto CollectOutput = [&](const BirNode& bir_node) -> Ok {
+        outputs.emplace_back(bir_node);
+        return adt::Ok{};
+      };
+      ADT_RETURN_IF_ERR(
+          matched_res_ptn_helper.VisitMatchedBirOutputOfRestPtnPackedIrOp(
+              res_ptn_ir_op, CollectOutput));
+    }
+    using Pair = std::pair<BirNode, OutTensorSource>;
+    std::vector<Pair> ret;
+    ret.reserve(outputs.size());
+    {
+      std::size_t output_idx = 0;
+      auto DoEachIndex = [&](std::size_t index) -> Ok {
+        ADT_CHECK(index < outputs.size());
+        const auto& bir_node = outputs.at(index);
+        OutTensorSource out_tensor_source{
+            TensorSource{NativeIrValueSource{output_idx}}};
+        ret.emplace_back(Pair{bir_node, out_tensor_source});
+        ++output_idx;
+        return adt::Ok{};
+      };
+      auto DoEachSlice = [&](std::size_t start, std::size_t end) -> Ok {
+        for (std::size_t i = start; i < end; ++i) {
+          ADT_CHECK(i < outputs.size());
+          const auto& bir_node = outputs.at(i);
+          OutTensorSource out_tensor_source{
+              TensorSource{PackedIrValueSource{output_idx, i}}};
+          ret.emplace_back(Pair{bir_node, out_tensor_source});
+        }
+        ++output_idx;
+        return adt::Ok{};
+      };
+      ADT_RETURN_IF_ERR(matched_res_ptn_helper.VisitApKernelOutputIndexOrSlice(
+          res_ptn_ir_op, DoEachIndex, DoEachSlice));
+    }
+    return ret;
+  }
+
+  template <typename TensorSourceT>
+  adt::Result<adt::Ok> CollectDimExprAndDimSourcePairs(
+      std::vector<std::pair<symbol::DimExpr, DimSource>>*
+          dim_expr_and_dim_source_pairs,
+      const std::vector<std::pair<BirNode, TensorSourceT>>& tensor_and_sources)
+      const {
+    for (const auto& [bir_node, tensor_source] : tensor_and_sources) {
+      ADT_LET_CONST_REF(
+          bir_value, matched_res_ptn_helper.CastToBirNativeIrValue(bir_node));
+      ADT_LET_CONST_REF(dim_exprs_ptr, bir_value.GetShapeDimExprsPtr());
+      for (int i = 0; i < dim_exprs_ptr->size(); ++i) {
+        const auto& dim_expr = dim_exprs_ptr->at(i);
+        using Pair = std::pair<symbol::DimExpr, DimSource>;
+        DimSource dim_source{ShapeDimSource{tensor_source, i}};
+        Pair pair{dim_expr, dim_source};
+        dim_expr_and_dim_source_pairs->emplace_back(pair);
+      }
+    }
+    return adt::Ok{};
+  }
+};
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/builtin_frame_util.h

@@ -0,0 +1,50 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/adt/adt.h"
+#include "paddle/ap/include/axpr/builtin_frame_util.h"
+#include "paddle/ap/include/axpr/dim_expr_method_class.h"
+#include "paddle/ap/include/code_gen/code_gen_result_method_class.h"
+#include "paddle/ap/include/code_module/code_module_method_class.h"
+#include "paddle/ap/include/code_module/func_declare_method_class.h"
+#include "paddle/ap/include/code_module/package_method_class.h"
+#include "paddle/ap/include/code_module/project_method_class.h"
+
+namespace ap::code_gen {
+
+template <typename ValueT, typename DoEachT>
+void VisitEachBuiltinFrameClass(const DoEachT& DoEach) {
+  DoEach(code_module::GetProjectClass());
+  DoEach(code_module::GetPackageClass());
+  DoEach(code_module::GetFuncDeclareClass());
+  DoEach(code_module::GetCodeModuleClass());
+  DoEach(axpr::GetDimExprClass<ValueT>());
+  DoEach(GetCodeGenResultClass());
+}
+
+template <typename ValueT>
+axpr::AttrMap<ValueT> MakeBuiltinFrameAttrMap() {
+  axpr::AttrMap<ValueT> attr_map;
+  axpr::VisitEachBuiltinFrameAttr<ValueT>(
+      [&](const std::string& k, const ValueT& v) { attr_map->Set(k, v); });
+  VisitEachBuiltinFrameClass<ValueT>([&](const auto& cls) {
+    attr_map->Set(cls.Name(), cls);
+    attr_map->Set(std::string("__builtin__") + cls.Name(), cls);
+  });
+  return attr_map;
+}
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/code_gen_ctx.h

@@ -0,0 +1,50 @@
+// 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.
+
+#pragma once
+#include "paddle/ap/include/adt/adt.h"
+#include "paddle/ap/include/axpr/attr_map.h"
+#include "paddle/ap/include/axpr/builtin_class_instance.h"
+#include "paddle/ap/include/axpr/core_expr.h"
+#include "paddle/ap/include/axpr/type.h"
+#include "paddle/ap/include/code_gen/arg_source_ctx.h"
+#include "paddle/ap/include/code_module/adt.h"
+#include "paddle/ap/include/code_module/arg_type.h"
+#include "paddle/ap/include/code_module/data_type.h"
+#include "paddle/ap/include/drr/value.h"
+#include "paddle/ap/include/ir_match/ir_match_ctx.h"
+
+namespace ap::code_gen {
+
+template <typename BirNode>
+struct CodeGenCtxImpl {
+  std::optional<ir_match::IrMatchCtx<BirNode>> ir_match_ctx{};
+
+  using DrrNode = drr::Node;
+  using DrrPackedIrOp = drr::PackedIrOp<DrrNode>;
+
+  DrrPackedIrOp res_ptn_ir_op;
+
+  ArgSourceCtx<BirNode> arg_source_ctx;
+
+  bool operator==(const CodeGenCtxImpl& other) const { return this == &other; }
+};
+
+template <typename BirNode>
+ADT_DEFINE_RC(CodeGenCtx, CodeGenCtxImpl<BirNode>);
+
+template <typename ValueT, typename BirNode>
+axpr::TypeImpl<axpr::BuiltinClassInstance<ValueT>> GetCodeGenCtxClass();
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/code_gen_ctx_method_class.h

@@ -0,0 +1,209 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/axpr/method_class.h"
+#include "paddle/ap/include/axpr/naive_class_ops.h"
+#include "paddle/ap/include/axpr/packed_args.h"
+#include "paddle/ap/include/code_gen/arg_source_helper.h"
+#include "paddle/ap/include/code_gen/cuda_code_gen_util.h"
+#include "paddle/ap/include/code_gen/dim_expr_kernel_arg_id_method_class.h"
+#include "paddle/ap/include/code_gen/in_tensor_data_ptr_kernel_arg_id_method_class.h"
+#include "paddle/ap/include/code_gen/ir_op.h"
+#include "paddle/ap/include/code_gen/kernel_arg_id_helper.h"
+#include "paddle/ap/include/code_gen/op_code_gen_ctx.h"
+#include "paddle/ap/include/code_gen/out_tensor_data_ptr_kernel_arg_id_method_class.h"
+#include "paddle/ap/include/code_module/code_module.h"
+#include "paddle/ap/include/ir_match/native_or_ref_ir_value.h"
+#include "paddle/ap/include/registry/registry_singleton.h"
+
+namespace ap::code_gen {
+
+using ap::axpr::BuiltinBinaryFunc;
+using ap::axpr::BuiltinFuncType;
+using ap::axpr::BuiltinUnaryFunc;
+using ap::axpr::CppDataType;
+using ap::axpr::CppPointerType;
+using ap::axpr::DataType;
+using ap::axpr::MethodClass;
+using ap::axpr::PointerType;
+
+template <typename ValueT, typename BirNode>
+struct CodeGenCtxMethodClass {
+  using This = CodeGenCtxMethodClass;
+  using Self = CodeGenCtx<BirNode>;
+
+  static adt::Result<ValueT> StaticMakeAndCheckOutTensorDataPtrKernelArgId(
+      const ValueT& self_val, const std::vector<ValueT>& args) {
+    ADT_LET_CONST_REF(self, self_val.template CastTo<Self>());
+    return This{}.MakeAndCheckOutTensorDataPtrKernelArgId(self, args);
+  }
+
+  adt::Result<ValueT> MakeAndCheckOutTensorDataPtrKernelArgId(
+      const Self& self, const std::vector<ValueT>& args) {
+    ADT_CHECK(args.size() == 1)
+        << adt::errors::TypeError{std::string() +
+                                  "out_tensor_data_ptr_kernel_"
+                                  "arg_id() takes 1 argument but " +
+                                  std::to_string(args.size()) + " were given."};
+    ADT_LET_CONST_REF(ir_value, CastToBirValue(args.at(0)))
+        << adt::errors::TypeError{
+               std::string() +
+               "the argument 1 of "
+               "out_tensor_data_ptr_kernel_arg_id() should be "
+               "'NativeIrValue' or 'RefIrValue' (not '" +
+               axpr::GetTypeName(args.at(0)) + "')."};
+    ADT_RETURN_IF_ERR(CheckOutTensorDataPtrRuntimeAvailable(self, ir_value));
+    OutTensorDataPtrKernelArgId<BirNode> uninitialized{ir_value, std::nullopt};
+    ArgSourceHelper<BirNode> helper{self->arg_source_ctx};
+    ADT_LET_CONST_REF(runtime_getter,
+                      helper.MakeRuntimeKerneArgGetter(uninitialized));
+    OutTensorDataPtrKernelArgId<BirNode> kernel_arg_id{ir_value,
+                                                       runtime_getter};
+    axpr::BuiltinClassInstance<ValueT> instance{
+        GetOutTensorDataPtrKernelArgIdClass<ValueT, BirNode>(), kernel_arg_id};
+    return instance;
+  }
+
+  adt::Result<adt::Ok> CheckOutTensorDataPtrRuntimeAvailable(
+      const Self& self, const BirNode& ir_value) {
+    ADT_CHECK(self->arg_source_ctx->GetOutputTensorSource(ir_value).has_value())
+        << adt::errors::TypeError{
+               std::string() +
+               "out_tensor_data_ptr_kernel_arg_id() failed. "
+               "please check whether the ir_value is an output value of the "
+               "current ap_pattern_fusion_op defined in drr result pattern "
+               "lambda."};
+    return adt::Ok{};
+  }
+
+  static adt::Result<ValueT> StaticMakeAndCheckInTensorDataPtrKernelArgId(
+      const ValueT& self_val, const std::vector<ValueT>& args) {
+    ADT_LET_CONST_REF(self, self_val.template CastTo<Self>());
+    return This{}.MakeAndCheckInTensorDataPtrKernelArgId(self, args);
+  }
+
+  adt::Result<ValueT> MakeAndCheckInTensorDataPtrKernelArgId(
+      const Self& self, const std::vector<ValueT>& args) {
+    ADT_CHECK(args.size() == 1)
+        << adt::errors::TypeError{std::string() +
+                                  "in_tensor_data_ptr_kernel_"
+                                  "arg_id() takes 1 argument but " +
+                                  std::to_string(args.size()) + " were given."};
+    ADT_LET_CONST_REF(ir_value, CastToBirValue(args.at(0)))
+        << adt::errors::TypeError{
+               std::string() +
+               "the argument 1 of "
+               "in_tensor_data_ptr_kernel_arg_id() should be "
+               "'NativeIrValue' or 'RefIrValue' (not '" +
+               axpr::GetTypeName(args.at(0)) + "')."};
+    ADT_RETURN_IF_ERR(CheckInTensorDataPtrRuntimeAvailable(self, ir_value));
+    InTensorDataPtrKernelArgId<BirNode> uninitialized{ir_value, std::nullopt};
+    ArgSourceHelper<BirNode> helper{self->arg_source_ctx};
+    ADT_LET_CONST_REF(runtime_getter,
+                      helper.MakeRuntimeKerneArgGetter(uninitialized));
+    InTensorDataPtrKernelArgId<BirNode> kernel_arg_id{ir_value, runtime_getter};
+    axpr::BuiltinClassInstance<ValueT> instance{
+        GetInTensorDataPtrKernelArgIdClass<ValueT, BirNode>(), kernel_arg_id};
+    return instance;
+  }
+
+  adt::Result<adt::Ok> CheckInTensorDataPtrRuntimeAvailable(
+      const Self& self, const BirNode& ir_value) {
+    ADT_CHECK(self->arg_source_ctx->GetInputTensorSource(ir_value).has_value())
+        << adt::errors::TypeError{
+               std::string() +
+               "in_tensor_data_ptr_kernel_arg_id() failed. "
+               "please check whether the ir_value is an input value of the "
+               "current ap_pattern_fusion_op defined in drr result pattern "
+               "lambda."};
+    return adt::Ok{};
+  }
+
+  adt::Result<BirNode> CastToBirValue(const ValueT& val) {
+    ADT_LET_CONST_REF(
+        instance, val.template CastTo<axpr::BuiltinClassInstance<ValueT>>());
+    if (instance.template Has<typename BirNode::native_value_type>()) {
+      ADT_LET_CONST_REF(
+          ret, instance.template TryGet<typename BirNode::native_value_type>());
+      return ret;
+    }
+    if (instance.template Has<typename BirNode::ref_value_type>()) {
+      ADT_LET_CONST_REF(
+          ret, instance.template TryGet<typename BirNode::ref_value_type>());
+      return ret;
+    }
+    return adt::errors::NotImplementedError{
+        std::string() +
+        "CastToBirValue() failed. only 'NativeIrValue' and 'RefIrValue' "
+        "argument is expected, but '" +
+        axpr::GetTypeName(val) + "' found."};
+  }
+
+  static adt::Result<ValueT> StaticMakeAndCheckDimExprKernelArgId(
+      const ValueT& self_val, const std::vector<ValueT>& args) {
+    ADT_LET_CONST_REF(self, self_val.template CastTo<Self>());
+    return This{}.MakeAndCheckDimExprKernelArgId(self, args);
+  }
+
+  adt::Result<ValueT> MakeAndCheckDimExprKernelArgId(
+      const Self& self, const std::vector<ValueT>& args) {
+    ADT_CHECK(args.size() == 1) << adt::errors::TypeError{
+        std::string() + "dim_expr_kernel_arg_id() takes 1 arguments but " +
+        std::to_string(args.size()) + " were given."};
+    ADT_LET_CONST_REF(dim_expr, args.at(0).template CastTo<symbol::DimExpr>())
+        << adt::errors::TypeError{std::string() +
+                                  "the argument 1 of dim_expr_kernel_arg_id() "
+                                  "should be 'DimExpr' (not '" +
+                                  axpr::GetTypeName(args.at(0)) + "')."};
+    ADT_RETURN_IF_ERR(CheckDimExprRuntimeAvailable(self, dim_expr));
+    DimExprKernelArgId<BirNode> uninitialized{dim_expr, std::nullopt};
+    ArgSourceHelper<BirNode> helper{self->arg_source_ctx};
+    ADT_LET_CONST_REF(runtime_getter,
+                      helper.MakeRuntimeKerneArgGetter(uninitialized));
+    DimExprKernelArgId<BirNode> kernel_arg_id{dim_expr, runtime_getter};
+    axpr::BuiltinClassInstance<ValueT> instance{
+        GetDimExprKernelArgIdClass<ValueT, BirNode>(), kernel_arg_id};
+    return instance;
+  }
+
+  adt::Result<adt::Ok> CheckDimExprRuntimeAvailable(
+      const Self& self, const symbol::DimExpr& dim_expr) {
+    ADT_CHECK(self->arg_source_ctx->HasDirectOrIndirectDimExprSource(dim_expr))
+        << adt::errors::ValueError{
+               std::string() +
+               "DimExpr could not evaluated in runtime. value: " +
+               symbol::ToString(dim_expr)};
+    return adt::Ok{};
+  }
+};
+
+template <typename ValueT, typename BirNode>
+axpr::TypeImpl<axpr::BuiltinClassInstance<ValueT>> GetCodeGenCtxClass() {
+  using ImplMethods = CodeGenCtxMethodClass<ValueT, BirNode>;
+  static auto cls(
+      axpr::MakeBuiltinClass<ValueT>("CodeGenCtx", [&](const auto& Define) {
+        Define("dim_expr_kernel_arg_id",
+               &ImplMethods::StaticMakeAndCheckDimExprKernelArgId);
+        Define("in_tensor_data_ptr_kernel_arg_id",
+               &ImplMethods::StaticMakeAndCheckInTensorDataPtrKernelArgId);
+        Define("out_tensor_data_ptr_kernel_arg_id",
+               &ImplMethods::StaticMakeAndCheckOutTensorDataPtrKernelArgId);
+      }));
+  using Self = typename ImplMethods::Self;
+  return axpr::MakeGlobalNaiveClassOps<Self>(cls);
+}
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/code_gen_result.h

@@ -0,0 +1,40 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/axpr/adt.h"
+#include "paddle/ap/include/axpr/builtin_class_instance.h"
+#include "paddle/ap/include/axpr/builtin_serializable_attr_map.h"
+#include "paddle/ap/include/axpr/core_expr.h"
+#include "paddle/ap/include/axpr/type.h"
+#include "paddle/ap/include/code_module/code_module.h"
+
+namespace ap::code_gen {
+
+template <typename ValueT>
+struct CodeGenResultImpl {
+  code_module::CodeModule code_module;
+  axpr::Function<axpr::SerializableValue> kernel_dispatch_func;
+  axpr::AttrMap<axpr::SerializableValue> kernel_dispatch_const_data;
+
+  bool operator==(const CodeGenResultImpl& other) const {
+    return this == &other;
+  }
+};
+
+template <typename ValueT>
+ADT_DEFINE_RC(CodeGenResult, CodeGenResultImpl<ValueT>);
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/code_gen_result_method_class.h

@@ -0,0 +1,25 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/axpr/method_class.h"
+#include "paddle/ap/include/axpr/naive_class_ops.h"
+#include "paddle/ap/include/code_gen/code_gen_result.h"
+
+namespace ap::code_gen {
+
+axpr::TypeImpl<axpr::BuiltinClassInstance<axpr::Value>> GetCodeGenResultClass();
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/cuda_code_gen_util.h

@@ -0,0 +1,32 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/code_gen/op_cuda_gen_impl.h"
+#include "paddle/ap/include/code_gen/value.h"
+
+namespace ap::code_gen {
+
+template <typename BirNode>
+adt::Result<axpr::ClassAttrs<axpr::SerializableValue>>
+ConvertFusionOpToClassAttrs(const OpCodeGenCtx<BirNode>& op_code_gen_ctx,
+                            const IrOp<BirNode>& ir_op) {
+  OpCudaCodeGenImpl<BirNode> impl{};
+  ADT_LET_CONST_REF(class_attrs,
+                    impl.ConvertFusionOpToClassAttrs(op_code_gen_ctx, ir_op));
+  return class_attrs;
+}
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/dim_expr_kernel_arg_id.h

@@ -0,0 +1,53 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/adt/adt.h"
+#include "paddle/ap/include/axpr/builtin_class_instance.h"
+#include "paddle/ap/include/axpr/function.h"
+#include "paddle/ap/include/axpr/serializable_value.h"
+#include "paddle/ap/include/axpr/type.h"
+#include "paddle/pir/include/dialect/shape/utils/dim_expr.h"
+
+namespace ap::code_gen {
+
+template <typename BirNode>
+struct DimExprKernelArgIdImpl {
+  symbol::DimExpr dim_expr;
+  std::optional<axpr::Function<axpr::SerializableValue>> runtime_getter;
+
+  bool operator==(const DimExprKernelArgIdImpl& other) const {
+    return this->dim_expr == other.dim_expr;
+  }
+
+  template <typename ValueT>
+  adt::Result<ValueT> CastData() const {
+    axpr::BuiltinClassInstance<ValueT> instance{axpr::GetDimExprClass<ValueT>(),
+                                                this->dim_expr};
+    return ValueT{instance};
+  }
+
+  std::size_t GetHashValue() const {
+    return std::hash<symbol::DimExpr>()(this->dim_expr);
+  }
+};
+
+template <typename BirNode>
+ADT_DEFINE_RC(DimExprKernelArgId, DimExprKernelArgIdImpl<BirNode>);
+
+template <typename ValueT, typename BirNode>
+axpr::TypeImpl<axpr::BuiltinClassInstance<ValueT>> GetDimExprKernelArgIdClass();
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/dim_expr_kernel_arg_id_method_class.h

@@ -0,0 +1,88 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/axpr/method_class.h"
+#include "paddle/ap/include/axpr/naive_class_ops.h"
+#include "paddle/ap/include/code_gen/code_gen_ctx.h"
+#include "paddle/ap/include/code_gen/dim_expr_kernel_arg_id.h"
+#include "paddle/ap/include/code_gen/kernel_arg_id_helper.h"
+
+namespace ap::code_gen {
+
+template <typename ValueT, typename BirNode /* background ir node */>
+struct DimExprKernelArgIdMethodClass {
+  using This = DimExprKernelArgIdMethodClass;
+  using Self = DimExprKernelArgId<BirNode>;
+
+  static adt::Result<ValueT> ToString(const ValueT& self_val,
+                                      const std::vector<ValueT>& args) {
+    ADT_LET_CONST_REF(self, self_val.template CastTo<Self>());
+    std::ostringstream ss;
+    ss << self->dim_expr;
+    return ss.str();
+  }
+
+  static adt::Result<ValueT> Hash(const ValueT& self_val,
+                                  const std::vector<ValueT>& args) {
+    ADT_LET_CONST_REF(self, self_val.template CastTo<Self>());
+    std::size_t hash_value = std::hash<symbol::DimExpr>()(self->dim_expr);
+    return static_cast<int64_t>(hash_value);
+  }
+
+  static adt::Result<ValueT> GetAttr(const ValueT& self_val,
+                                     const std::vector<ValueT>& args) {
+    ADT_LET_CONST_REF(self, self_val.template CastTo<Self>());
+    ADT_CHECK(args.size() == 1);
+    const auto& attr_name_val = args.at(0);
+    ADT_LET_CONST_REF(attr_name, attr_name_val.template TryGet<std::string>());
+    if (attr_name == "value") {
+      return self->template CastData<ValueT>();
+    }
+    if (attr_name == "type") {
+      return This{}.GetArgType(self);
+    }
+    if (attr_name == "runtime_getter") {
+      ADT_CHECK(self->runtime_getter.has_value())
+          << adt::errors::ValueError{"no runtime getter initialized"};
+      return self->runtime_getter.value();
+    }
+    return adt::errors::AttributeError{
+        std::string() + "'DimExprKernelArgId' instance has no attribute '" +
+        attr_name + "'."};
+  }
+
+  adt::Result<ValueT> GetArgType(const Self& self) {
+    KernelArgIdHelper<BirNode> helper;
+    ADT_LET_CONST_REF(arg_type, helper.GetArgType(self));
+    return arg_type.template CastTo<ValueT>();
+  }
+};
+
+template <typename ValueT, typename BirNode>
+axpr::TypeImpl<axpr::BuiltinClassInstance<ValueT>>
+GetDimExprKernelArgIdClass() {
+  using ImplMethods = DimExprKernelArgIdMethodClass<ValueT, BirNode>;
+  static auto cls(axpr::MakeBuiltinClass<ValueT>(
+      "DimExprKernelArgId", [&](const auto& Define) {
+        Define("__str__", &ImplMethods::ToString);
+        Define("__hash__", &ImplMethods::Hash);
+        Define("__getattr__", &ImplMethods::GetAttr);
+      }));
+  using Self = typename ImplMethods::Self;
+  return axpr::MakeGlobalNaiveClassOps<Self>(cls);
+}
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/in_tensor_data_ptr_kernel_arg_id.h

@@ -0,0 +1,67 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/adt/adt.h"
+#include "paddle/ap/include/axpr/builtin_class_instance.h"
+#include "paddle/ap/include/axpr/function.h"
+#include "paddle/ap/include/axpr/serializable_value.h"
+#include "paddle/ap/include/axpr/type.h"
+
+namespace ap::code_gen {
+
+template <typename BirNode>
+struct InTensorDataPtrKernelArgIdImpl {
+  BirNode ir_value;
+  std::optional<axpr::Function<axpr::SerializableValue>> runtime_getter;
+
+  bool operator==(const InTensorDataPtrKernelArgIdImpl& other) const {
+    return this->ir_value == other.ir_value;
+  }
+
+  template <typename ValueT>
+  adt::Result<ValueT> CastData() const {
+    using RetT = adt::Result<ValueT>;
+    return this->ir_value.Match(
+        [&](const typename BirNode::native_value_type& impl) -> RetT {
+          return impl;
+        },
+        [&](const typename BirNode::ref_value_type& impl) -> RetT {
+          return impl;
+        },
+        [&](const auto& impl) -> RetT {
+          using T = std::decay_t<decltype(impl)>;
+          return adt::errors::NotImplementedError{
+              std::string() +
+              "CastData() failed, only NativeIrValue and RefIrValue supported, "
+              "but '" +
+              typeid(T).name() + "' found."};
+        });
+  }
+
+  std::size_t GetHashValue() const {
+    return std::hash<BirNode>()(this->ir_value);
+  }
+};
+
+template <typename BirNode>
+ADT_DEFINE_RC(InTensorDataPtrKernelArgId,
+              InTensorDataPtrKernelArgIdImpl<BirNode>);
+
+template <typename ValueT, typename BirNode>
+axpr::TypeImpl<axpr::BuiltinClassInstance<ValueT>>
+GetInTensorDataPtrKernelArgIdClass();
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/in_tensor_data_ptr_kernel_arg_id_method_class.h

@@ -0,0 +1,86 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/axpr/method_class.h"
+#include "paddle/ap/include/axpr/naive_class_ops.h"
+#include "paddle/ap/include/code_gen/in_tensor_data_ptr_kernel_arg_id.h"
+#include "paddle/ap/include/code_gen/kernel_arg_id_helper.h"
+
+namespace ap::code_gen {
+
+template <typename ValueT, typename BirNode /* background ir node */>
+struct InTensorDataPtrKernelArgIdMethodClass {
+  using This = InTensorDataPtrKernelArgIdMethodClass;
+  using Self = InTensorDataPtrKernelArgId<BirNode>;
+
+  static adt::Result<ValueT> ToString(const ValueT& self_val,
+                                      const std::vector<ValueT>& args) {
+    ADT_LET_CONST_REF(self, self_val.template CastTo<Self>());
+    const void* ptr = self.__adt_rc_shared_ptr_raw_ptr();
+    std::ostringstream ss;
+    ss << "<InTensorDataPtrKernelArgId object at " << ptr << ">";
+    return ss.str();
+  }
+
+  static adt::Result<ValueT> Hash(const ValueT& self_val,
+                                  const std::vector<ValueT>& args) {
+    ADT_LET_CONST_REF(self, self_val.template CastTo<Self>());
+    std::size_t hash_value = self->ir_value.GetHashValue();
+    return static_cast<int64_t>(hash_value);
+  }
+
+  static adt::Result<ValueT> GetAttr(const ValueT& self_val,
+                                     const std::vector<ValueT>& args) {
+    ADT_LET_CONST_REF(self, self_val.template CastTo<Self>());
+    ADT_CHECK(args.size() == 1);
+    const auto& attr_name_val = args.at(0);
+    ADT_LET_CONST_REF(attr_name, attr_name_val.template CastTo<std::string>());
+    if (attr_name == "type") {
+      return This{}.GetArgType(self);
+    }
+    if (attr_name == "runtime_getter") {
+      ADT_CHECK(self->runtime_getter.has_value())
+          << adt::errors::ValueError{"no runtime getter initialized"};
+      return self->runtime_getter.value();
+    }
+    return adt::errors::AttributeError{
+        std::string() +
+        "'InTensorDataPtrKernelArgId' instance has no attribute '" + attr_name +
+        "'."};
+  }
+
+  adt::Result<ValueT> GetArgType(const Self& self) {
+    KernelArgIdHelper<BirNode> helper;
+    ADT_LET_CONST_REF(arg_type, helper.GetArgType(self));
+    return arg_type.template CastTo<ValueT>();
+  }
+};
+
+template <typename ValueT, typename BirNode>
+axpr::TypeImpl<axpr::BuiltinClassInstance<ValueT>>
+GetInTensorDataPtrKernelArgIdClass() {
+  using ImplMethods = InTensorDataPtrKernelArgIdMethodClass<ValueT, BirNode>;
+  static auto cls(axpr::MakeBuiltinClass<ValueT>(
+      "InTensorDataPtrKernelArgId", [&](const auto& Define) {
+        Define("__str__", &ImplMethods::ToString);
+        Define("__hash__", &ImplMethods::Hash);
+        Define("__getattr__", &ImplMethods::GetAttr);
+      }));
+  using Self = typename ImplMethods::Self;
+  return axpr::MakeGlobalNaiveClassOps<Self>(cls);
+}
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/ir_op.h

@@ -0,0 +1,55 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/adt/adt.h"
+#include "paddle/ap/include/axpr/builtin_class_instance.h"
+
+namespace ap::code_gen {
+
+template <typename BirNode>
+using IrOpImpl = std::variant<typename BirNode::native_op_type,
+                              typename BirNode::packed_op_type,
+                              typename BirNode::ref_op_type>;
+
+template <typename BirNode>
+struct IrOp : public IrOpImpl<BirNode> {
+  using IrOpImpl<BirNode>::IrOpImpl;
+  ADT_DEFINE_VARIANT_METHODS(IrOpImpl<BirNode>);
+
+  template <typename ValueT>
+  static adt::Result<IrOp> CastFrom(const ValueT& val) {
+    ADT_LET_CONST_REF(
+        instance, val.template CastTo<axpr::BuiltinClassInstance<ValueT>>());
+    if (instance.template Has<typename BirNode::native_op_type>()) {
+      ADT_LET_CONST_REF(
+          ret, instance.template TryGet<typename BirNode::native_op_type>());
+      return ret;
+    }
+    if (instance.template Has<typename BirNode::packed_op_type>()) {
+      ADT_LET_CONST_REF(
+          ret, instance.template TryGet<typename BirNode::packed_op_type>());
+      return ret;
+    }
+    if (instance.template Has<typename BirNode::ref_op_type>()) {
+      ADT_LET_CONST_REF(
+          ret, instance.template TryGet<typename BirNode::ref_op_type>());
+      return ret;
+    }
+    return adt::errors::ValueError{"IrOp::CastFrom failed."};
+  }
+};
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/kernel_arg.h

@@ -0,0 +1,36 @@
+// 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.
+
+#pragma once
+#include "paddle/ap/include/axpr/attr_map.h"
+#include "paddle/ap/include/axpr/core_expr.h"
+#include "paddle/ap/include/axpr/type.h"
+#include "paddle/ap/include/code_gen/kernel_arg_id.h"
+#include "paddle/ap/include/code_module/adt.h"
+
+namespace ap::code_gen {
+
+struct KernelArgImpl {
+  KernelArgId kernel_arg_id;
+  axpr::Lambda<axpr::CoreExpr> getter_lambda;
+
+  bool operator==(const KernelArgImpl& other) const {
+    return this->kernel_arg_id == other.kernel_arg_id &&
+           this->getter_lambda == other.getter_lambda;
+  }
+};
+
+ADT_DEFINE_RC(KernelArg, KernelArgImpl);
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/kernel_arg_id.h

@@ -0,0 +1,75 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/adt/adt.h"
+#include "paddle/ap/include/code_gen/dim_expr_kernel_arg_id.h"
+#include "paddle/ap/include/code_gen/in_tensor_data_ptr_kernel_arg_id.h"
+#include "paddle/ap/include/code_gen/out_tensor_data_ptr_kernel_arg_id.h"
+
+namespace ap::code_gen {
+
+template <typename BirNode>
+using KernelArgIdImpl = std::variant<DimExprKernelArgId<BirNode>,
+                                     InTensorDataPtrKernelArgId<BirNode>,
+                                     OutTensorDataPtrKernelArgId<BirNode>>;
+
+template <typename BirNode>
+struct KernelArgId : public KernelArgIdImpl<BirNode> {
+  using KernelArgIdImpl<BirNode>::KernelArgIdImpl;
+
+  ADT_DEFINE_VARIANT_METHODS(KernelArgIdImpl<BirNode>);
+
+  template <typename ValueT>
+  ValueT CastTo() const {
+    return Match([](const auto& impl) -> ValueT { return impl; });
+  }
+
+  template <typename ValueT>
+  static adt::Result<KernelArgId> CastFrom(const ValueT& val) {
+    using RetT = adt::Result<KernelArgId>;
+    return val.Match(
+        [](const DimExprKernelArgId<BirNode>& impl) -> RetT { return impl; },
+        [](const InTensorDataPtrKernelArgId<BirNode>& impl) -> RetT {
+          return impl;
+        },
+        [](const OutTensorDataPtrKernelArgId<BirNode>& impl) -> RetT {
+          return impl;
+        },
+        [](const auto& impl) -> RetT {
+          return adt::errors::TypeError{"KernelArgId::CastFrom() failed."};
+        });
+  }
+
+  std::size_t GetHashValue() const {
+    std::size_t hash_value = Match(
+        [&](const auto& impl) -> std::size_t { return impl->GetHashValue(); });
+    return adt::hash_combine(this->index(), hash_value);
+  }
+};
+
+}  // namespace ap::code_gen
+
+namespace std {
+
+template <typename BirNode>
+struct hash<ap::code_gen::KernelArgId<BirNode>> {
+  std::size_t operator()(
+      const ap::code_gen::KernelArgId<BirNode>& arg_id) const {
+    return arg_id.GetHashValue();
+  }
+};
+
+}  // namespace std

paddle/ap/include/code_gen/kernel_arg_id_helper.h

@@ -0,0 +1,87 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/axpr/pointer_type_util.h"
+#include "paddle/ap/include/code_gen/kernel_arg_id.h"
+#include "paddle/ap/include/code_module/arg_type.h"
+
+namespace ap::code_gen {
+
+template <typename BirNode>
+struct KernelArgIdHelper {
+  using BirNativeIrValue = typename BirNode::native_value_type;
+
+  template <typename ValueT>
+  adt::Result<KernelArgId<BirNode>> CastToKernelArgId(const ValueT& val) {
+    using RetT = adt::Result<KernelArgId<BirNode>>;
+    return val.Match(
+        [&](const DimExprKernelArgId<BirNode>& impl) -> RetT { return impl; },
+        [&](const InTensorDataPtrKernelArgId<BirNode>& impl) -> RetT {
+          return impl;
+        },
+        [&](const OutTensorDataPtrKernelArgId<BirNode>& impl) -> RetT {
+          return impl;
+        },
+        [&](const auto& impl) -> RetT {
+          return adt::errors::TypeError{
+              std::string() +
+              "only DimExprKernelArgId, InTensorDataPtrKernelArgId and "
+              "OutTensorDataPtrKernelArgId (not including '" +
+              axpr::GetTypeName(val) + "') can be cast to KernelArgId"};
+        });
+  }
+
+  adt::Result<code_module::ArgType> GetArgType(
+      const KernelArgId<BirNode>& arg_id) {
+    using RetT = adt::Result<code_module::ArgType>;
+    return arg_id.Match(
+        [](const DimExprKernelArgId<BirNode>&) -> RetT {
+          return axpr::CppDataType<int64_t>();
+        },
+        [&](const InTensorDataPtrKernelArgId<BirNode>& in_data_ptr) -> RetT {
+          ADT_LET_CONST_REF(ir_value,
+                            GetBirNativeIrValue(in_data_ptr->ir_value));
+          ADT_LET_CONST_REF(data_type, ir_value.GetDataType());
+          return axpr::GetConstPointerTypeFromDataType(data_type);
+        },
+        [&](const OutTensorDataPtrKernelArgId<BirNode>& out_data_ptr) -> RetT {
+          ADT_LET_CONST_REF(ir_value,
+                            GetBirNativeIrValue(out_data_ptr->ir_value));
+          ADT_LET_CONST_REF(data_type, ir_value.GetDataType());
+          return axpr::GetMutablePointerTypeFromDataType(data_type);
+        });
+  }
+
+  adt::Result<BirNativeIrValue> GetBirNativeIrValue(
+      const BirNode& bir_node) const {
+    using RetT = adt::Result<BirNativeIrValue>;
+    return bir_node.Match(
+        [&](const BirNativeIrValue& impl) -> RetT { return impl; },
+        [&](const typename BirNode::ref_value_type& impl) -> RetT {
+          return impl.GetOwnerNativeIrValue();
+        },
+        [&](const auto& impl) -> RetT {
+          using T = std::decay_t<decltype(impl)>;
+          return adt::errors::NotImplementedError{
+              std::string() +
+              "GetBirNativeIrValue() failed. only 'NativeIrValue' and "
+              "'RefIrValue' argument expected, but '" +
+              typeid(T).name() + "' found."};
+        });
+  }
+};
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/loop_anchor_flags.h

@@ -0,0 +1,25 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/adt/adt.h"
+
+namespace ap::code_gen {
+
+ADT_DEFINE_TAG(tLoopAnchorFlag);
+
+using LoopAnchorFlags = adt::List<tLoopAnchorFlag<bool>>;
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/matched_result_pattern_helper.h

@@ -0,0 +1,236 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/axpr/anf_expr_util.h"
+#include "paddle/ap/include/axpr/atomic.h"
+#include "paddle/ap/include/axpr/data_type_util.h"
+#include "paddle/ap/include/axpr/lambda_expr_builder.h"
+#include "paddle/ap/include/code_gen/arg_source_ctx.h"
+#include "paddle/ap/include/drr/drr_graph_descriptor.h"
+#include "paddle/ap/include/drr/drr_node_descriptor.h"
+#include "paddle/ap/include/drr/res_ptn_packed_ir_op_declare_data.h"
+#include "paddle/ap/include/drr/result_pattern_helper.h"
+#include "paddle/ap/include/drr/value.h"
+#include "paddle/ap/include/graph/graph_helper.h"
+#include "paddle/ap/include/ir_match/graph_match_ctx.h"
+#include "paddle/ap/include/ir_match/graph_matcher.h"
+#include "paddle/ap/include/ir_match/ir_match_ctx.h"
+
+namespace ap::code_gen {
+
+template <typename BirNode /* backend ir node*/>
+struct MatchedResultPatternHelper {
+  using DrrNode = drr::Node;
+
+  using DrrCtx = drr::DrrCtx;
+
+  using DrrNativeIrValue = drr::NativeIrValue<DrrNode>;
+  using DrrPackedIrValue = drr::PackedIrValue<DrrNode>;
+  using DrrIrValue = drr::IrValue;
+
+  using DrrNativeIrOp = drr::NativeIrOp<DrrNode>;
+  using DrrNativeIrOpOperand = drr::NativeIrOpOperand<DrrNode>;
+  using DrrNativeIrOpResult = drr::NativeIrOpResult<DrrNode>;
+  using DrrPackedIrOp = drr::PackedIrOp<DrrNode>;
+  using DrrPackedIrOpOperand = drr::PackedIrOpOperand<DrrNode>;
+  using DrrPackedIrOpResult = drr::PackedIrOpResult<DrrNode>;
+  using DrrOptPackedIrOp = drr::OptPackedIrOp<DrrNode>;
+  using DrrOptPackedIrOpOperand = drr::OptPackedIrOpOperand<DrrNode>;
+  using DrrOptPackedIrOpResult = drr::OptPackedIrOpResult<DrrNode>;
+
+  using DrrIrOpImpl = std::variant<DrrNativeIrOp, DrrPackedIrOp>;
+
+  using IrMatchCtx = ir_match::IrMatchCtx<BirNode>;
+
+  using GraphMatchCtx = ir_match::GraphMatchCtx<BirNode>;
+
+  const GraphMatchCtx& match_ctx_;
+  const DrrCtx& drr_ctx_;
+
+  template <typename DoEachT>
+  adt::Result<adt::Ok> VisitMatchedBirInputOfRestPtnPackedIrOp(
+      const DrrPackedIrOp& res_ptn_ir_op, const DoEachT& DoEach) const {
+    auto CollectInput =
+        [&](const DrrIrValue& drr_ir_value) -> adt::Result<adt::Ok> {
+      return VisitMatchedBirValueOfResPtnIrValue(drr_ir_value, DoEach);
+    };
+    ADT_RETURN_IF_ERR(
+        VisitResPtnInputIrValueByResPtnIrOp(res_ptn_ir_op, CollectInput));
+    return adt::Ok{};
+  }
+
+  template <typename DoEachT>
+  adt::Result<adt::Ok> VisitMatchedBirOutputOfRestPtnPackedIrOp(
+      const DrrPackedIrOp& res_ptn_ir_op, const DoEachT& DoEach) const {
+    auto DoEachDrrIrValue =
+        [&](const DrrIrValue& drr_ir_value) -> adt::Result<adt::Ok> {
+      return VisitMatchedBirValueOfResPtnIrValue(drr_ir_value, DoEach);
+    };
+    ADT_RETURN_IF_ERR(
+        VisitResPtnOutputIrValueByResPtnIrOp(res_ptn_ir_op, DoEachDrrIrValue));
+    return adt::Ok{};
+  }
+
+  template <typename DoEachT>
+  adt::Result<adt::Ok> VisitResPtnInputIrValueByResPtnIrOp(
+      const DrrPackedIrOp& res_ptn_ir_op, const DoEachT& DoEach) const {
+    drr::ResultPatternHelper helper{drr_ctx_};
+    return helper.VisitResPtnInputIrValueByResPtnIrOp(res_ptn_ir_op, DoEach);
+  }
+
+  template <typename DoEachT>
+  adt::Result<adt::Ok> VisitMatchedBirValueOfResPtnIrValue(
+      const DrrIrValue& res_ptn_ir_value, const DoEachT& DoEach) const {
+    using Ok = adt::Result<adt::Ok>;
+    const auto& opt_ir_value = SrcPtnIrValue4ResPtnIrValue(res_ptn_ir_value);
+    ADT_CHECK(opt_ir_value.has_value());
+    const auto& ir_value = opt_ir_value.value();
+    ADT_RETURN_IF_ERR(
+        match_ctx_->VisitBigGraphIrValueNode(ir_value.node(), DoEach));
+    return adt::Ok{};
+  }
+
+  template <typename DoEachIndexT, typename DoEachSliceT>
+  adt::Result<adt::Ok> VisitApKernelInputIndexOrSlice(
+      const DrrPackedIrOp& res_ptn_ir_op,
+      const DoEachIndexT& DoEachIndex,
+      const DoEachSliceT& DoEachSlice) const {
+    std::size_t start = 0;
+    using Ok = adt::Result<adt::Ok>;
+    auto DoEachIrValue = [&](const DrrIrValue& drr_ir_value) -> Ok {
+      ADT_LET_CONST_REF(num_ir_values, GetResPtnNumBirValues(drr_ir_value));
+      ADT_RETURN_IF_ERR(drr_ir_value.Match(
+          [&](const DrrNativeIrValue&) -> Ok {
+            ADT_RETURN_IF_ERR(DoEachIndex(start));
+            return adt::Ok{};
+          },
+          [&](const DrrPackedIrValue&) -> Ok {
+            ADT_RETURN_IF_ERR(DoEachSlice(start, start + num_ir_values));
+            return adt::Ok{};
+          }));
+      start += num_ir_values;
+      return adt::Ok{};
+    };
+    ADT_RETURN_IF_ERR(
+        VisitResPtnInputIrValueByResPtnIrOp(res_ptn_ir_op, DoEachIrValue));
+    return adt::Ok{};
+  }
+
+  template <typename DoEachIndexT, typename DoEachSliceT>
+  adt::Result<adt::Ok> VisitApKernelOutputIndexOrSlice(
+      const DrrPackedIrOp& res_ptn_ir_op,
+      const DoEachIndexT& DoEachIndex,
+      const DoEachSliceT& DoEachSlice) const {
+    std::size_t start = 0;
+    using Ok = adt::Result<adt::Ok>;
+    auto DoEachIrValue = [&](const DrrIrValue& drr_ir_value) -> Ok {
+      ADT_LET_CONST_REF(num_ir_values, GetResPtnNumBirValues(drr_ir_value));
+      ADT_RETURN_IF_ERR(drr_ir_value.Match(
+          [&](const DrrNativeIrValue&) -> Ok {
+            ADT_RETURN_IF_ERR(DoEachIndex(start));
+            return adt::Ok{};
+          },
+          [&](const DrrPackedIrValue&) -> Ok {
+            ADT_RETURN_IF_ERR(DoEachSlice(start, start + num_ir_values));
+            return adt::Ok{};
+          }));
+      start += num_ir_values;
+      return adt::Ok{};
+    };
+    ADT_RETURN_IF_ERR(
+        VisitResPtnOutputIrValueByResPtnIrOp(res_ptn_ir_op, DoEachIrValue));
+    return adt::Ok{};
+  }
+
+  adt::Result<std::size_t> GetApKernelNumOutputs(
+      const DrrPackedIrOp& res_ptn_ir_op) const {
+    std::size_t num_outputs = 0;
+    auto AccNumOutputs =
+        [&](const DrrIrValue& drr_ir_value) -> adt::Result<adt::Ok> {
+      ADT_LET_CONST_REF(num_ir_values, GetResPtnNumBirValues(drr_ir_value));
+      num_outputs += num_ir_values;
+      return adt::Ok{};
+    };
+    ADT_RETURN_IF_ERR(
+        VisitResPtnOutputIrValueByResPtnIrOp(res_ptn_ir_op, AccNumOutputs));
+    return num_outputs;
+  }
+
+  template <typename T, typename IrOpT, typename DoEachT>
+  adt::Result<adt::Ok> VisitEachMatchedDrrIrValueAndOutputSlice(
+      const std::vector<T>& output_values,
+      const IrOpT& res_ptn_ir_op,
+      const DoEachT& DoEach) const {
+    std::size_t offset = 0;
+    auto DoEachSlice =
+        [&](const DrrIrValue& drr_ir_value) -> adt::Result<adt::Ok> {
+      ADT_LET_CONST_REF(num_ir_values, GetResPtnNumBirValues(drr_ir_value));
+      ADT_CHECK(offset + num_ir_values <= output_values.size());
+      std::vector<T> slice{output_values.begin() + offset,
+                           output_values.begin() + offset + num_ir_values};
+      ADT_RETURN_IF_ERR(DoEach(drr_ir_value, slice));
+      offset += num_ir_values;
+      return adt::Ok{};
+    };
+    return VisitResPtnOutputIrValueByResPtnIrOp(res_ptn_ir_op, DoEachSlice);
+  }
+
+  template <typename IrOpT, typename DoEachT>
+  adt::Result<adt::Ok> VisitResPtnOutputIrValueByResPtnIrOp(
+      const IrOpT& res_ptn_ir_op, const DoEachT& DoEach) const {
+    drr::ResultPatternHelper helper{drr_ctx_};
+    return helper.VisitResPtnOutputIrValueByResPtnIrOp(res_ptn_ir_op, DoEach);
+  }
+
+  adt::Result<std::size_t> GetResPtnNumBirValues(
+      const DrrIrValue& res_ptn_ir_value) const {
+    const auto& opt_src_ptn_ir_value =
+        SrcPtnIrValue4ResPtnIrValue(res_ptn_ir_value);
+    if (!opt_src_ptn_ir_value.has_value()) {
+      // internal ir value in result pattern.
+      return 1U;
+    }
+    return match_ctx_->GetNumBigGraphIrValueNodes(
+        opt_src_ptn_ir_value.value().node());
+  }
+
+  std::optional<DrrIrValue> SrcPtnIrValue4ResPtnIrValue(
+      const DrrIrValue& res_ptn_ir_value) const {
+    drr::ResultPatternHelper helper{drr_ctx_};
+    return helper.SrcPtnIrValue4ResPtnIrValue(res_ptn_ir_value);
+  }
+
+  using BirNativeIrValue = typename BirNode::native_value_type;
+
+  adt::Result<BirNativeIrValue> CastToBirNativeIrValue(
+      const BirNode& bir_node) const {
+    using RetT = adt::Result<BirNativeIrValue>;
+    return bir_node.Match(
+        [&](const typename BirNode::native_value_type& bir_value) -> RetT {
+          return bir_value;
+        },
+        [&](const typename BirNode::ref_value_type& ref_value) -> RetT {
+          return ref_value.GetOwnerNativeIrValue();
+        },
+        [&](const auto&) -> RetT {
+          return adt::errors::TypeError{
+              "bir_node is not an PirNode::native_value_type or "
+              "BirNode::ref_value_type"};
+        });
+  }
+};
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/op_code_gen_ctx.h

@@ -0,0 +1,43 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/adt/adt.h"
+#include "paddle/ap/include/axpr/type.h"
+#include "paddle/ap/include/code_gen/kernel_arg_id.h"
+#include "paddle/ap/include/code_gen/loop_anchor_flags.h"
+#include "paddle/ap/include/ir_match/native_or_ref_ir_value.h"
+
+namespace ap::code_gen {
+
+template <typename BirNode>
+struct CodeGenCtxImpl;
+
+template <typename BirNode>
+struct OpCodeGenCtxImpl {
+  std::weak_ptr<CodeGenCtxImpl<BirNode>> code_gen_ctx;
+
+  LoopAnchorFlags input_index_loop_anchor_flags;
+  LoopAnchorFlags output_index_loop_anchor_flags;
+
+  bool operator==(const OpCodeGenCtxImpl& other) const {
+    return this == &other;
+  }
+};
+
+template <typename BirNode>
+ADT_DEFINE_RC(OpCodeGenCtx, OpCodeGenCtxImpl<BirNode>);
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/op_cuda_gen_impl.h

@@ -0,0 +1,34 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/adt/adt.h"
+#include "paddle/ap/include/axpr/class_attrs.h"
+#include "paddle/ap/include/axpr/serializable_value.h"
+#include "paddle/ap/include/code_gen/ir_op.h"
+#include "paddle/ap/include/code_gen/op_code_gen_ctx.h"
+
+namespace ap::code_gen {
+
+template <typename BirNode>
+struct OpCudaCodeGenImpl {
+  adt::Result<std::string> CodeGen(const OpCodeGenCtx<BirNode>& op_code_gen_ctx,
+                                   const IrOp<BirNode>& ir_op);
+  adt::Result<axpr::ClassAttrs<axpr::SerializableValue>>
+  ConvertFusionOpToClassAttrs(const OpCodeGenCtx<BirNode>& op_code_gen_ctx,
+                              const IrOp<BirNode>& ir_op);
+};
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/out_tensor_data_ptr_kernel_arg_id.h

@@ -0,0 +1,67 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/adt/adt.h"
+#include "paddle/ap/include/axpr/builtin_class_instance.h"
+#include "paddle/ap/include/axpr/function.h"
+#include "paddle/ap/include/axpr/serializable_value.h"
+#include "paddle/ap/include/axpr/type.h"
+
+namespace ap::code_gen {
+
+template <typename BirNode>
+struct OutTensorDataPtrKernelArgIdImpl {
+  BirNode ir_value;
+  std::optional<axpr::Function<axpr::SerializableValue>> runtime_getter{};
+
+  bool operator==(const OutTensorDataPtrKernelArgIdImpl& other) const {
+    return this->ir_value == other.ir_value;
+  }
+
+  template <typename ValueT>
+  adt::Result<ValueT> CastData() const {
+    using RetT = adt::Result<ValueT>;
+    return this->ir_value.Match(
+        [&](const typename BirNode::native_value_type& impl) -> RetT {
+          return impl;
+        },
+        [&](const typename BirNode::ref_value_type& impl) -> RetT {
+          return impl;
+        },
+        [&](const auto& impl) -> RetT {
+          using T = std::decay_t<decltype(impl)>;
+          return adt::errors::NotImplementedError{
+              std::string() +
+              "CastData() failed, only NativeIrValue and RefIrValue supported, "
+              "but '" +
+              typeid(T).name() + "' found."};
+        });
+  }
+
+  std::size_t GetHashValue() const {
+    return std::hash<BirNode>()(this->ir_value);
+  }
+};
+
+template <typename BirNode>
+ADT_DEFINE_RC(OutTensorDataPtrKernelArgId,
+              OutTensorDataPtrKernelArgIdImpl<BirNode>);
+
+template <typename ValueT, typename BirNode>
+axpr::TypeImpl<axpr::BuiltinClassInstance<ValueT>>
+GetOutTensorDataPtrKernelArgIdClass();
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/out_tensor_data_ptr_kernel_arg_id_method_class.h

@@ -0,0 +1,86 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/axpr/method_class.h"
+#include "paddle/ap/include/axpr/naive_class_ops.h"
+#include "paddle/ap/include/code_gen/in_tensor_data_ptr_kernel_arg_id.h"
+#include "paddle/ap/include/code_gen/kernel_arg_id_helper.h"
+
+namespace ap::code_gen {
+
+template <typename ValueT, typename BirNode /* background ir node */>
+struct OutTensorDataPtrKernelArgIdMethodClass {
+  using This = OutTensorDataPtrKernelArgIdMethodClass;
+  using Self = OutTensorDataPtrKernelArgId<BirNode>;
+
+  static adt::Result<ValueT> ToString(const ValueT& self_val,
+                                      const std::vector<ValueT>& args) {
+    ADT_LET_CONST_REF(self, self_val.template CastTo<Self>());
+    const void* ptr = self.__adt_rc_shared_ptr_raw_ptr();
+    std::ostringstream ss;
+    ss << "<OutTensorDataPtrKernelArgId object at " << ptr << ">";
+    return ss.str();
+  }
+
+  static adt::Result<ValueT> Hash(const ValueT& self_val,
+                                  const std::vector<ValueT>& args) {
+    ADT_LET_CONST_REF(self, self_val.template CastTo<Self>());
+    std::size_t hash_value = self->ir_value.GetHashValue();
+    return static_cast<int64_t>(hash_value);
+  }
+
+  static adt::Result<ValueT> GetAttr(const ValueT& self_val,
+                                     const std::vector<ValueT>& args) {
+    ADT_LET_CONST_REF(self, self_val.template CastTo<Self>());
+    ADT_CHECK(args.size() == 1);
+    const auto& attr_name_val = args.at(0);
+    ADT_LET_CONST_REF(attr_name, attr_name_val.template TryGet<std::string>());
+    if (attr_name == "type") {
+      return This{}.GetArgType(self);
+    }
+    if (attr_name == "runtime_getter") {
+      ADT_CHECK(self->runtime_getter.has_value())
+          << adt::errors::ValueError{"no runtime getter initialized"};
+      return self->runtime_getter.value();
+    }
+    return adt::errors::AttributeError{
+        std::string() +
+        "'OutTensorDataPtrKernelArgId' instance has no attribute '" +
+        attr_name + "'."};
+  }
+
+  adt::Result<ValueT> GetArgType(const Self& self) {
+    KernelArgIdHelper<BirNode> helper;
+    ADT_LET_CONST_REF(arg_type, helper.GetArgType(self));
+    return arg_type.template CastTo<ValueT>();
+  }
+};
+
+template <typename ValueT, typename BirNode>
+axpr::TypeImpl<axpr::BuiltinClassInstance<ValueT>>
+GetOutTensorDataPtrKernelArgIdClass() {
+  using ImplMethods = OutTensorDataPtrKernelArgIdMethodClass<ValueT, BirNode>;
+  static auto cls(axpr::MakeBuiltinClass<ValueT>(
+      "OutTensorDataPtrKernelArgId", [&](const auto& Define) {
+        Define("__str__", &ImplMethods::ToString);
+        Define("__hash__", &ImplMethods::Hash);
+        Define("__getattr__", &ImplMethods::GetAttr);
+      }));
+  using Self = typename ImplMethods::Self;
+  return axpr::MakeGlobalNaiveClassOps<Self>(cls);
+}
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/value.h

@@ -0,0 +1,35 @@
+// 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.
+
+#pragma once
+#include "paddle/ap/include/adt/adt.h"
+#include "paddle/ap/include/axpr/builtin_serializable_attr_map.h"
+#include "paddle/ap/include/axpr/dim_expr.h"
+#include "paddle/ap/include/axpr/value.h"
+#include "paddle/ap/include/code_gen/code_gen_ctx.h"
+#include "paddle/ap/include/code_gen/code_gen_result.h"
+#include "paddle/ap/include/code_gen/dim_expr_kernel_arg_id.h"
+#include "paddle/ap/include/code_gen/in_tensor_data_ptr_kernel_arg_id.h"
+#include "paddle/ap/include/code_gen/out_tensor_data_ptr_kernel_arg_id.h"
+#include "paddle/ap/include/code_module/adt.h"
+#include "paddle/ap/include/code_module/code_module.h"
+#include "paddle/ap/include/code_module/data_type.h"
+#include "paddle/ap/include/ir_match/op_match_ctx.h"
+#include "paddle/ap/include/ir_match/tensor_match_ctx.h"
+
+namespace ap::code_gen {
+
+using axpr::Value;
+
+}  // namespace ap::code_gen

paddle/ap/include/code_gen/value_method_class.h

@@ -0,0 +1,24 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ap/include/axpr/dim_expr_method_class.h"
+#include "paddle/ap/include/code_gen/code_gen_ctx_method_class.h"
+#include "paddle/ap/include/code_gen/code_gen_result_method_class.h"
+#include "paddle/ap/include/code_gen/dim_expr_kernel_arg_id_method_class.h"
+#include "paddle/ap/include/code_gen/in_tensor_data_ptr_kernel_arg_id_method_class.h"
+#include "paddle/ap/include/code_gen/out_tensor_data_ptr_kernel_arg_id_method_class.h"
+#include "paddle/ap/include/ir_match/op_match_ctx_method_class.h"
+#include "paddle/ap/include/ir_match/tensor_match_ctx_method_class.h"