diff --git a/tensorflow/compiler/xla/service/hlo_module_config.h b/tensorflow/compiler/xla/service/hlo_module_config.h index 61ea8392d94..833d0fe59d0 100644 --- a/tensorflow/compiler/xla/service/hlo_module_config.h +++ b/tensorflow/compiler/xla/service/hlo_module_config.h @@ -204,6 +204,14 @@ class HloModuleConfig { 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: // If you add new members, be sure to update compilation_cache_key. @@ -241,6 +249,9 @@ class HloModuleConfig { FusionConfigCollection fusion_config_collection_ = 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 // in computation c must be fused to all its consumers (true) or not (false). 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 // convolution. 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