145 lines
4.8 KiB
C++
145 lines
4.8 KiB
C++
// Copyright (c) 2021 CINN 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.
|
|
|
|
#include "paddle/cinn/ir/operation.h"
|
|
|
|
#include <memory>
|
|
|
|
#include "paddle/cinn/common/common.h"
|
|
|
|
namespace cinn {
|
|
namespace ir {
|
|
|
|
Operation PlaceholderOp::Make(const std::string &name,
|
|
const std::vector<Expr> &shape,
|
|
Type dtype) {
|
|
auto n = make_shared<PlaceholderOp>();
|
|
n->name = name;
|
|
n->shape = shape;
|
|
n->set_type(dtype);
|
|
return Operation(n);
|
|
}
|
|
|
|
Operation PlaceholderOp::Make(const std::string &name,
|
|
const std::vector<Dim> &sym_shape,
|
|
Type dtype) {
|
|
auto n = make_shared<PlaceholderOp>();
|
|
n->name = name;
|
|
n->sym_shape = sym_shape;
|
|
for (int i = 0; i < sym_shape.size(); i++) {
|
|
n->shape.emplace_back(sym_shape[i]->dim_expr);
|
|
}
|
|
n->set_type(dtype);
|
|
return Operation(n);
|
|
}
|
|
|
|
const char *PlaceholderOp::func_type() const { return "placeholder_op"; }
|
|
|
|
const char *ComputeOp::func_type() const { return "compute_op"; }
|
|
|
|
Operation ComputeOp::Make(const std::string &name,
|
|
ComputeOp::handle_t handle,
|
|
const std::vector<Expr> &shape,
|
|
const std::vector<Expr> &domain,
|
|
const std::vector<Var> &reduce_axis,
|
|
const std::map<std::string, IrNodeRef> &attrs,
|
|
const std::string &tag) {
|
|
auto n = make_shared<ComputeOp>();
|
|
n->name = name;
|
|
n->producer_fn = handle;
|
|
n->shape = domain;
|
|
n->reduce_axis = reduce_axis;
|
|
n->tag = tag;
|
|
n->attrs = attrs;
|
|
n->axis = cinn::common::GenDefaultAxis(domain.size());
|
|
std::vector<Expr> tmp_axis;
|
|
for (auto &x : n->axis) {
|
|
tmp_axis.push_back(x);
|
|
}
|
|
n->body = {handle(tmp_axis)};
|
|
n->reduce_axis = reduce_axis;
|
|
return Operation(n);
|
|
}
|
|
|
|
Operation CallOp::Make(const std::string &call_target, Expr call_op) {
|
|
auto n = make_shared<CallOp>();
|
|
n->call_expr = call_op;
|
|
return Operation(n);
|
|
}
|
|
|
|
Operation PrecedingViewOp::Make(const Tensor &tensor, int preceding_axis) {
|
|
return Operation();
|
|
}
|
|
|
|
const char *PrecedingViewOp::func_type() const {
|
|
return PrecedingViewOp::__func_type__;
|
|
}
|
|
|
|
const char *CallOp::func_type() const { return __func_type__; }
|
|
|
|
const char *ComputeOp::__func_type__ = "compute_op";
|
|
const char *PlaceholderOp::__func_type__ = "placeholder_op";
|
|
const char *CallOp::__func_type__ = "call_op";
|
|
|
|
const std::string &CallOp::target() const {
|
|
auto *call = call_expr.As<ir::Call>();
|
|
PADDLE_ENFORCE_NOT_NULL(call,
|
|
::common::errors::InvalidArgument(
|
|
"The 'call_expr' must be of type 'ir::Call'."));
|
|
return call->name;
|
|
}
|
|
std::vector<Expr> &CallOp::write_args() {
|
|
auto *call = call_expr.As<ir::Call>();
|
|
PADDLE_ENFORCE_NOT_NULL(call,
|
|
::common::errors::InvalidArgument(
|
|
"The 'call_expr' must be of type 'ir::Call'."));
|
|
return call->write_args;
|
|
}
|
|
std::vector<Expr> &CallOp::read_args() {
|
|
auto *call = call_expr.As<ir::Call>();
|
|
PADDLE_ENFORCE_NOT_NULL(call,
|
|
::common::errors::InvalidArgument(
|
|
"The 'call_expr' must be of type 'ir::Call'."));
|
|
return call->read_args;
|
|
}
|
|
const std::vector<Expr> &CallOp::write_args() const {
|
|
auto *call = call_expr.As<ir::Call>();
|
|
PADDLE_ENFORCE_NOT_NULL(call,
|
|
::common::errors::InvalidArgument(
|
|
"The 'call_expr' must be of type 'ir::Call'."));
|
|
return call->write_args;
|
|
}
|
|
const std::vector<Expr> &CallOp::read_args() const {
|
|
auto *call = call_expr.As<ir::Call>();
|
|
PADDLE_ENFORCE_NOT_NULL(call,
|
|
::common::errors::InvalidArgument(
|
|
"The 'call_expr' must be of type 'ir::Call'."));
|
|
return call->read_args;
|
|
}
|
|
std::vector<Expr> CallOp::args() const {
|
|
std::vector<Expr> args;
|
|
auto &rargs = read_args();
|
|
auto &wargs = write_args();
|
|
args.insert(std::end(args), rargs.begin(), rargs.end());
|
|
args.insert(std::end(args), wargs.begin(), wargs.end());
|
|
return args;
|
|
}
|
|
const char *PrecedingViewOp::__func_type__ = "preceding_view_op";
|
|
|
|
const char *BufferShareOp::__func_type__ = "buffer_share_op";
|
|
const char *BufferShareOp::func_type() const { return __func_type__; }
|
|
|
|
} // namespace ir
|
|
} // namespace cinn
|