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
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
|
#pragma once
#include <torch/csrc/WindowsTorchApiMacro.h>
#include <torch/csrc/jit/codegen/cuda/kernel.h>
#include <torch/csrc/jit/codegen/cuda/kernel_ir.h>
#include <memory>
namespace torch {
namespace jit {
namespace fuser {
namespace kir {
// Simple classification helpers
bool isLoweredScalar(const Val* val);
bool isLoweredVal(const Val* val);
//! Kernel IR builder interface
//!
//! The only way to create new Kernel IR nodes is through the
//! kir::IrBuilder interface. An IrBuilder instance is attached to a
//! particular Kernel instance and it provides methods for creating
//! single nodes (kir::IrBuilder::create()) or basic composite expressions
//! (ex. kir::IrBuilder::addExpr()).
//!
//! If the Kernel object is readily available, an IrBuilder can be "wrapped"
//! around it directly:
//!
//! kir::IrBuilder ir_builder(kernel);
//!
//! During lowering, another option is to create an IrBuilder for the
//! kernel that is being created:
//!
//! kir::IrBuilder ir_builder(GpuLower::current()->kernel());
//!
//! Once we have an IR builder instance, creating nodes looks like:
//!
//! auto new_node = ir_builder.create<kir::Int>(1));
//! auto result = ir_builder.mulExpr(lhs, rhs);
//!
class IrBuilder {
public:
explicit IrBuilder(Kernel* kernel) : kernel_(kernel) {}
//! Allocate a new Kernel IR node, forwarding the arguments
//! to the appropriate constructor
template <class T, class... Args>
T* create(Args&&... args) {
// TODO(kir): switch this to Kernel registration
return new T(kir::Passkey(), std::forward<Args>(args)...);
}
// Binary expressions
Val* andExpr(Val* lhs, Val* rhs);
Val* eqExpr(Val* lhs, Val* rhs);
Val* ltExpr(Val* lhs, Val* rhs);
Val* addExpr(Val* lhs, Val* rhs);
Val* subExpr(Val* lhs, Val* rhs);
Val* mulExpr(Val* lhs, Val* rhs);
Val* divExpr(Val* lhs, Val* rhs);
Val* ceilDivExpr(Val* lhs, Val* rhs);
Val* modExpr(Val* lhs, Val* rhs);
private:
Val* newResult(const Val* lhs, const Val* rhs);
Val* newArithmeticExpr(BinaryOpType op_type, Val* lhs, Val* rhs);
Val* newLogicExpr(BinaryOpType op_type, Val* lhs, Val* rhs);
private:
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wunused-private-field"
// Non-owning pointer to the kernel to be modified
Kernel* kernel_ = nullptr;
#pragma clang diagnostic pop
};
} // namespace kir
} // namespace fuser
} // namespace jit
} // namespace torch
|
