mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
109bc1047e
Summary: This is the initial skeleton for C++ codegen, it includes generations for Allocate and Free. Pull Request resolved: https://github.com/pytorch/pytorch/pull/51070 Test Plan: New unit tests are added to `test_cpp_codegen.cpp`. Reviewed By: ZolotukhinM Differential Revision: D26061818 Pulled By: cheng-chang fbshipit-source-id: b5256b2dcee6b2583ba73b6c9684994dbe7cdc1f
45 lines
1.2 KiB
C++
45 lines
1.2 KiB
C++
#include <torch/csrc/jit/tensorexpr/cpp_codegen.h>
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
namespace tensorexpr {
|
|
|
|
void CppPrinter::visit(const Allocate* alloc) {
|
|
constexpr size_t kAllocOnStackThresholdSize = 512;
|
|
|
|
size_t size = 1;
|
|
for (auto dim : alloc->dims()) {
|
|
const IntImm* v = dynamic_cast<const IntImm*>(dim);
|
|
if (v) {
|
|
size *= v->value();
|
|
} else {
|
|
throw std::runtime_error("Only IntImm dimensions are supported for now");
|
|
}
|
|
}
|
|
|
|
emitIndent();
|
|
if (size <= kAllocOnStackThresholdSize) {
|
|
os() << alloc->dtype().ToCppString() << " " << (*alloc->buffer_var()) << "["
|
|
<< size << "];" << std::endl;
|
|
} else {
|
|
size *= alloc->dtype().byte_size();
|
|
os() << alloc->dtype().ToCppString() << "* " << (*alloc->buffer_var())
|
|
<< " = static_cast<" << alloc->dtype().ToCppString() << "*>(malloc("
|
|
<< size << "));" << std::endl;
|
|
allocated_on_heap_.insert(alloc->buffer_var());
|
|
}
|
|
}
|
|
|
|
void CppPrinter::visit(const Free* free) {
|
|
const Var* var = free->buffer_var();
|
|
if (allocated_on_heap_.count(var)) {
|
|
emitIndent();
|
|
os() << "free(" << name_manager()->get_unique_name(var) << ");"
|
|
<< std::endl;
|
|
}
|
|
}
|
|
|
|
} // namespace tensorexpr
|
|
} // namespace jit
|
|
} // namespace torch
|