Skip to content

Commit 109bc10

Browse files
Cheng Changfacebook-github-bot
Cheng Chang
authored and
facebook-github-bot
committed
[NNC] Generate C++ code for Allocate and Free (#51070)
Summary: This is the initial skeleton for C++ codegen, it includes generations for Allocate and Free. Pull Request resolved: #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
1 parent 642afcb commit 109bc10

5 files changed

Lines changed: 129 additions & 0 deletions

File tree

‎test/cpp/tensorexpr/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ set(TENSOREXPR_TEST_SRCS
44
${TENSOREXPR_TEST_ROOT}/test_aten.cpp
55
${TENSOREXPR_TEST_ROOT}/test_boundsinference.cpp
66
${TENSOREXPR_TEST_ROOT}/test_conv.cpp
7+
${TENSOREXPR_TEST_ROOT}/test_cpp_codegen.cpp
78
${TENSOREXPR_TEST_ROOT}/test_expr.cpp
89
${TENSOREXPR_TEST_ROOT}/test_ir_printer.cpp
910
${TENSOREXPR_TEST_ROOT}/test_kernel.cpp

‎test/cpp/tensorexpr/test_cpp_codegen.cpp‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
#include <gtest/gtest.h>
2+
3+
#include <test/cpp/tensorexpr/test_base.h>
4+
5+
#include <torch/csrc/jit/testing/file_check.h>
6+
#include "torch/csrc/jit/tensorexpr/cpp_codegen.h"
7+
#include "torch/csrc/jit/tensorexpr/mem_arena.h"
8+
#include "torch/csrc/jit/tensorexpr/stmt.h"
9+
10+
namespace torch {
11+
namespace jit {
12+
13+
using namespace torch::jit::tensorexpr;
14+
15+
TEST(CppPrinter, AllocateOnStackThenFree) {
16+
constexpr int dim0 = 2, dim1 = 3;
17+
KernelScope kernel_scope;
18+
VarHandle var("x", kHandle);
19+
Allocate* alloc = Allocate::make(var, kInt, {dim0, dim1});
20+
Free* free = Free::make(var);
21+
Block* block = Block::make({alloc, free});
22+
23+
std::stringstream ss;
24+
CppPrinter printer(&ss);
25+
printer.visit(block);
26+
const std::string expected = R"(
27+
# CHECK: {
28+
# CHECK: int x[6];
29+
# CHECK: }
30+
)";
31+
torch::jit::testing::FileCheck().run(expected, ss.str());
32+
}
33+
34+
TEST(CppPrinter, AllocateOnHeapThenFree) {
35+
constexpr int dim0 = 20, dim1 = 50, dim2 = 3;
36+
KernelScope kernel_scope;
37+
VarHandle var("y", kHandle);
38+
Allocate* alloc = Allocate::make(var, kLong, {dim0, dim1, dim2});
39+
Free* free = Free::make(var);
40+
Block* block = Block::make({alloc, free});
41+
42+
std::stringstream ss;
43+
CppPrinter printer(&ss);
44+
printer.visit(block);
45+
// size(long) = 8;
46+
// dim0 * dim1 * dim2 * size(long) = 24000.
47+
const std::string expected = R"(
48+
# CHECK: {
49+
# CHECK: int64_t* y = static_cast<int64_t*>(malloc(24000));
50+
# CHECK: free(y);
51+
# CHECK: }
52+
)";
53+
torch::jit::testing::FileCheck().run(expected, ss.str());
54+
}
55+
56+
} // namespace jit
57+
} // namespace torch

‎tools/build_variables.bzl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -241,6 +241,7 @@ core_sources_full_mobile = [
241241
"torch/csrc/jit/tensorexpr/bounds_overlap.cpp",
242242
"torch/csrc/jit/tensorexpr/mem_dependency_checker.cpp",
243243
"torch/csrc/jit/tensorexpr/codegen.cpp",
244+
"torch/csrc/jit/tensorexpr/cpp_codegen.cpp",
244245
"torch/csrc/jit/tensorexpr/eval.cpp",
245246
"torch/csrc/jit/tensorexpr/expr.cpp",
246247
"torch/csrc/jit/tensorexpr/hash_provider.cpp",

‎torch/csrc/jit/tensorexpr/cpp_codegen.cpp‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
#include <torch/csrc/jit/tensorexpr/cpp_codegen.h>
2+
3+
namespace torch {
4+
namespace jit {
5+
namespace tensorexpr {
6+
7+
void CppPrinter::visit(const Allocate* alloc) {
8+
constexpr size_t kAllocOnStackThresholdSize = 512;
9+
10+
size_t size = 1;
11+
for (auto dim : alloc->dims()) {
12+
const IntImm* v = dynamic_cast<const IntImm*>(dim);
13+
if (v) {
14+
size *= v->value();
15+
} else {
16+
throw std::runtime_error("Only IntImm dimensions are supported for now");
17+
}
18+
}
19+
20+
emitIndent();
21+
if (size <= kAllocOnStackThresholdSize) {
22+
os() << alloc->dtype().ToCppString() << " " << (*alloc->buffer_var()) << "["
23+
<< size << "];" << std::endl;
24+
} else {
25+
size *= alloc->dtype().byte_size();
26+
os() << alloc->dtype().ToCppString() << "* " << (*alloc->buffer_var())
27+
<< " = static_cast<" << alloc->dtype().ToCppString() << "*>(malloc("
28+
<< size << "));" << std::endl;
29+
allocated_on_heap_.insert(alloc->buffer_var());
30+
}
31+
}
32+
33+
void CppPrinter::visit(const Free* free) {
34+
const Var* var = free->buffer_var();
35+
if (allocated_on_heap_.count(var)) {
36+
emitIndent();
37+
os() << "free(" << name_manager()->get_unique_name(var) << ");"
38+
<< std::endl;
39+
}
40+
}
41+
42+
} // namespace tensorexpr
43+
} // namespace jit
44+
} // namespace torch

‎torch/csrc/jit/tensorexpr/cpp_codegen.h‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
#pragma once
2+
3+
#include <torch/csrc/jit/tensorexpr/ir_printer.h>
4+
5+
#include <unordered_set>
6+
7+
namespace torch {
8+
namespace jit {
9+
namespace tensorexpr {
10+
11+
// Generates C++ code from the IR.
12+
class TORCH_API CppPrinter : public IRPrinter {
13+
public:
14+
explicit CppPrinter(std::ostream* os) : IRPrinter(*os) {}
15+
16+
using IRPrinter::visit;
17+
void visit(const Allocate*) override;
18+
void visit(const Free*) override;
19+
20+
private:
21+
std::unordered_set<const Var*> allocated_on_heap_;
22+
};
23+
24+
} // namespace tensorexpr
25+
} // namespace jit
26+
} // namespace torch

0 commit comments

Comments
 (0)