chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,610 @@
|
||||
// Copyright (c) 2024 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/backends/codegen_gpu_dev.h"
|
||||
|
||||
#include <glog/logging.h>
|
||||
#include <paddle/cinn/utils/string.h>
|
||||
|
||||
#include <fstream>
|
||||
#include <set>
|
||||
#include <unordered_set>
|
||||
|
||||
#include "paddle/cinn/common/ir_util.h"
|
||||
#include "paddle/cinn/ir/op/ir_operators.h"
|
||||
#include "paddle/cinn/ir/utils/ir_verify.h"
|
||||
#include "paddle/cinn/optim/ir_simplify.h"
|
||||
#include "paddle/common/enforce.h"
|
||||
#include "paddle/common/errors.h"
|
||||
namespace cinn {
|
||||
namespace backends {
|
||||
|
||||
CodeGenGpuDev::CodeGenGpuDev(Target target) : CodeGenC(target) {}
|
||||
|
||||
std::string CodeGenGpuDev::Compile(const ir::Module &module, bool use_rtc) {
|
||||
use_rtc_ = use_rtc;
|
||||
auto source = Compile(module, OutputKind::CImpl);
|
||||
|
||||
return source;
|
||||
}
|
||||
|
||||
void CodeGenGpuDev::Compile(const ir::Module &module, const Outputs &outputs) {
|
||||
ir::ir_utils::IrVerify(module.As<ir::_Module_>());
|
||||
|
||||
CodeGenC::inline_builtin_codes_ = false;
|
||||
if (!outputs.c_header_name.empty()) {
|
||||
auto source = Compile(module, OutputKind::CHeader);
|
||||
str_ = "";
|
||||
std::ofstream file(outputs.c_header_name);
|
||||
PADDLE_ENFORCE_EQ(file.is_open(),
|
||||
true,
|
||||
::common::errors::PreconditionNotMet(
|
||||
"failed to open file %s", outputs.c_header_name));
|
||||
file << source;
|
||||
file.close();
|
||||
VLOG(5) << "Output C header to file " << outputs.c_header_name;
|
||||
}
|
||||
|
||||
if (!outputs.cuda_source_name.empty()) {
|
||||
auto source = Compile(module, OutputKind::CImpl);
|
||||
str_ = "";
|
||||
std::ofstream file(outputs.cuda_source_name);
|
||||
PADDLE_ENFORCE_EQ(file.is_open(),
|
||||
true,
|
||||
::common::errors::PreconditionNotMet(
|
||||
"failed to open file %s", outputs.cuda_source_name));
|
||||
file << source;
|
||||
file.close();
|
||||
VLOG(5) << "Output C source to file " << outputs.cuda_source_name;
|
||||
}
|
||||
}
|
||||
|
||||
void CodeGenGpuDev::Compile(const ir::LoweredFunc &func) {
|
||||
dyn_shared_mem_offset_ = Expr(-1);
|
||||
Visit(func.As<ir::_LoweredFunc_>());
|
||||
}
|
||||
|
||||
std::vector<ir::stmt::StmtRef> CodeGenGpuDev::GenerateBufferAliasStmts(
|
||||
const ir::_LoweredFunc_ *op, const std::vector<ir::Buffer> &temp_buffers) {
|
||||
std::set<ir::Buffer> temp_buffer_set(temp_buffers.begin(),
|
||||
temp_buffers.end());
|
||||
// prepare temp buffer alias
|
||||
std::vector<ir::stmt::StmtRef> buffer_alias;
|
||||
auto tensors =
|
||||
ir::ir_utils::CollectIRNodes(op->body_block, [&](const Expr *x) {
|
||||
return x->as_tensor() && x->as_tensor()->buffer.defined() &&
|
||||
temp_buffer_set.count(x->as_tensor()->buffer);
|
||||
});
|
||||
|
||||
// unique tensors
|
||||
std::set<ir::Tensor> unique_tensors;
|
||||
for (auto &e : tensors) {
|
||||
unique_tensors.insert(e.as_tensor_ref());
|
||||
}
|
||||
|
||||
for (auto &t : unique_tensors) {
|
||||
auto tensor_type = t->type();
|
||||
auto tensor_ptr_type = tensor_type;
|
||||
tensor_ptr_type.set_cpp_handle();
|
||||
|
||||
auto buffer_type = t->buffer->dtype;
|
||||
auto buffer_ptr_type = buffer_type;
|
||||
buffer_ptr_type.set_cpp_handle();
|
||||
|
||||
Expr t_var = Var(t->name, tensor_ptr_type);
|
||||
Expr buf_var = Var(t->buffer->name, buffer_ptr_type);
|
||||
|
||||
// A tensor and its buffer may have different types when multiple tensors
|
||||
// share the same buffer. In this case, add a Cast before aliasing.
|
||||
if (tensor_type != buffer_type) {
|
||||
buf_var = common::cast(buf_var, tensor_ptr_type);
|
||||
}
|
||||
|
||||
buffer_alias.push_back(ir::stmt::Let(t_var, buf_var));
|
||||
}
|
||||
|
||||
return buffer_alias;
|
||||
}
|
||||
|
||||
std::vector<ir::stmt::StmtRef> CodeGenGpuDev::FilterDeallocTempBuffers(
|
||||
const std::vector<ir::stmt::StmtRef> &frees) {
|
||||
std::vector<ir::stmt::StmtRef> filtered;
|
||||
for (const auto &free : frees) {
|
||||
PADDLE_ENFORCE_EQ(
|
||||
free.isa<ir::stmt::Free>(),
|
||||
true,
|
||||
::common::errors::InvalidArgument("Free is not a free node"));
|
||||
const auto op = free.as<ir::stmt::Free>();
|
||||
bool has_symbolic_constant = false;
|
||||
const ir::_Buffer_ *buffer = op->destination().As<ir::_Buffer_>();
|
||||
for (Expr shape : buffer->shape) {
|
||||
shape = optim::ArithSimplify(shape);
|
||||
ir::ir_utils::CollectIRNodes(shape, [&](const Expr *x) {
|
||||
if (x->as_var()) {
|
||||
PADDLE_ENFORCE_EQ(
|
||||
x->as_var()->is_symbolic_constant,
|
||||
true,
|
||||
::common::errors::PreconditionNotMet(
|
||||
"var in buffer shape must be symbolic constant."));
|
||||
has_symbolic_constant = true;
|
||||
}
|
||||
return false;
|
||||
});
|
||||
}
|
||||
if (has_symbolic_constant &&
|
||||
buffer->memory_type == ir::MemoryType::GPULocal) {
|
||||
filtered.emplace_back(free);
|
||||
}
|
||||
}
|
||||
return filtered;
|
||||
}
|
||||
|
||||
void CodeGenGpuDev::Visit(const ir::_LoweredFunc_ *op) {
|
||||
// clear names valid within scope when enter a new function
|
||||
vectorized_tensor_names_.clear();
|
||||
dynamic_shape_map_.clear();
|
||||
for (const auto &arg : op->args) {
|
||||
if (arg.is_var()) {
|
||||
dynamic_shape_map_.emplace(arg.name(), arg.type());
|
||||
}
|
||||
}
|
||||
str_ += "__global__\n";
|
||||
|
||||
PrintFunctionDeclaration(op);
|
||||
str_ += "\n";
|
||||
|
||||
DoIndent();
|
||||
|
||||
std::vector<ir::stmt::StmtRef> new_body_stmts;
|
||||
|
||||
auto axis_range_assumption_stmts = op->PrepareAxisRangeAssumptionStmts();
|
||||
auto alloca_temp_buffer_stmts = op->PrepareAllocTempBufferStmts();
|
||||
auto temp_buffer_alias_stmts = GenerateBufferAliasStmts(op, op->temp_bufs);
|
||||
auto alias_var_stmts = op->CudaAliasVarStmts();
|
||||
auto dealloc_temp_buffer_stmts =
|
||||
FilterDeallocTempBuffers(op->PrepareDeallocTempBufferStmts());
|
||||
|
||||
#define APPEND_TO_NEW_BODY_STMTS(field__) \
|
||||
new_body_stmts.insert( \
|
||||
std::end(new_body_stmts), std::begin(field__), std::end(field__));
|
||||
APPEND_TO_NEW_BODY_STMTS(axis_range_assumption_stmts)
|
||||
APPEND_TO_NEW_BODY_STMTS(alloca_temp_buffer_stmts)
|
||||
APPEND_TO_NEW_BODY_STMTS(temp_buffer_alias_stmts)
|
||||
APPEND_TO_NEW_BODY_STMTS(alias_var_stmts)
|
||||
APPEND_TO_NEW_BODY_STMTS(op->body_block->stmts())
|
||||
APPEND_TO_NEW_BODY_STMTS(dealloc_temp_buffer_stmts);
|
||||
|
||||
ir::stmt::BlockRef func_body_block = ir::stmt::BlockRef(new_body_stmts);
|
||||
|
||||
// Use ir_simplify when pass updated.
|
||||
// optim::SimplifyUnitBlock(&func_body);
|
||||
// // Make sure that the function's body is wrapped by a block
|
||||
// if (!func_body.As<ir::Block>()) {
|
||||
// func_body = ir::Block::Make({func_body});
|
||||
// }
|
||||
|
||||
CodeGenC::VisitBlock(func_body_block);
|
||||
}
|
||||
|
||||
void CodeGenGpuDev::VisitStmt(const ir::stmt::Free &stmt) {
|
||||
str_ += "delete [] ";
|
||||
str_ += stmt->destination().As<ir::_Buffer_>()->name;
|
||||
str_ += ";\n";
|
||||
}
|
||||
|
||||
void CodeGenGpuDev::Visit(const ir::_Var_ *op) {
|
||||
if (utils::StartsWith(op->name, "threadIdx") ||
|
||||
utils::StartsWith(op->name, "blockIdx")) {
|
||||
str_ += "(int)";
|
||||
str_ += op->name;
|
||||
} else {
|
||||
str_ += op->name;
|
||||
}
|
||||
}
|
||||
|
||||
void CodeGenGpuDev::VisitStmt(const ir::stmt::Alloc &stmt) {
|
||||
PrintTempBufferCreation(stmt->destination().as_buffer_ref());
|
||||
}
|
||||
|
||||
void CodeGenGpuDev::Visit(const ir::Min *op) {
|
||||
str_ += "min(";
|
||||
ir::Expr a = op->a(), b = op->b();
|
||||
auto [unify_bit, both_dyn] =
|
||||
common::UnifiedOperandTypeBits(&dynamic_shape_map_, op);
|
||||
ProcessMinMaxOperand(&a, &b, unify_bit, both_dyn);
|
||||
IrPrinter::Visit(a);
|
||||
str_ += ", ";
|
||||
IrPrinter::Visit(b);
|
||||
str_ += ")";
|
||||
}
|
||||
|
||||
void CodeGenGpuDev::Visit(const ir::Max *op) {
|
||||
str_ += "max(";
|
||||
ir::Expr a = op->a(), b = op->b();
|
||||
auto [unify_bit, both_dyn] =
|
||||
common::UnifiedOperandTypeBits(&dynamic_shape_map_, op);
|
||||
ProcessMinMaxOperand(&a, &b, unify_bit, both_dyn);
|
||||
IrPrinter::Visit(a);
|
||||
str_ += ", ";
|
||||
IrPrinter::Visit(b);
|
||||
str_ += ")";
|
||||
}
|
||||
|
||||
void CodeGenGpuDev::PrintFunctionDeclaration(const ir::_LoweredFunc_ *op) {
|
||||
str_ += "void ";
|
||||
if (op->cuda_axis_info.valid()) {
|
||||
int max_threads_per_block = op->cuda_axis_info.max_threads_per_block();
|
||||
if (max_threads_per_block > 0) {
|
||||
str_ += "__launch_bounds__(";
|
||||
str_ += std::to_string(max_threads_per_block);
|
||||
int min_blocks_per_sm = op->cuda_axis_info.min_blocks_per_sm();
|
||||
if (min_blocks_per_sm > 0) {
|
||||
str_ += ", ";
|
||||
str_ += std::to_string(min_blocks_per_sm);
|
||||
}
|
||||
str_ += ") ";
|
||||
}
|
||||
}
|
||||
|
||||
str_ += op->name;
|
||||
str_ += "(";
|
||||
for (int i = 0; i < op->args.size() - 1; i++) {
|
||||
auto &arg = op->args[i];
|
||||
PrintFuncArg(arg);
|
||||
str_ += ", ";
|
||||
}
|
||||
if (!op->args.empty()) {
|
||||
PrintFuncArg(op->args.back());
|
||||
}
|
||||
str_ += ")";
|
||||
}
|
||||
|
||||
void CodeGenGpuDev::PrintFuncArg(const ir::Argument &arg) {
|
||||
if (arg.is_buffer()) {
|
||||
// In CUDA/HIP kernel, only primitive type is supported, so we replace the
|
||||
// buffer with T*j
|
||||
if (arg.is_input()) str_ += "const ";
|
||||
str_ += GetTypeRepr(arg.buffer_arg()->dtype);
|
||||
str_ += "* ";
|
||||
str_ += kCKeywordRestrict;
|
||||
str_ += " ";
|
||||
str_ += ir::BufferGetTensorName(arg.buffer_arg().As<ir::_Buffer_>());
|
||||
} else if (arg.is_var()) {
|
||||
if (arg.var_arg()->type().is_cpp_handle()) {
|
||||
str_ += kCKeywordRestrict;
|
||||
}
|
||||
str_ += GetTypeRepr(arg.type());
|
||||
str_ += " ";
|
||||
str_ += arg.name();
|
||||
} else {
|
||||
CINN_NOT_IMPLEMENTED
|
||||
}
|
||||
}
|
||||
|
||||
void CodeGenGpuDev::PrintBuiltinCodes() {}
|
||||
|
||||
std::string CodeGenGpuDev::Compile(const ir::Module &module,
|
||||
CodeGenC::OutputKind output_kind) {
|
||||
if (output_kind == OutputKind::CHeader) {
|
||||
GenerateHeaderFile(module);
|
||||
} else if (output_kind == OutputKind::CImpl) {
|
||||
if (use_rtc_) {
|
||||
str_ += "\nextern \"C\" {\n\n";
|
||||
}
|
||||
|
||||
PrintBuiltinCodes();
|
||||
|
||||
for (auto &func : module.functions()) {
|
||||
Compile(func);
|
||||
}
|
||||
} else {
|
||||
PADDLE_THROW(::common::errors::InvalidArgument("Not supported OutputKind"));
|
||||
}
|
||||
|
||||
if (use_rtc_) {
|
||||
str_ += "\n\n}";
|
||||
}
|
||||
return str_;
|
||||
}
|
||||
|
||||
void CodeGenGpuDev::PrintTempBufferCreation(const ir::Buffer &buffer) {
|
||||
PADDLE_ENFORCE_NE(
|
||||
buffer->type(),
|
||||
Void(),
|
||||
::common::errors::InvalidArgument("buffer type should not be void"));
|
||||
// Calculate buffer size and determine if it contains a symbolic constant
|
||||
Expr buffer_size(1);
|
||||
for (int i = 0; i < buffer->shape.size(); i++) {
|
||||
buffer_size = buffer_size * buffer->shape[i];
|
||||
}
|
||||
buffer_size = optim::ArithSimplify(buffer_size);
|
||||
bool has_symbolic_constant = false;
|
||||
ir::ir_utils::CollectIRNodes(buffer_size, [&](const Expr *x) {
|
||||
if (x->as_var()) {
|
||||
PADDLE_ENFORCE_EQ(x->as_var()->is_symbolic_constant,
|
||||
true,
|
||||
::common::errors::PreconditionNotMet(
|
||||
"var in buffer shape must be symbolic constant."));
|
||||
has_symbolic_constant = true;
|
||||
}
|
||||
return false;
|
||||
});
|
||||
|
||||
if (buffer->memory_type == ir::MemoryType::GPUShared) {
|
||||
if (MathEqual(dyn_shared_mem_offset_, Expr(-1))) {
|
||||
// The first shared memory buffer, uint8_t as a byte
|
||||
str_ += "extern __shared__ uint8_t dyn_shared_buffer[];\n ";
|
||||
dyn_shared_mem_offset_ = Expr(0);
|
||||
}
|
||||
std::string type_name = GetTypeRepr(buffer->dtype);
|
||||
str_ += type_name;
|
||||
str_ += " *";
|
||||
str_ += buffer->name;
|
||||
str_ += " = (";
|
||||
str_ += type_name;
|
||||
str_ += "*)&dyn_shared_buffer[ ";
|
||||
IrPrinter::Visit(dyn_shared_mem_offset_);
|
||||
str_ += " ]";
|
||||
|
||||
int type_bytes = buffer->dtype.bytes();
|
||||
dyn_shared_mem_offset_ =
|
||||
dyn_shared_mem_offset_ + buffer_size * Expr(type_bytes);
|
||||
dyn_shared_mem_offset_ = optim::ArithSimplify(dyn_shared_mem_offset_);
|
||||
VLOG(6) << "dyn_shared_mem_offset_ = " << dyn_shared_mem_offset_;
|
||||
} else if (buffer->memory_type == ir::MemoryType::GPULocal) {
|
||||
// print func of static allocation
|
||||
auto print_gpu_memory = [&](const std::string &mark) {
|
||||
str_ += mark;
|
||||
str_ += GetTypeRepr(buffer->dtype);
|
||||
str_ += " ";
|
||||
str_ += buffer->name;
|
||||
str_ += " ";
|
||||
|
||||
str_ += "[ ";
|
||||
IrPrinter::Visit(buffer_size);
|
||||
str_ += " ]";
|
||||
};
|
||||
// print func of dynamic allocation
|
||||
auto print_gpu_local_memory_dynamic_allocation = [&]() {
|
||||
str_ += GetTypeRepr(buffer->dtype);
|
||||
str_ += " *";
|
||||
str_ += buffer->name;
|
||||
str_ += " = new ";
|
||||
str_ += GetTypeRepr(buffer->dtype);
|
||||
str_ += "[ ";
|
||||
IrPrinter::Visit(buffer_size);
|
||||
str_ += " ]";
|
||||
};
|
||||
if (has_symbolic_constant) {
|
||||
print_gpu_local_memory_dynamic_allocation();
|
||||
} else {
|
||||
print_gpu_memory("");
|
||||
}
|
||||
} else {
|
||||
std::stringstream ss;
|
||||
ss << "CUDA/HIP device codegen not support memory " << buffer->name
|
||||
<< ", type " << buffer->memory_type;
|
||||
PADDLE_THROW(::common::errors::InvalidArgument(
|
||||
"CUDA/HIP codegen error in CINN: %s", ss.str()));
|
||||
}
|
||||
}
|
||||
|
||||
void CodeGenGpuDev::Visit(const ir::Call *op) {
|
||||
str_ += op->name;
|
||||
str_ += "(";
|
||||
|
||||
if (!op->read_args.empty()) {
|
||||
for (int i = 0; i < op->read_args.size() - 1; i++) {
|
||||
auto &arg = op->read_args[i];
|
||||
if (arg.as_tensor()) {
|
||||
str_ += arg.as_tensor()->name;
|
||||
str_ += ", ";
|
||||
} else {
|
||||
IrPrinter::Visit(arg);
|
||||
str_ += ", ";
|
||||
}
|
||||
}
|
||||
if (op->read_args.back().as_tensor()) {
|
||||
str_ += op->read_args.back().as_tensor()->name;
|
||||
} else {
|
||||
IrPrinter::Visit(op->read_args.back());
|
||||
}
|
||||
}
|
||||
|
||||
if (!op->write_args.empty()) {
|
||||
str_ += ", ";
|
||||
for (int i = 0; i < op->write_args.size() - 1; i++) {
|
||||
auto &arg = op->write_args[i];
|
||||
if (arg.as_tensor()) {
|
||||
str_ += arg.as_tensor()->name;
|
||||
str_ += ", ";
|
||||
} else {
|
||||
IrPrinter::Visit(arg);
|
||||
str_ += ", ";
|
||||
}
|
||||
}
|
||||
if (op->write_args.back().as_tensor()) {
|
||||
str_ += op->write_args.back().as_tensor()->name;
|
||||
} else {
|
||||
IrPrinter::Visit(op->write_args.back());
|
||||
}
|
||||
}
|
||||
|
||||
str_ += ")";
|
||||
}
|
||||
|
||||
void CodeGenGpuDev::VisitStmt(const ir::stmt::Let &stmt) {
|
||||
PADDLE_ENFORCE_EQ(
|
||||
stmt->type().valid(),
|
||||
true,
|
||||
::common::errors::PreconditionNotMet("Let op type must be valid."));
|
||||
// identify vectorized tensors by checking their dtypes are customized_type
|
||||
// with customized_type::kcuda_builtin_vector_t prefix, and save their names
|
||||
if (stmt->type().is_customized() &&
|
||||
utils::StartsWith(
|
||||
stmt->type().customized_type(),
|
||||
cinn::common::customized_type::kcuda_builtin_vector_t)) {
|
||||
str_ += GetTypeRepr(stmt->type());
|
||||
if (stmt->type().is_cpp_handle()) {
|
||||
str_ += " ";
|
||||
str_ += kCKeywordRestrict;
|
||||
}
|
||||
str_ += " ";
|
||||
IrPrinter::Visit(stmt->symbol());
|
||||
vectorized_tensor_names_.insert(utils::GetStreamCnt(stmt->symbol()));
|
||||
// skip "=0" in "half8 temp = 0;" since the operator= of half8 may not
|
||||
// overloaded.
|
||||
if (stmt->body().As<ir::IntImm>() &&
|
||||
stmt->body().As<ir::IntImm>()->value == 0) {
|
||||
return;
|
||||
}
|
||||
str_ += " = ";
|
||||
IrPrinter::Visit(stmt->body());
|
||||
} else {
|
||||
CodeGenC::VisitStmt(stmt);
|
||||
}
|
||||
}
|
||||
|
||||
bool CodeGenGpuDev::PrintBuiltinVectorAccess(const ir::LoadStoreAddrMnger *op,
|
||||
ir::Expr index_expr,
|
||||
bool is_store) {
|
||||
static constexpr char index2suffix[8] = {
|
||||
'x', 'y', 'z', 'w', 'v', 'u', 't', 's'};
|
||||
|
||||
// addr of op should be a place of tensor and the index is simple int number
|
||||
if (!op->is_addr_tensor() || !index_expr.As<ir::IntImm>()) {
|
||||
return false;
|
||||
}
|
||||
auto *tensor = op->tensor.As<ir::_Tensor_>();
|
||||
PADDLE_ENFORCE_NOT_NULL(
|
||||
tensor,
|
||||
::common::errors::InvalidArgument(
|
||||
"LoadStoreAddrMnger contains NULL tensor, which is an "
|
||||
"illegal argument"));
|
||||
|
||||
// identify vectorized tensors by their names
|
||||
if (!vectorized_tensor_names_.count(tensor->name)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// the index can't exceed the range of cuda/hip built-in vector type
|
||||
int index = index_expr.As<ir::IntImm>()->value;
|
||||
if (index < 0 || index >= 8) {
|
||||
return false;
|
||||
}
|
||||
if (is_store && tensor->type().is_cpp_handle()) {
|
||||
str_ += tensor->name;
|
||||
str_ += "[";
|
||||
str_ += std::to_string(index);
|
||||
str_ += "]";
|
||||
} else {
|
||||
str_ += tensor->name;
|
||||
str_ += (tensor->type().is_cpp_handle() ? "->" : ".");
|
||||
str_ += index2suffix[index];
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
bool CodeGenGpuDev::PrintBuiltinVectorAccess(const ir::stmt::Store &stmt,
|
||||
ir::Expr index_expr,
|
||||
bool is_store) {
|
||||
static constexpr char index2suffix[8] = {
|
||||
'x', 'y', 'z', 'w', 'v', 'u', 't', 's'};
|
||||
|
||||
// addr of op should be a place of tensor and the index is simple int number
|
||||
if (!stmt->is_addr_tensor() || !index_expr.As<ir::IntImm>()) {
|
||||
return false;
|
||||
}
|
||||
auto *tensor = stmt->tensor().As<ir::_Tensor_>();
|
||||
PADDLE_ENFORCE_NOT_NULL(
|
||||
tensor,
|
||||
::common::errors::InvalidArgument(
|
||||
"LoadStoreAddrMnger contains NULL tensor, which is an "
|
||||
"illegal argument"));
|
||||
|
||||
// identify vectorized tensors by their names
|
||||
if (!vectorized_tensor_names_.count(tensor->name)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// the index can't exceed the range of cuda/hip built-in vector type
|
||||
int index = index_expr.As<ir::IntImm>()->value;
|
||||
if (index < 0 || index >= 8) {
|
||||
return false;
|
||||
}
|
||||
if (is_store && tensor->type().is_cpp_handle()) {
|
||||
str_ += tensor->name;
|
||||
str_ += "[";
|
||||
str_ += std::to_string(index);
|
||||
str_ += "]";
|
||||
} else {
|
||||
str_ += tensor->name;
|
||||
str_ += (tensor->type().is_cpp_handle() ? "->" : ".");
|
||||
str_ += index2suffix[index];
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
void CodeGenGpuDev::Visit(const ir::Load *op) {
|
||||
// overload this visit function to especially deal with the case when it
|
||||
// accesses element at a cuda/hip built-in vector, others still resolve to
|
||||
// CodeGenC
|
||||
if (!PrintBuiltinVectorAccess(op, op->index(), false)) {
|
||||
CodeGenC::Visit(op);
|
||||
}
|
||||
}
|
||||
|
||||
void CodeGenGpuDev::VisitStmt(const ir::stmt::Store &stmt) {
|
||||
// overload this visit function to especially deal with the case when it
|
||||
// accesses element at a cuda/hip built-in vector, others still resolve to
|
||||
// CodeGenC
|
||||
if (PrintBuiltinVectorAccess(stmt, stmt->index(), true)) {
|
||||
str_ += " = ";
|
||||
IrPrinter::Visit(stmt->value());
|
||||
} else {
|
||||
CodeGenC::VisitStmt(stmt);
|
||||
}
|
||||
}
|
||||
|
||||
ir::Expr CalculateSharedMemory(const ir::Buffer &buffer) {
|
||||
Expr buffer_size(1);
|
||||
for (int i = 0; i < buffer->shape.size(); i++) {
|
||||
buffer_size = buffer_size * buffer->shape[i];
|
||||
}
|
||||
int type_bytes = buffer->dtype.bytes();
|
||||
return buffer_size * Expr(type_bytes);
|
||||
}
|
||||
|
||||
ir::Expr CalculateSharedMemory(const ir::LoweredFunc &func) {
|
||||
auto alloc_temp_buffers = func->PrepareAllocTempBufferStmts();
|
||||
ir::Expr shm_size{0};
|
||||
for (const auto &alloc : alloc_temp_buffers) {
|
||||
PADDLE_ENFORCE_EQ(
|
||||
alloc.isa<ir::stmt::Alloc>(),
|
||||
true,
|
||||
::common::errors::InvalidType("stmt is not a Alloc node"));
|
||||
PADDLE_ENFORCE_NOT_NULL(
|
||||
alloc.as<ir::stmt::Alloc>()->destination().as_buffer(),
|
||||
::common::errors::InvalidType("stmt is not a Buffer node"));
|
||||
|
||||
auto buffer = alloc.as<ir::stmt::Alloc>()->destination().as_buffer_ref();
|
||||
if (buffer->memory_type == ir::MemoryType::GPUShared) {
|
||||
shm_size = shm_size + CalculateSharedMemory(buffer);
|
||||
}
|
||||
}
|
||||
return optim::ArithSimplify(shm_size);
|
||||
}
|
||||
|
||||
} // namespace backends
|
||||
} // namespace cinn
|
||||
Reference in New Issue
Block a user