experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add layout config to HloModuleConfig.

Browse Source 
PiperOrigin-RevId: 311193903
Change-Id: I9b6680c5a9919804e449ed617bd6bc310800183e
This commit is contained in:
A. Unique TensorFlower 2020-05-12 13:43:09 -07:00 committed by TensorFlower Gardener
parent e036f1bd8f
commit 117c75d311
1 changed files with 15 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
tensorflow/compiler/xla/service/hlo_module_config.h
View File

@ -204,6 +204,14 @@ class HloModuleConfig {
std::vector<std::vector<int64>>* mutable_dot_config() { return &dot_config_; } std::vector<std::vector<int64>>* mutable_dot_config() { return &dot_config_; }
absl::Span<const std::vector<std::vector<int64>>> layout_config() const {
return layout_config_;
}
std::vector<std::vector<std::vector<int64>>>* mutable_layout_config() {
return &layout_config_;
}
private: private:
// If you add new members, be sure to update compilation_cache_key. // If you add new members, be sure to update compilation_cache_key.
@ -241,6 +249,9 @@ class HloModuleConfig {
FusionConfigCollection fusion_config_collection_ = FusionConfigCollection fusion_config_collection_ =
FusionConfigCollection::kOff; FusionConfigCollection::kOff;
// TODO(b/155665133): Consolidate fusion, dot, and layout config into a proto
// similar to backend config.
// Custom fusion configuration, where fusion_config_[c][v] control if node v // Custom fusion configuration, where fusion_config_[c][v] control if node v
// in computation c must be fused to all its consumers (true) or not (false). // in computation c must be fused to all its consumers (true) or not (false).
std::vector<std::vector<bool>> fusion_config_; std::vector<std::vector<bool>> fusion_config_;
@ -249,6 +260,10 @@ class HloModuleConfig {
// how to convert dot operation v (sorted topologically and by computation) to // how to convert dot operation v (sorted topologically and by computation) to
// convolution. // convolution.
std::vector<std::vector<int64>> dot_config_; std::vector<std::vector<int64>> dot_config_;
// Layout configuration, where layout_config_[v][i] controls the layout
// decision i of operation v.
std::vector<std::vector<std::vector<int64>>> layout_config_;
}; };
} // namespace xla } // namespace xla