forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test_ir_printer.cpp
49 lines (40 loc) · 1.14 KB
/
test_ir_printer.cpp
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
#include <stdexcept>
#include "test/cpp/tensorexpr/test_base.h"
#include "torch/csrc/jit/tensorexpr/expr.h"
#include "torch/csrc/jit/tensorexpr/ir.h"
#include "torch/csrc/jit/tensorexpr/ir_printer.h"
#include <sstream>
namespace torch {
namespace jit {
using namespace torch::jit::tensorexpr;
void testIRPrinterBasicValueTest() {
KernelScope kernel_scope;
ExprHandle a = IntImm::make(2), b = IntImm::make(3);
ExprHandle c = Add::make(a, b);
std::stringstream ss;
ss << c;
ASSERT_EQ(ss.str(), "2 + 3");
}
void testIRPrinterBasicValueTest02() {
KernelScope kernel_scope;
ExprHandle a(2.0f);
ExprHandle b(3.0f);
ExprHandle c(4.0f);
ExprHandle d(5.0f);
ExprHandle f = (a + b) - (c + d);
std::stringstream ss;
ss << f;
ASSERT_EQ(ss.str(), "(2.f + 3.f) - (4.f + 5.f)");
}
void testIRPrinterCastTest() {
KernelScope kernel_scope;
VarHandle x("x", kHalf);
VarHandle y("y", kFloat);
ExprHandle body = ExprHandle(2.f) +
(Cast::make(kFloat, x) * ExprHandle(3.f) + ExprHandle(4.f) * y);
std::stringstream ss;
ss << body;
ASSERT_EQ(ss.str(), "2.f + (float(x) * 3.f + 4.f * y)");
}
} // namespace jit
} // namespace torch