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
|
#pragma once
#include <c10/core/ScalarType.h>
#include <torch/csrc/Export.h>
#include <vector>
#include <torch/csrc/jit/tensorexpr/ir_mutator.h>
namespace torch {
namespace jit {
namespace tensorexpr {
class TORCH_API IRCloner : public IRMutator {
public:
~IRCloner() override = default;
ExprPtr mutate(AddPtr v) override;
ExprPtr mutate(SubPtr v) override;
ExprPtr mutate(MulPtr v) override;
ExprPtr mutate(DivPtr v) override;
ExprPtr mutate(ModPtr v) override;
ExprPtr mutate(MaxPtr v) override;
ExprPtr mutate(MinPtr v) override;
ExprPtr mutate(AndPtr v) override;
ExprPtr mutate(OrPtr v) override;
ExprPtr mutate(XorPtr v) override;
ExprPtr mutate(LshiftPtr v) override;
ExprPtr mutate(RshiftPtr v) override;
ExprPtr mutate(CompareSelectPtr v) override;
#define IMM_MUTATE_DECLARE(Type, Name) ExprPtr mutate(Name##ImmPtr v) override;
AT_FORALL_SCALAR_TYPES_AND3(Bool, Half, BFloat16, IMM_MUTATE_DECLARE);
#undef IMM_MUTATE_DECLARE
ExprPtr mutate(CastPtr v) override;
ExprPtr mutate(BitCastPtr v) override;
ExprPtr mutate(VarPtr v) override;
ExprPtr mutate(BufPtr v) override;
ExprPtr mutate(RampPtr v) override;
ExprPtr mutate(LoadPtr v) override;
ExprPtr mutate(BroadcastPtr v) override;
ExprPtr mutate(IfThenElsePtr v) override;
ExprPtr mutate(IntrinsicsPtr v) override;
ExprPtr mutate(TermPtr v) override;
ExprPtr mutate(PolynomialPtr v) override;
ExprPtr mutate(RoundOffPtr v) override;
ExprPtr mutate(MaxTermPtr v) override;
ExprPtr mutate(MinTermPtr v) override;
ExprPtr mutate(ReduceOpPtr v) override;
StmtPtr mutate(ForPtr v) override;
StmtPtr mutate(BlockPtr v) override;
StmtPtr mutate(StorePtr v) override;
StmtPtr mutate(AtomicAddPtr v) override;
StmtPtr mutate(SyncThreadsPtr v) override;
StmtPtr mutate(ExternalCallPtr v) override;
StmtPtr mutate(ExternalCallWithAllocPtr v) override;
StmtPtr mutate(AllocatePtr v) override;
StmtPtr mutate(FreePtr v) override;
StmtPtr mutate(LetPtr v) override;
StmtPtr mutate(CondPtr v) override;
};
} // namespace tensorexpr
} // namespace jit
} // namespace torch
|
