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
|
#pragma once
#include <ATen/ATen.h>
#include <ATen/core/functional.h> // fmap
#include <c10/util/hash.h>
#include <torch/csrc/Export.h>
#include <torch/csrc/jit/codegen/fuser/tensor_desc.h>
#include <cstdint>
#include <vector>
namespace torch {
namespace jit {
namespace fuser {
// Describes the (runtime) arguments to a kernel.
// ArgSpecs are also used as keys to lookup instantiated kernels, so
// they are hashable.
// Note: the device to run on is included in the arg spec because kernels
// are compiled per-device.
struct TORCH_API ArgSpec {
// NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
ArgSpec(at::TensorList inputs, const int _device)
: descs_{c10::fmap<TensorDesc>(inputs)},
hash_code_{c10::get_hash(_device, inputs.size(), descs_)},
device_{_device} {}
// (Common) hash function
static size_t hash(const ArgSpec& spec) {
return spec.hash_code_;
}
// Comparators
bool operator==(const ArgSpec& other) const {
return (descs_ == other.descs_ && device_ == other.device_);
}
bool operator!=(const ArgSpec& spec) const {
return !(*this == spec);
}
// Getters
size_t hashCode() const {
return hash_code_;
}
const std::vector<TensorDesc>& descs() const {
return descs_;
}
int device() const {
return device_;
}
private:
std::vector<TensorDesc> descs_;
size_t hash_code_;
int device_;
};
} // namespace fuser
} // namespace jit
} // namespace torch
|
