chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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
|
||||
Reference in New Issue
Block a user