From 117c75d3117ac8babe84393dce32dbc2dd2dbe36 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 12 May 2020 13:43:09 -0700
Subject: [PATCH] Add layout config to HloModuleConfig.
PiperOrigin-RevId: 311193903
Change-Id: I9b6680c5a9919804e449ed617bd6bc310800183e
---
.../compiler/xla/service/hlo_module_config.h | 15 +++++++++++++++
1 file changed, 15 insertions(+)
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