pytorch/test/cpp/tensorexpr/test_cpp_codegen.cpp at master · DavidBert/pytorch · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
#include <gtest/gtest.h>

#include <test/cpp/tensorexpr/test_base.h>

#include <torch/csrc/jit/tensorexpr/cpp_codegen.h>
#include <torch/csrc/jit/tensorexpr/mem_arena.h>
#include <torch/csrc/jit/tensorexpr/stmt.h>
#include <torch/csrc/jit/testing/file_check.h>

namespace torch {
namespace jit {

using namespace torch::jit::tensorexpr;

TEST(CppPrinter, AllocateOnStackThenFree) {
  KernelScope kernel_scope;
  std::vector<ExprPtr> dims = {alloc<IntImm>(2), alloc<IntImm>(3)};
  BufPtr buf = alloc<Buf>("x", dims, kInt);
  AllocatePtr alloc_ = alloc<Allocate>(buf);
  FreePtr free_ = alloc<Free>(buf);
  BlockPtr block = Block::make({alloc_, free_});

  std::stringstream ss;
  CppPrinter printer(&ss);
  printer.visit(block);
  const std::string expected = R"(
    # CHECK: {
    # CHECK:   int x[6];
    # CHECK: }
  )";
  torch::jit::testing::FileCheck().run(expected, ss.str());
}

TEST(CppPrinter, AllocateOnHeapThenFree) {
  KernelScope kernel_scope;
  std::vector<ExprPtr> dims = {
      alloc<IntImm>(20), alloc<IntImm>(50), alloc<IntImm>(3)};
  BufPtr buf = alloc<Buf>("y", dims, kLong);
  AllocatePtr alloc_ = alloc<Allocate>(buf);
  FreePtr free_ = alloc<Free>(buf);
  BlockPtr block = Block::make({alloc_, free_});

  std::stringstream ss;
  CppPrinter printer(&ss);
  printer.visit(block);
  // size(long) = 8;
  // dim0 * dim1 * dim2 * size(long) = 24000.
  const std::string expected = R"(
    # CHECK: {
    # CHECK:   int64_t* y = static_cast<int64_t*>(malloc(24000));
    # CHECK:   free(y);
    # CHECK: }
  )";
  torch::jit::testing::FileCheck().run(expected, ss.str());
}

} // namespace jit
} // namespace torch