wehub-resource-sync
247 lines
9.7 KiB
C++
247 lines
9.7 KiB
C++
// 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
|