140 lines
4.0 KiB
C++
140 lines
4.0 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/module.h"
|
|
|
|
#include <memory>
|
|
|
|
#include "paddle/cinn/ir/ir_printer.h"
|
|
#include "paddle/cinn/optim/ir_simplify.h"
|
|
#include "paddle/cinn/optim/optimize.h"
|
|
#include "paddle/common/enforce.h"
|
|
|
|
namespace cinn {
|
|
namespace ir {
|
|
|
|
void Module::Builder::AddFunction(ir::LoweredFunc func) {
|
|
module_->functions.push_back(func);
|
|
}
|
|
|
|
void Module::Builder::AddFunctionWithoutOptim(const ir::LoweredFunc &func) {
|
|
module_->functions.push_back(func);
|
|
}
|
|
|
|
std::optional<int> GetDataAlignmentImpl(common::UnknownArch arch) {
|
|
return std::nullopt;
|
|
}
|
|
|
|
std::optional<int> GetDataAlignmentImpl(common::X86Arch arch) { return 32; }
|
|
|
|
std::optional<int> GetDataAlignmentImpl(common::ARMArch arch) {
|
|
return std::nullopt;
|
|
}
|
|
|
|
std::optional<int> GetDataAlignmentImpl(common::NVGPUArch) {
|
|
return std::nullopt;
|
|
}
|
|
std::optional<int> GetDataAlignmentImpl(common::CustomDeviceArch) {
|
|
return std::nullopt;
|
|
}
|
|
|
|
std::optional<int> GetDataAlignmentImpl(common::HygonDCUArchHIP arch) {
|
|
return std::nullopt;
|
|
}
|
|
|
|
std::optional<int> GetDataAlignmentImpl(common::HygonDCUArchSYCL arch) {
|
|
return std::nullopt;
|
|
}
|
|
|
|
std::optional<int> GetDataAlignment(common::Arch arch) {
|
|
return std::visit([](const auto &impl) { return GetDataAlignmentImpl(impl); },
|
|
arch.variant());
|
|
}
|
|
|
|
void Module::Builder::AddBuffer(ir::Buffer buffer) {
|
|
PADDLE_ENFORCE_EQ(
|
|
buffer->target.defined(),
|
|
true,
|
|
::common::errors::InvalidArgument(
|
|
"The target of buffer [%s] is undefined. Please define the target.",
|
|
buffer->name));
|
|
if (std::find_if(
|
|
module_->buffers.begin(), module_->buffers.end(), [&](const Expr &x) {
|
|
return x.as_buffer()->name == buffer->name;
|
|
}) == std::end(module_->buffers)) {
|
|
module_->buffers.push_back(buffer);
|
|
if (auto alignment = GetDataAlignment(module_->target.arch)) {
|
|
module_->buffers.back().as_buffer()->data_alignment = alignment.value();
|
|
}
|
|
}
|
|
}
|
|
|
|
void Module::Builder::AddPredicate(ir::Expr predicate) {
|
|
module_->predicates.push_back(predicate);
|
|
}
|
|
|
|
void Module::Builder::AddPriority(int priority) {
|
|
module_->priorities.push_back(priority);
|
|
}
|
|
|
|
void Module::Builder::SetInferShapeFunc(ir::LoweredFunc infer_shape_func) {
|
|
module_->infer_shape_func = infer_shape_func;
|
|
}
|
|
|
|
void Module::Builder::Clear() {
|
|
module_->buffers.clear();
|
|
module_->functions.clear();
|
|
module_->submodules.clear();
|
|
module_->predicates.clear();
|
|
}
|
|
|
|
common::Arch Module::Builder::GetTargetArch() { return module_->target.arch; }
|
|
|
|
Module Module::Builder::Build() {
|
|
if (module_->functions.empty()) {
|
|
VLOG(1) << "Module has no functions";
|
|
}
|
|
|
|
auto res = ir::Module(module_.get());
|
|
return res;
|
|
}
|
|
|
|
ir::_Module_ *Module::self() { return p_->as<ir::_Module_>(); }
|
|
const ir::_Module_ *Module::self() const { return p_->as<ir::_Module_>(); }
|
|
|
|
const Target &Module::target() const { return self()->target; }
|
|
|
|
const std::string &Module::name() const { return self()->name; }
|
|
|
|
std::vector<ir::Buffer> Module::buffers() const {
|
|
std::vector<ir::Buffer> buffers;
|
|
for (auto &buffer : self()->buffers) {
|
|
buffers.emplace_back(buffer.as_buffer_ref());
|
|
}
|
|
return buffers;
|
|
}
|
|
|
|
const std::vector<ir::LoweredFunc> &Module::functions() const {
|
|
return self()->functions;
|
|
}
|
|
|
|
const std::vector<Module> &Module::submodules() const {
|
|
return self()->submodules;
|
|
}
|
|
|
|
void Module::Compile(const backends::Outputs &outputs) const {}
|
|
|
|
} // namespace ir
|
|
} // namespace cinn
|