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
|
#pragma once
#include <torch/csrc/jit/ir/ir.h>
#include <torch/csrc/jit/passes/symbolic_shape_analysis.h>
namespace torch {
namespace jit {
struct TORCH_API CanonicalizedSymbolicShape {
// TODO: Consider in the future if it is reasonable to
// merge code with SymbolicShape or VaryingShape while keeping
// the two not implicitly convertable (and cause bugs).
CanonicalizedSymbolicShape(
const c10::SymbolicShape& orig_shape,
std::unordered_map<int64_t, int64_t>& ss_map) {
init(orig_shape, ss_map);
}
CanonicalizedSymbolicShape(c10::SymbolicShape& orig_shape) {
std::unordered_map<int64_t, int64_t> new_ssmap;
init(orig_shape, new_ssmap);
}
size_t hash() const;
c10::SymbolicShape toSymbolicShape(
std::unordered_map<int64_t, int64_t>& inverse_ss_map) const;
TORCH_API friend bool operator==(
const CanonicalizedSymbolicShape& a,
const CanonicalizedSymbolicShape& b);
private:
c10::optional<std::vector<int64_t>> values_;
void init(
const c10::SymbolicShape& orig_shape,
std::unordered_map<int64_t, int64_t>& ss_map);
};
// SHAPE CACHE API
TORCH_API c10::optional<std::vector<at::SymbolicShape>>
get_cached_shape_function(
const FunctionSchema* schema,
const std::vector<SSAInput>& arg_vec);
TORCH_API void cache_shape_function(
const FunctionSchema* schema,
const std::vector<SSAInput>& arg_vec,
const std::vector<at::SymbolicShape>& ret_vec);
// For use in test code
TORCH_API void clear_shape_cache();
TORCH_API size_t get_shape_cache_size();
} // namespace jit
} // namespace torch
|
