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
|
#pragma once
#include <torch/csrc/WindowsTorchApiMacro.h>
#include <torch/csrc/jit/codegen/cuda/dispatch.h>
#include <torch/csrc/jit/codegen/cuda/instrumentation.h>
#include <torch/csrc/jit/codegen/cuda/ir_all_nodes.h>
#include <torch/csrc/jit/codegen/cuda/kernel_ir_builder.h>
#include <vector>
namespace torch {
namespace jit {
namespace fuser {
class TORCH_CUDA_API IndexLowering : public OptInDispatch {
public:
static std::vector<Expr*> getIndexedExprs(
Fusion* fusion,
std::vector<Expr*> incoming_exprs) {
FUSER_PERF_SCOPE("IndexLowering::getIndexedExprs");
FusionGuard fg(fusion);
IndexLowering il;
il.generate(incoming_exprs);
return il.lowered_exprs;
}
private:
IndexLowering();
// Wrap pushBack, if active_scope is null we want it to go
// straight to lower_exprs
void pushBack(Expr*);
// Open the for loop.
void handle(kir::ForLoop*) final;
// Open the for loop.
void handle(kir::IfThenElse*) final;
// Remake operations with TensorIndex
void handle(UnaryOp*) final;
void handle(BinaryOp*) final;
void handle(TernaryOp*) final;
void handle(ReductionOp*) final;
void handle(BroadcastOp*) final;
void handle(kir::Allocate*) final;
void handle(kir::Sync*) final;
void generate(const std::vector<Expr*>& exprs);
Val* lowerOperand(Val* op, Val* out) const;
Val* lowerOutput(Expr* expr) const;
private:
std::vector<Expr*> lowered_exprs;
// This is a slight work around as scope has a couple definitions, we have the
// Scope that's in ForLoop/IfThenElse which is really just a wrapper around
// std::vector<Expr*> and then we have the actual ForLoop/IfThenElse. We want
// to be able to carry both around because when we push back to a scope it
// could be either the body or else body of the IfThenElse. However, we want
// to understand the nesting of IfThenElse/ForLoop nodes.
kir::Scope* active_scope = nullptr;
Expr* active_scope_expr = nullptr;
kir::IrBuilder ir_builder_;
};
} // namespace fuser
} // namespace jit
} // namespace torch
|
