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
|
#pragma once
#include <c10/macros/Export.h>
#include <torch/csrc/jit/codegen/cuda/dynamic_type.h>
#include <torch/csrc/jit/codegen/cuda/ir_interface_nodes.h>
#include <torch/csrc/jit/codegen/cuda/iter_visitor.h>
#include <c10/util/Optional.h>
#include <unordered_map>
namespace torch {
namespace jit {
namespace fuser {
namespace cuda {
class FusionPrecomputedValues;
//! Calculate Fusion IR expressions
class TORCH_CUDA_CU_API ExpressionEvaluator : private OptOutDispatch {
public:
// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
explicit ExpressionEvaluator(Fusion* fusion) : fusion_(fusion) {}
//! Returns the associated fusion object
Fusion* fusion() const {
return fusion_;
}
//! Bind a concrete value to an IR variable
void bind(Val* value, const IntOrDouble& concrete_value);
//! Try to evaluate a Fusion IR value
c10::optional<IntOrDouble> evaluate(Val* value);
//! Debugging helper, prints all the currently known values
void print() const;
void bindPrecomputedValues(FusionPrecomputedValues* precomputed_values) {
evaluator_precomputed_values_ = precomputed_values;
}
auto precomputedValues() {
return evaluator_precomputed_values_;
}
private:
c10::optional<IntOrDouble> getValue(Val* value);
void handle(UnaryOp*) final;
void handle(BinaryOp*) final;
private:
std::unordered_map<const Val*, IntOrDouble> known_values_;
Fusion* fusion_ = nullptr;
FusionPrecomputedValues* evaluator_precomputed_values_ = nullptr;
};
} // namespace cuda
} // namespace fuser
} // namespace jit
} // namespace torch
|
