1 | #pragma once |
2 | |
3 | #include <c10/macros/Export.h> |
4 | |
5 | #include <ir_all_nodes.h> |
6 | |
7 | namespace torch { |
8 | namespace jit { |
9 | namespace fuser { |
10 | namespace cuda { |
11 | |
12 | struct VectorizedSetInfo { |
13 | //! Producer of a vectorized set |
14 | TensorView* producer_tv = nullptr; |
15 | //! Consumer of a vectorized set |
16 | TensorView* consumer_tv = nullptr; |
17 | //! Number of elements to vectorize |
18 | int word_size = -1; |
19 | //! Vectorized domain |
20 | IterDomain* vectorized_leaf_id = nullptr; |
21 | //! Right-most root dependent domain of the leaf domain |
22 | IterDomain* vectorized_root_id = nullptr; |
23 | //! All of the dependent root domains that are contiguously merged |
24 | std::unordered_set<IterDomain*> contig_root_ids; |
25 | }; |
26 | |
27 | } // namespace cuda |
28 | } // namespace fuser |
29 | } // namespace jit |
30 | } // namespace torch |
31 | |