1#pragma once
2
3#include <c10/macros/Export.h>
4
5#include <ir_all_nodes.h>
6
7namespace torch {
8namespace jit {
9namespace fuser {
10namespace cuda {
11
12struct 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