Files
paddlepaddle--paddle/paddle/cinn/ir/module.cc
T
2026-07-13 12:40:42 +08:00

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