chore: import upstream snapshot with attribution

paddle/ap/include/kernel_dispatch/ap_variadic_kernel.h

@@ -0,0 +1,41 @@
+// 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 <string>
+#include <vector>
+
+#include "paddle/ap/include/adt/adt.h"
+#include "paddle/ap/include/kernel_dispatch/device_ctx.h"
+
+namespace phi {
+
+class DenseTensor;
+
+}
+
+namespace ap::kernel_dispatch {
+
+adt::Result<adt::Ok> ApVariadicKernel(
+    const DeviceCtx& device_ctx,
+    const std::vector<const phi::DenseTensor*>& xs,
+    int num_outputs,
+    const std::string& kernel_define_lambda,
+    const std::string& infer_meta_lambda,
+    const std::string& kernel_dispatch_lambda,
+    const std::string& kernel_dispatch_const_data_lambda,
+    std::vector<phi::DenseTensor*> outs);
+
+}  // namespace ap::kernel_dispatch

paddle/ap/include/kernel_dispatch/arg_value.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/adt/adt.h"
+#include "paddle/ap/include/axpr/value.h"
+#include "paddle/ap/include/kernel_dispatch/typed_buffer.h"
+#include "paddle/ap/include/rt_module/arg_value.h"
+#include "paddle/phi/core/dense_tensor.h"
+
+namespace phi {
+
+class DenseTensor;
+
+}
+
+namespace ap::kernel_dispatch {
+
+using code_module::ArgType;
+
+using rt_module::ArgValue;
+
+using rt_module::CastToArgValue;
+
+}  // namespace ap::kernel_dispatch

paddle/ap/include/kernel_dispatch/builtin_frame_util.h

@@ -0,0 +1,42 @@
+// 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/kernel_dispatch/const_tensor_method_class.h"
+#include "paddle/ap/include/kernel_dispatch/dispatch_ctx_method_class.h"
+#include "paddle/ap/include/kernel_dispatch/mutable_tensor_method_class.h"
+
+namespace ap::kernel_dispatch {
+
+template <typename ValueT, typename DoEachT>
+void VisitEachBuiltinFrameAttr(const DoEachT& DoEach) {
+  // Do Nothing.
+}
+
+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); });
+  VisitEachBuiltinFrameAttr<ValueT>([&](const std::string& k, const ValueT& v) {
+    attr_map->Set(k, v);
+    attr_map->Set(std::string("__builtin__") + k, v);
+  });
+  return attr_map;
+}
+
+}  // namespace ap::kernel_dispatch

paddle/ap/include/kernel_dispatch/const_tensor.h

@@ -0,0 +1,73 @@
+// 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_module/data_type.h"
+#include "paddle/ap/include/kernel_dispatch/typed_buffer.h"
+#include "paddle/phi/core/dense_tensor.h"
+
+namespace phi {
+
+class DenseTensor;
+
+}
+
+namespace ap::kernel_dispatch {
+
+using ConstTensorDataImpl = std::variant<const phi::DenseTensor*, TypedBuffer>;
+struct ConstTensorData : public ConstTensorDataImpl {
+  using ConstTensorDataImpl::ConstTensorDataImpl;
+  ADT_DEFINE_VARIANT_METHODS(ConstTensorDataImpl);
+
+  template <typename T>
+  const T* data() const {
+    return Match(
+        [](const phi::DenseTensor* tensor) -> const T* {
+          return reinterpret_cast<const T*>(tensor->data());
+        },
+        [](const TypedBuffer& buffer) -> const T* {
+          return reinterpret_cast<T*>(buffer->buffer);
+        });
+  }
+  const void* data() const {
+    return Match(
+        [](const phi::DenseTensor* tensor) -> const void* {
+          return tensor->data();
+        },
+        [](const TypedBuffer& buffer) -> const void* {
+          return buffer->buffer;
+        });
+  }
+  phi::DataType dtype() const {
+    return Match([](const phi::DenseTensor* tensor) { return tensor->dtype(); },
+                 [](const TypedBuffer& buffer) { return buffer->dtype; });
+  }
+};
+
+template <typename ValueT>
+struct ConstTensorImpl {
+  ConstTensorData tensor_data;
+  adt::List<ValueT> dims;
+
+  bool operator==(const ConstTensorImpl& other) const {
+    return other.tensor_data == this->tensor_data && other.dims == this->dims;
+  }
+};
+
+template <typename ValueT>
+ADT_DEFINE_RC(ConstTensor, ConstTensorImpl<ValueT>);
+
+}  // namespace ap::kernel_dispatch

paddle/ap/include/kernel_dispatch/const_tensor_method_class.h

@@ -0,0 +1,110 @@
+// 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/data_type_util.h"
+#include "paddle/ap/include/axpr/method_class.h"
+#include "paddle/ap/include/axpr/naive_class_ops.h"
+#include "paddle/ap/include/kernel_dispatch/const_tensor.h"
+
+namespace ap::kernel_dispatch {
+
+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::Method;
+using ap::axpr::MethodClass;
+using ap::axpr::PointerType;
+using ap::axpr::PointerValue;
+
+namespace detail {
+
+template <typename Val>
+Result<Val> ConstTensorShapeGetAttr(const ConstTensor<Val>& tensor,
+                                    const std::string&) {
+  return tensor->dims;
+}
+
+template <typename T>
+const T* GetConstTensorDataPtr(const ap::axpr::CppDataType<T>&,
+                               const ConstTensorData& tensor) {
+  return tensor.template data<T>();
+}
+
+template <typename Val>
+Result<Val> ConstTensorDataGetAttr(const ConstTensor<Val>& tensor,
+                                   const std::string&) {
+  phi::DataType dtype = tensor->tensor_data.dtype();
+  const auto& data_type = ap::axpr::GetDataTypeFromPhiDataType(dtype);
+  ADT_RETURN_IF_ERR(data_type);
+  return data_type.GetOkValue().Match(
+      [&](const adt::Undefined&) -> Result<Val> {
+        return TypeError{"dtype is invalid."};
+      },
+      [&](const auto& impl) -> Result<Val> {
+        return PointerValue{GetConstTensorDataPtr(impl, tensor->tensor_data)};
+      });
+}
+
+template <typename Val>
+using ConstTensorGetAttrT = Result<Val> (*)(const ConstTensor<Val>& tensor,
+                                            const std::string&);
+
+template <typename Val>
+Result<Val> TensorGetAttr(const ConstTensor<Val>& tensor,
+                          const std::string& name) {
+  static const std::unordered_map<std::string, ConstTensorGetAttrT<Val>> map{
+      {"shape", &ConstTensorShapeGetAttr<Val>},
+      {"data_ptr", &ConstTensorDataGetAttr<Val>},
+  };
+  const auto& iter = map.find(name);
+  if (iter == map.end()) {
+    return AttributeError{std::string("'Tensor' has no attribute '") + name +
+                          "'"};
+  }
+  return iter->second(tensor, name);
+}
+
+}  // namespace detail
+
+template <typename ValueT>
+struct ConstTensorMethodClass {
+  using Self = ConstTensor<ValueT>;
+
+  static adt::Result<ValueT> GetAttr(const ValueT& obj_val,
+                                     const std::vector<ValueT>& args) {
+    ADT_CHECK(args.size() == 1);
+    const auto& attr_name_val = args.at(0);
+    ADT_LET_CONST_REF(obj, axpr::Get<ConstTensor<ValueT>>(obj_val));
+    ADT_LET_CONST_REF(attr_name, attr_name_val.template TryGet<std::string>());
+    return detail::TensorGetAttr<ValueT>(obj, attr_name);
+  }
+};
+
+template <typename ValueT>
+axpr::TypeImpl<axpr::BuiltinClassInstance<ValueT>> GetConstTensorClass() {
+  using ImplMethods = ConstTensorMethodClass<ValueT>;
+  static auto cls(
+      axpr::MakeBuiltinClass<ValueT>("ConstTensor", [&](const auto& DoEach) {
+        DoEach("__getattr__", &ImplMethods::GetAttr);
+      }));
+  using Self = typename ImplMethods::Self;
+  return axpr::MakeGlobalNaiveClassOps<Self>(cls);
+}
+
+}  // namespace ap::kernel_dispatch

paddle/ap/include/kernel_dispatch/device_ctx.h

@@ -0,0 +1,36 @@
+// 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 "paddle/ap/include/adt/adt.h"
+#include "paddle/ap/include/axpr/pointer_value.h"
+
+namespace ap::kernel_dispatch {
+
+class DeviceCtxImpl {
+ public:
+  virtual ~DeviceCtxImpl() {}
+
+  virtual adt::Result<axpr::PointerValue> GetStreamAddrAsVoidPtr() = 0;
+
+  bool operator==(const DeviceCtxImpl& other) const { return this == &other; }
+
+ protected:
+  DeviceCtxImpl() {}
+};
+
+ADT_DEFINE_RC(DeviceCtx, DeviceCtxImpl);
+
+}  // namespace ap::kernel_dispatch

paddle/ap/include/kernel_dispatch/device_ctx_method_class.h

@@ -0,0 +1,25 @@
+// 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 "paddle/ap/include/axpr/naive_class_ops.h"
+#include "paddle/ap/include/axpr/value.h"
+#include "paddle/ap/include/kernel_dispatch/device_ctx.h"
+
+namespace ap::kernel_dispatch {
+
+axpr::TypeImpl<axpr::BuiltinClassInstance<axpr::Value>> GetDeviceCtxClass();
+
+}

paddle/ap/include/kernel_dispatch/dispatch_ctx.h

@@ -0,0 +1,46 @@
+// 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/type.h"
+#include "paddle/ap/include/axpr/value.h"
+#include "paddle/ap/include/code_module/arg_type.h"
+#include "paddle/ap/include/code_module/data_type.h"
+#include "paddle/ap/include/kernel_dispatch/dispatch_raw_ctx.h"
+#include "paddle/ap/include/kernel_dispatch/typed_buffer.h"
+#include "paddle/phi/core/dense_tensor.h"
+
+namespace phi {
+
+class DenseTensor;
+
+}
+
+namespace ap::kernel_dispatch {
+
+template <typename ValueT>
+struct DispatchCtxImpl {
+  DispatchRawCtx<ValueT> raw_ctx;
+
+  axpr::AttrMap<axpr::SerializableValue> kernel_dispatch_const_data;
+
+  bool operator==(const DispatchCtxImpl& other) const { return &other == this; }
+};
+
+template <typename ValueT>
+ADT_DEFINE_RC(DispatchCtx, DispatchCtxImpl<ValueT>);
+
+}  // namespace ap::kernel_dispatch

paddle/ap/include/kernel_dispatch/dispatch_ctx_method_class.h

@@ -0,0 +1,246 @@
+// 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/data_type_util.h"
+#include "paddle/ap/include/axpr/method_class.h"
+#include "paddle/ap/include/axpr/naive_class_ops.h"
+#include "paddle/ap/include/axpr/value.h"
+#include "paddle/ap/include/axpr/value_method_class.h"
+#include "paddle/ap/include/kernel_dispatch/device_ctx_method_class.h"
+#include "paddle/ap/include/kernel_dispatch/dispatch_ctx.h"
+#include "paddle/ap/include/rt_module/function_method_class.h"
+
+namespace ap::kernel_dispatch {
+
+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::DataValue;
+using ap::axpr::MethodClass;
+using ap::axpr::PointerType;
+using ap::axpr::PointerValue;
+
+namespace detail {
+
+template <typename Val>
+Result<Val> DispatchCtxGetInputs(const DispatchCtx<Val>& ctx,
+                                 const std::string& attr_name) {
+  return ctx->raw_ctx->inputs;
+}
+
+template <typename Val>
+Result<Val> DispatchCtxGetOutputs(const DispatchCtx<Val>& ctx,
+                                  const std::string& attr_name) {
+  return ctx->raw_ctx->outputs;
+}
+
+template <typename Val>
+Result<Val> DispatchCtxGetDeviceCtx(const DispatchCtx<Val>& ctx,
+                                    const std::string& attr_name) {
+  return GetDeviceCtxClass().New(ctx->raw_ctx->device_ctx);
+}
+
+template <typename Val>
+Result<adt::List<ArgValue>> GetKernelArgs(const Val& args) {
+  const Result<adt::List<Val>>& arg_list =
+      args.template TryGet<adt::List<Val>>();
+  ADT_RETURN_IF_ERR(arg_list);
+  adt::List<ArgValue> ret;
+  ret->reserve(arg_list.GetOkValue()->size());
+  for (const auto& arg : *arg_list.GetOkValue()) {
+    const Result<ArgValue>& arg_value = CastToArgValue(arg);
+    ADT_RETURN_IF_ERR(arg_value);
+    ret->emplace_back(arg_value.GetOkValue());
+  }
+  return ret;
+}
+
+template <typename Val>
+Result<Val> LaunchCuda(const Val& self_val, const std::vector<Val>& args) {
+  ADT_CHECK(args.size() == 4) << TypeError{
+      std::string() +
+      "DispatchCtx.launch_cuda take 6 arguments (including self) but " +
+      std::to_string(args.size()) + " were given."};
+  ADT_LET_CONST_REF(ctx, axpr::Get<DispatchCtx<Val>>(self_val));
+  ADT_LET_CONST_REF(func_name, args.at(0).template TryGet<std::string>());
+  ADT_LET_CONST_REF(num_blocks, args.at(1).template TryGet<int64_t>());
+  ADT_LET_CONST_REF(num_threads, args.at(2).template TryGet<int64_t>());
+  ADT_LET_CONST_REF(kernel_args, GetKernelArgs(args.at(3)));
+  ADT_RETURN_IF_ERR(ctx->raw_ctx->LaunchCudaKernel(
+      func_name, num_blocks, num_threads, kernel_args));
+  return adt::Nothing{};
+}
+
+template <typename Val, BuiltinFuncType<Val> BuiltinFunc>
+Result<Val> MakeDispatchCtxMethod(const DispatchCtx<Val>& ctx,
+                                  const std::string&) {
+  return ap::axpr::Method<Val>{ctx, BuiltinFuncType<Val>{BuiltinFunc}};
+}
+
+template <typename Val, typename T>
+Result<Val> DispatchCtxType(const DispatchCtx<Val>& ctx, const std::string&) {
+  return ap::axpr::TypeImpl<T>{};
+}
+
+template <typename Val>
+using KernelCtxGettAttrT = Result<Val> (*)(const DispatchCtx<Val>& ctx,
+                                           const std::string&);
+
+template <typename Val>
+Result<Val> DispatchCtxGetAttr(const DispatchCtx<Val>& ctx,
+                               const std::string& name) {
+  static const std::unordered_map<std::string, KernelCtxGettAttrT<Val>> map{
+      {ap::axpr::TypeImpl<ap::axpr::DataValue>{}.Name(),
+       &DispatchCtxType<Val, ap::axpr::DataValue>},
+      {"inputs", &DispatchCtxGetInputs<Val>},
+      {"outputs", &DispatchCtxGetOutputs<Val>},
+      {"device_ctx", &DispatchCtxGetDeviceCtx},
+  };
+  const auto& iter = map.find(name);
+  if (iter == map.end()) {
+    return AttributeError{std::string("'DispatchCtx' has no attribute '") +
+                          name + "'"};
+  }
+  return iter->second(ctx, name);
+}
+
+}  // namespace detail
+
+template <typename ValueT>
+struct DispatchCtxMethodClass {
+  using This = DispatchCtxMethodClass;
+  using Self = DispatchCtx<ValueT>;
+
+  static adt::Result<axpr::Value> ToString(
+      const axpr::Value& self_val, const std::vector<axpr::Value>& 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 << "<DispatchCtx object at " << ptr << ">";
+    return ss.str();
+  }
+
+  static adt::Result<ValueT> GetAttr(const ValueT& self_val,
+                                     const std::vector<ValueT>& args) {
+    ADT_LET_CONST_REF(self, axpr::Get<Self>(self_val));
+    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 == "kernel_dispatch_const_data") {
+      return self->kernel_dispatch_const_data;
+    }
+    return detail::DispatchCtxGetAttr<ValueT>(self, attr_name);
+  }
+
+  static adt::Result<ValueT> StaticGetInputIndexByName(
+      const ValueT& self_val, const std::vector<ValueT>& args) {
+    ADT_LET_CONST_REF(self, axpr::Get<Self>(self_val));
+    ADT_CHECK(args.size() == 1) << adt::errors::TypeError{
+        std::string() +
+        "'DispatchCtx.get_input_index_by_name' takes 1 argument but " +
+        std::to_string(args.size()) + " were given."};
+    ADT_LET_CONST_REF(tensor_name, args.at(0).template TryGet<std::string>())
+        << adt::errors::TypeError{
+               std::string() +
+               "the argument 1 of 'DispatchCtx.get_input_index_by_name' should "
+               "be str (not '" +
+               axpr::GetTypeName(args.at(0)) + "')."};
+    return This{}.GetInputIndexByName(self, tensor_name);
+  }
+
+  adt::Result<ValueT> GetInputIndexByName(const Self& self,
+                                          const std::string& tensor_name) {
+    const auto& data = self->kernel_dispatch_const_data;
+    ADT_LET_CONST_REF(
+        name2idx,
+        data->template TryGet<axpr::AttrMap<axpr::SerializableValue>>(
+            "__builtin_ap_kernel_input_name_to_index"));
+    ADT_LET_CONST_REF(index, name2idx->template TryGet<int64_t>(tensor_name));
+    return index;
+  }
+
+  static adt::Result<ValueT> StaticGetOutputIndexByName(
+      const ValueT& self_val, const std::vector<ValueT>& args) {
+    ADT_LET_CONST_REF(self, axpr::Get<Self>(self_val));
+    ADT_CHECK(args.size() == 1) << adt::errors::TypeError{
+        std::string() +
+        "'DispatchCtx.get_output_index_by_name' takes 1 argument but " +
+        std::to_string(args.size()) + " were given."};
+    ADT_LET_CONST_REF(tensor_name, args.at(0).template TryGet<std::string>())
+        << adt::errors::TypeError{
+               std::string() +
+               "the argument 1 of 'DispatchCtx.get_output_index_by_name' "
+               "should be str (not '" +
+               axpr::GetTypeName(args.at(0)) + "')."};
+    return This{}.GetOutputIndexByName(self, tensor_name);
+  }
+
+  static adt::Result<ValueT> StaticGetSoFunction(
+      const ValueT& self_val, const std::vector<ValueT>& args) {
+    ADT_LET_CONST_REF(self, axpr::Get<Self>(self_val));
+    ADT_CHECK(args.size() == 1) << adt::errors::TypeError{
+        std::string() +
+        "'DispatchCtx.get_so_function()' takes 1 argument but " +
+        std::to_string(args.size()) + " were given."};
+    ADT_LET_CONST_REF(function_name, args.at(0).template TryGet<std::string>())
+        << adt::errors::TypeError{
+               std::string() +
+               "the argument 1 of 'DispatchCtx.get_so_function()' "
+               "should be str (not '" +
+               axpr::GetTypeName(args.at(0)) + "')."};
+    ADT_LET_CONST_REF(
+        rt_module,
+        self->raw_ctx->rt_module
+            .template TryGet<std::shared_ptr<const ap::rt_module::Module>>());
+    ADT_LET_CONST_REF(function, rt_module->Get(function_name))
+        << adt::errors::TypeError{
+               std::string() +
+               "DispatchCtx.get_so_function() failed. so function '" +
+               function_name + "' not found"};
+    return rt_module::GetSoFunctionClass().New(function);
+  }
+
+  adt::Result<ValueT> GetOutputIndexByName(const Self& self,
+                                           const std::string& tensor_name) {
+    const auto& data = self->kernel_dispatch_const_data;
+    ADT_LET_CONST_REF(
+        name2idx,
+        data->template TryGet<axpr::AttrMap<axpr::SerializableValue>>(
+            "__builtin_ap_kernel_output_name_to_index"));
+    ADT_LET_CONST_REF(index, name2idx->template TryGet<int64_t>(tensor_name));
+    return index;
+  }
+};
+
+template <typename ValueT>
+axpr::TypeImpl<axpr::BuiltinClassInstance<ValueT>> GetDispatchCtxClass() {
+  using Methods = DispatchCtxMethodClass<ValueT>;
+  static auto cls(
+      axpr::MakeBuiltinClass<ValueT>("DispatchCtx", [&](const auto& Yield) {
+        Yield("__str__", &Methods::ToString);
+        Yield("__getattr__", &Methods::GetAttr);
+        Yield("get_input_index_by_name", &Methods::StaticGetInputIndexByName);
+        Yield("get_output_index_by_name", &Methods::StaticGetOutputIndexByName);
+        Yield("get_so_function", &Methods::StaticGetSoFunction);
+      }));
+  using Self = typename Methods::Self;
+  return axpr::MakeGlobalNaiveClassOps<Self>(cls);
+}
+
+}  // namespace ap::kernel_dispatch

paddle/ap/include/kernel_dispatch/dispatch_raw_ctx.h

@@ -0,0 +1,73 @@
+// 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_module/data_type.h"
+#include "paddle/ap/include/kernel_dispatch/arg_value.h"
+#include "paddle/ap/include/kernel_dispatch/device_ctx.h"
+#include "paddle/ap/include/kernel_dispatch/typed_buffer.h"
+#include "paddle/ap/include/rt_module/module.h"
+#include "paddle/phi/core/dense_tensor.h"
+
+namespace phi {
+
+class DenseTensor;
+
+}
+
+namespace ap::kernel_dispatch {
+
+using RtModuleImpl = std::variant<std::shared_ptr<const rt_module::Module>>;
+
+struct RtModule : public RtModuleImpl {
+  using RtModuleImpl::RtModuleImpl;
+  ADT_DEFINE_VARIANT_METHODS(RtModuleImpl);
+};
+
+template <typename ValueT>
+struct DispatchRawCtxImpl {
+  DeviceCtx device_ctx;
+  adt::List<ValueT> inputs;
+  adt::List<ValueT> outputs;
+  RtModule rt_module;
+
+  bool operator==(const DispatchRawCtxImpl& other) const {
+    return &other == this;
+  }
+
+  Result<adt::Ok> LaunchCudaKernel(
+      const std::string& func_name,
+      int64_t num_blocks,
+      int64_t num_threads,
+      const adt::List<ArgValue>& kernel_args) const;
+};
+
+template <typename ValueT>
+ADT_DEFINE_RC(DispatchRawCtx, DispatchRawCtxImpl<ValueT>);
+
+}  // namespace ap::kernel_dispatch
+
+namespace ap::axpr {
+
+template <typename ValueT>
+struct TypeImpl<ap::kernel_dispatch::DispatchRawCtx<ValueT>>
+    : public std::monostate {
+  using value_type = ap::kernel_dispatch::DispatchRawCtx<ValueT>;
+
+  const char* Name() const { return "DispatchRawCtx"; }
+};
+
+}  // namespace ap::axpr

paddle/ap/include/kernel_dispatch/mutable_tensor.h

@@ -0,0 +1,68 @@
+// 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_module/data_type.h"
+#include "paddle/ap/include/kernel_dispatch/typed_buffer.h"
+#include "paddle/phi/core/dense_tensor.h"
+
+namespace phi {
+
+class DenseTensor;
+
+}
+
+namespace ap::kernel_dispatch {
+
+using MutableTensorDataImpl = std::variant<phi::DenseTensor*, TypedBuffer>;
+struct MutableTensorData : public MutableTensorDataImpl {
+  using MutableTensorDataImpl::MutableTensorDataImpl;
+  ADT_DEFINE_VARIANT_METHODS(MutableTensorDataImpl);
+
+  template <typename T>
+  T* data() const {
+    return Match(
+        [](phi::DenseTensor* tensor) -> T* {
+          return reinterpret_cast<T*>(tensor->data());
+        },
+        [](const TypedBuffer& buffer) -> T* {
+          return reinterpret_cast<T*>(buffer->buffer);
+        });
+  }
+  void* data() const {
+    return Match(
+        [](phi::DenseTensor* tensor) -> void* { return tensor->data(); },
+        [](const TypedBuffer& buffer) -> void* { return buffer->buffer; });
+  }
+  phi::DataType dtype() const {
+    return Match([](phi::DenseTensor* tensor) { return tensor->dtype(); },
+                 [](const TypedBuffer& buffer) { return buffer->dtype; });
+  }
+};
+
+template <typename ValueT>
+struct MutableTensorImpl {
+  MutableTensorData tensor_data;
+  adt::List<ValueT> dims;
+
+  bool operator==(const MutableTensorImpl& other) const {
+    return other.tensor_data == this->tensor_data && other.dims == this->dims;
+  }
+};
+template <typename ValueT>
+ADT_DEFINE_RC(MutableTensor, MutableTensorImpl<ValueT>);
+
+}  // namespace ap::kernel_dispatch

paddle/ap/include/kernel_dispatch/mutable_tensor_method_class.h

@@ -0,0 +1,110 @@
+// 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/data_type_util.h"
+#include "paddle/ap/include/axpr/method_class.h"
+#include "paddle/ap/include/axpr/naive_class_ops.h"
+#include "paddle/ap/include/kernel_dispatch/mutable_tensor.h"
+
+namespace ap::kernel_dispatch {
+
+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::DataValue;
+using ap::axpr::Method;
+using ap::axpr::MethodClass;
+using ap::axpr::PointerType;
+using ap::axpr::PointerValue;
+
+namespace detail {
+
+template <typename Val>
+Result<Val> MutableTensorShapeGetAttr(const MutableTensor<Val>& tensor,
+                                      const std::string&) {
+  return tensor->dims;
+}
+
+template <typename T>
+T* GetMutableTensorDataPtr(const ap::axpr::CppDataType<T>&,
+                           const MutableTensorData& tensor) {
+  return tensor.template data<T>();
+}
+
+template <typename Val>
+Result<Val> MutableTensorDataGetAttr(const MutableTensor<Val>& tensor,
+                                     const std::string&) {
+  phi::DataType dtype = tensor->tensor_data.dtype();
+  const auto& data_type = ap::axpr::GetDataTypeFromPhiDataType(dtype);
+  ADT_RETURN_IF_ERR(data_type);
+  return data_type.GetOkValue().Match(
+      [&](const adt::Undefined&) -> Result<Val> {
+        return TypeError{"dtype is invalid."};
+      },
+      [&](const auto& impl) -> Result<Val> {
+        return PointerValue{GetMutableTensorDataPtr(impl, tensor->tensor_data)};
+      });
+}
+
+template <typename Val>
+using MutableTensorGetAttrT = Result<Val> (*)(const MutableTensor<Val>& tensor,
+                                              const std::string&);
+
+template <typename Val>
+Result<Val> TensorGetAttr(const MutableTensor<Val>& tensor,
+                          const std::string& name) {
+  static const std::unordered_map<std::string, MutableTensorGetAttrT<Val>> map{
+      {"shape", &MutableTensorShapeGetAttr<Val>},
+      {"data_ptr", &MutableTensorDataGetAttr<Val>},
+  };
+  const auto& iter = map.find(name);
+  if (iter == map.end()) {
+    return AttributeError{std::string("'Tensor' has no attribute '") + name +
+                          "'"};
+  }
+  return iter->second(tensor, name);
+}
+
+}  // namespace detail
+
+template <typename ValueT>
+struct MutableTensorMethodClass {
+  using Self = MutableTensor<ValueT>;
+
+  static adt::Result<ValueT> GetAttr(const ValueT& obj_val,
+                                     const std::vector<ValueT>& args) {
+    ADT_CHECK(args.size() == 1);
+    const auto& attr_name_val = args.at(0);
+    ADT_LET_CONST_REF(obj, axpr::Get<MutableTensor<ValueT>>(obj_val));
+    ADT_LET_CONST_REF(attr_name, attr_name_val.template TryGet<std::string>());
+    return detail::TensorGetAttr<ValueT>(obj, attr_name);
+  }
+};
+
+template <typename ValueT>
+axpr::TypeImpl<axpr::BuiltinClassInstance<ValueT>> GetMutableTensorClass() {
+  using Methods = MutableTensorMethodClass<ValueT>;
+  static auto cls(axpr::MakeBuiltinClass<ValueT>(
+      "MutableTensor",
+      [&](const auto& Yield) { Yield("__getattr__", &Methods::GetAttr); }));
+  using Self = typename Methods::Self;
+  return axpr::MakeGlobalNaiveClassOps<Self>(cls);
+}
+
+}  // namespace ap::kernel_dispatch

paddle/ap/include/kernel_dispatch/typed_buffer.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/code_module/adt.h"
+#include "paddle/ap/include/code_module/data_type.h"
+#include "paddle/phi/core/dense_tensor.h"
+
+namespace phi {
+
+class DenseTensor;
+
+}
+
+namespace ap::kernel_dispatch {
+
+struct TypedBufferImpl {
+  void* buffer;
+  phi::DataType dtype;
+  size_t size;
+
+  bool operator==(const TypedBufferImpl& other) const {
+    return other.buffer == this->buffer && other.dtype == this->dtype &&
+           other.size == this->size;
+  }
+};
+ADT_DEFINE_RC(TypedBuffer, TypedBufferImpl);
+
+}  // namespace ap::kernel_dispatch

paddle/ap/include/kernel_dispatch/value.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/adt/adt.h"
+#include "paddle/ap/include/axpr/builtin_serializable_attr_map.h"
+#include "paddle/ap/include/axpr/value.h"
+#include "paddle/ap/include/code_module/data_type.h"
+#include "paddle/ap/include/kernel_dispatch/arg_value.h"
+#include "paddle/ap/include/kernel_dispatch/const_tensor.h"
+#include "paddle/ap/include/kernel_dispatch/dispatch_ctx.h"
+#include "paddle/ap/include/kernel_dispatch/mutable_tensor.h"
+#include "paddle/ap/include/kernel_dispatch/typed_buffer.h"
+
+namespace ap::kernel_dispatch {
+
+using axpr::Value;
+
+using Val = Value;
+
+}  // namespace ap::kernel_dispatch