Automated Code Change

PiperOrigin-RevId: 643541642
tensorflow · Jun 18, 2024 · 518a4f7 · 518a4f7
1 parent c395b48
commit 518a4f7
Show file tree

Hide file tree

Showing 11 changed files with 263 additions and 67 deletions.
diff --git a/third_party/xla/third_party/tsl/tsl/platform/env.cc b/third_party/xla/third_party/tsl/tsl/platform/env.cc
@@ -328,7 +328,8 @@ absl::Status Env::HasAtomicMove(const string& path, bool* has_atomic_move) {
   return fs->HasAtomicMove(path, has_atomic_move);
 }
 
-Status Env::CanCreateTempFile(const string& fname, bool* can_create_temp_file) {
+absl::Status Env::CanCreateTempFile(const string& fname,
+                                    bool* can_create_temp_file) {
   FileSystem* fs;
   TF_RETURN_IF_ERROR(GetFileSystemForFile(fname, &fs));
   return fs->CanCreateTempFile(fname, can_create_temp_file);

diff --git a/third_party/xla/third_party/tsl/tsl/platform/env.h b/third_party/xla/third_party/tsl/tsl/platform/env.h
@@ -352,8 +352,8 @@ class Env {
   /// If this returns false, TensorFlow will write directly to output files
   /// instead of creating a temporary file and swapping it in. This may mean
   /// that incomplete writes are visible to consumers.
-  Status CanCreateTempFile(const std::string& fname,
-                           bool* can_create_temp_file);
+  absl::Status CanCreateTempFile(const std::string& fname,
+                                 bool* can_create_temp_file);
 
   /// Stores the size of `fname` in `*file_size`.
   absl::Status GetFileSize(const std::string& fname, uint64* file_size);

diff --git a/third_party/xla/third_party/tsl/tsl/platform/file_system.cc b/third_party/xla/third_party/tsl/tsl/platform/file_system.cc
@@ -95,10 +95,10 @@ absl::Status FileSystem::HasAtomicMove(const string& path,
   return absl::OkStatus();
 }
 
-Status FileSystem::CanCreateTempFile(const std::string& fname,
-                                     bool* can_create_temp_file) {
+absl::Status FileSystem::CanCreateTempFile(const std::string& fname,
+                                           bool* can_create_temp_file) {
   *can_create_temp_file = true;
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 void FileSystem::FlushCaches(TransactionToken* token) {}

diff --git a/third_party/xla/third_party/tsl/tsl/platform/file_system.h b/third_party/xla/third_party/tsl/tsl/platform/file_system.h
@@ -392,8 +392,8 @@ class FileSystem {
   /// to determine if there needs to be a temp location to safely write objects.
   /// If the file system cannot create a temp file, it's possibile that
   /// uncomplete result may appear in the given file.
-  virtual Status CanCreateTempFile(const std::string& fname,
-                                   bool* can_create_temp_file);
+  virtual absl::Status CanCreateTempFile(const std::string& fname,
+                                         bool* can_create_temp_file);
 
   /// \brief Flushes any cached filesystem objects from memory.
   virtual void FlushCaches() { FlushCaches(nullptr); }

diff --git a/third_party/xla/third_party/tsl/tsl/platform/retrying_file_system.h b/third_party/xla/third_party/tsl/tsl/platform/retrying_file_system.h
@@ -151,8 +151,8 @@ class RetryingFileSystem : public FileSystem {
     return base_file_system_->HasAtomicMove(path, has_atomic_move);
   }
 
-  Status CanCreateTempFile(const std::string& fname,
-                           bool* can_create_temp_file) override {
+  absl::Status CanCreateTempFile(const std::string& fname,
+                                 bool* can_create_temp_file) override {
     // this method does not need to be retried
     return base_file_system_->CanCreateTempFile(fname, can_create_temp_file);
   }

diff --git a/third_party/xla/xla/service/gpu/model/BUILD b/third_party/xla/xla/service/gpu/model/BUILD
@@ -691,6 +691,7 @@ xla_cc_test(
         "//xla/service/gpu:hlo_traversal",
         "//xla/tests:hlo_test_base",
         "//xla/tests:verified_hlo_module",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest_main",

diff --git a/third_party/xla/xla/service/gpu/model/symbolic_tile.cc b/third_party/xla/xla/service/gpu/model/symbolic_tile.cc
@@ -106,43 +106,6 @@ AffineMap SubstituteAllIndicesAndRangeVarSymbolsWithSameValue(
   return simplifyAffineMap(affine_map.replace(indices, num_dims, num_symbols));
 }
 
-// Merges `maybe_first_map` and `second_map` if
-//  (1) `maybe_first_map` is present, and
-//  (2) `second_map` and `*maybe_first_map` have distinct sets of keys.
-// Otherwise, returns `std::nullopt`.
-//
-//
-// The behaviour of this function is in spirit equivalent to using C++23's
-// `std::optional<T>::and_then` to merge a collection of `ConstraintMap`s.
-//
-// We pass `maybe_first_map` by value here in order to exploit move semantics
-// to avoid copies when possible.
-//
-// TODO(bchetioui): allow merging constraints in more edge cases, e.g. if one
-// of the intervals is contained within the other.
-std::optional<ConstraintMap> MergeConstraintMapIfPresentAndCompatible(
-    std::optional<ConstraintMap> maybe_first_map,
-    const ConstraintMap& second_map) {
-  if (!maybe_first_map.has_value()) {
-    return std::nullopt;
-  }
-
-  ConstraintMap& first_map = *maybe_first_map;
-
-  for (const auto& [expr, interval] : second_map) {
-    if (first_map.contains(expr)) {
-      AffineMapPrinter printer;
-      VLOG(1) << "Got two different constraints for expression "
-              << printer.ToString(expr);
-      return std::nullopt;
-    }
-
-    first_map.insert({expr, interval});
-  }
-
-  return first_map;
-}
-
 struct SizeAndStrideExpression {
   AffineExpr size;
   AffineExpr stride;
@@ -620,6 +583,29 @@ AffineExpr SimplifyAffineExpr(const AffineExpr& expr,
 
 }  // anonymous namespace
 
+std::optional<ConstraintMap> MergeConstraintMapIfPresentAndCompatible(
+    std::optional<ConstraintMap> maybe_first_map,
+    const ConstraintMap& second_map) {
+  if (!maybe_first_map.has_value()) {
+    return std::nullopt;
+  }
+
+  ConstraintMap& first_map = *maybe_first_map;
+
+  for (const auto& [expr, interval] : second_map) {
+    if (first_map.contains(expr)) {
+      AffineMapPrinter printer;
+      VLOG(1) << "Got two different constraints for expression "
+              << printer.ToString(expr);
+      return std::nullopt;
+    }
+
+    first_map.insert({expr, interval});
+  }
+
+  return first_map;
+}
+
 /*static*/ std::optional<SymbolicTile> SymbolicTile::FromIndexingMap(
     const IndexingMap& indexing_map) {
   VLOG(1) << "SymbolicTile::FromIndexingMap: " << indexing_map.ToString();

diff --git a/third_party/xla/xla/service/gpu/model/symbolic_tile.h b/third_party/xla/xla/service/gpu/model/symbolic_tile.h
@@ -225,6 +225,25 @@ class SymbolicTile {
         is_satisfiable_(is_satisfiable) {}
 };
 
+// Merges `maybe_first_map` and `second_map` if
+//  (1) `maybe_first_map` is present, and
+//  (2) `second_map` and `*maybe_first_map` have distinct sets of keys.
+// Otherwise, returns `std::nullopt`.
+//
+//
+// The behaviour of this function is in spirit equivalent to using C++23's
+// `std::optional<T>::and_then` to merge a collection of `ConstraintMap`s.
+//
+// We pass `maybe_first_map` by value here in order to exploit move semantics
+// to avoid copies when possible.
+//
+// TODO(bchetioui): allow merging constraints in more edge cases, e.g. if one
+// of the intervals is contained within the other.
+std::optional<SymbolicTile::ConstraintMap>
+MergeConstraintMapIfPresentAndCompatible(
+    std::optional<SymbolicTile::ConstraintMap> maybe_first_map,
+    const SymbolicTile::ConstraintMap& second_map);
+
 }  // namespace gpu
 }  // namespace xla
 

diff --git a/third_party/xla/xla/service/gpu/model/symbolic_tile_analysis.cc b/third_party/xla/xla/service/gpu/model/symbolic_tile_analysis.cc
@@ -35,6 +35,8 @@ limitations under the License.
 #include "absl/strings/str_join.h"
 #include "absl/types/span.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
 #include "mlir/IR/AffineExpr.h"  // from @llvm-project
 #include "mlir/IR/AffineMap.h"  // from @llvm-project
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
@@ -61,6 +63,7 @@ namespace {
 
 using ::mlir::AffineExpr;
 using ::mlir::MLIRContext;
+using ConstraintMap = SymbolicTile::ConstraintMap;
 
 // Computes indexing map from program id into the tile offset for the given
 // shape and tile sizes.
@@ -162,6 +165,8 @@ absl::StatusOr<IndexingMap> ComputeBlockIdToTileOffsetIndexing(
       const HloInstructionAdaptor&, IndexingMap)>
       get_tiled_hlo_instruction;
 
+  std::optional<ConstraintMap> constraints = ConstraintMap();
+
   // Create a new tiled hlo instruction or return existing instruction from
   // cache for the given hlo and indexing map.
   get_tiled_hlo_instruction =
@@ -180,8 +185,6 @@ absl::StatusOr<IndexingMap> ComputeBlockIdToTileOffsetIndexing(
     // line. This is not an inherent limitation of the approach, but simply
     // issues to be resolved in the current implementation.
     if (hlo->opcode() == HloOpcode::kDot ||
-        hlo->opcode() == HloOpcode::kReshape ||
-        hlo->opcode() == HloOpcode::kBitcast ||
         hlo->opcode() == HloOpcode::kConcatenate) {
       return FusionDecision{} << "Bailing out on " << hlo->ToString();
     }
@@ -199,6 +202,22 @@ absl::StatusOr<IndexingMap> ComputeBlockIdToTileOffsetIndexing(
                               << hlo->ToString();
     }
 
+    if (!symbolic_tile->is_satisfiable()) {
+      return FusionDecision{} << "Symbolic tile " << symbolic_tile->ToString()
+                              << " is not satisfiable for "
+                              << indexing_map.ToString() << " for HLO "
+                              << hlo->ToString();
+    }
+
+    constraints = MergeConstraintMapIfPresentAndCompatible(
+        std::move(constraints), symbolic_tile->constraints());
+
+    if (!constraints.has_value()) {
+      return FusionDecision{} << "Failed to merge constraints of "
+                              << hlo->ToString() << " in pre-existing "
+                              << "constraint map";
+    }
+
     tiled_hlo_instructions.push_back(
         std::make_unique<SymbolicTiledHloInstruction>(
             hlo, std::move(indexing_map), std::move(*symbolic_tile)));
@@ -258,12 +277,53 @@ absl::StatusOr<IndexingMap> ComputeBlockIdToTileOffsetIndexing(
     return topological_order.at(i1.get()) < topological_order.at(i2.get());
   });
 
-  return SymbolicTileAnalysis(std::move(tiled_hlo_instructions), ctx);
+  return SymbolicTileAnalysis(std::move(tiled_hlo_instructions),
+                              std::move(*constraints), ctx);
+}
+
+absl::StatusOr<bool> SymbolicTileAnalysis::ParametersSatisfyConstraints(
+    const std::vector<int64_t>& tile_parameters) const {
+  // Populate parameter map.
+  llvm::SmallVector<AffineExpr> parameters = llvm::to_vector(
+      llvm::map_range(tile_parameters, [this](const int64_t v) -> AffineExpr {
+        return mlir::getAffineConstantExpr(v, context_);
+      }));
+
+  for (auto [constrained_expr, interval] : constraints_) {
+    AffineExpr constrained_expr_value =
+        constrained_expr.replaceSymbols(parameters);
+    if (constrained_expr_value.getKind() != mlir::AffineExprKind::Constant) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Failed to reduce ", AffineMapPrinter().ToString(constrained_expr),
+          " to a constant with tile parameters ",
+          absl::StrJoin(tile_parameters, ", ")));
+    }
+
+    int64_t constrained_value =
+        llvm::cast<mlir::AffineConstantExpr>(constrained_expr_value).getValue();
+
+    if (constrained_value < interval.lower ||
+        constrained_value > interval.upper) {
+      return false;
+    }
+  }
+  return true;
 }
 
 absl::StatusOr<TiledHloComputation>
 SymbolicTileAnalysis::ComputeTiledHloInstructions(
-    const std::vector<int64_t>& tile_parameters) const {
+    const std::vector<int64_t>& tile_parameters,
+    bool constraints_are_known_satisfied) const {
+  if (!constraints_are_known_satisfied) {
+    TF_ASSIGN_OR_RETURN(bool constraints_are_satisfied,
+                        ParametersSatisfyConstraints(tile_parameters));
+    if (!constraints_are_satisfied) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Tile parameters ", absl::StrJoin(tile_parameters, ", "),
+          " do not satisfy the SymbolicTileAnalysis's constraints."));
+    }
+  }
+
   IndexingMap block_id_to_root_tile_offset = ComputeBlockIdToOutputTileIndexing(
       GetRoot()->hlo()->shape().dimensions(), tile_parameters, context_);
 

diff --git a/third_party/xla/xla/service/gpu/model/symbolic_tile_analysis.h b/third_party/xla/xla/service/gpu/model/symbolic_tile_analysis.h
@@ -29,6 +29,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/service/gpu/hlo_traversal.h"
 #include "xla/service/gpu/model/affine_map_printer.h"
+#include "xla/service/gpu/model/symbolic_tile.h"
 #include "xla/service/gpu/model/symbolic_tiled_hlo_instruction.h"
 #include "xla/service/gpu/model/tiled_hlo_computation.h"
 #include "xla/service/instruction_fusion.h"
@@ -56,8 +57,13 @@ class SymbolicTileAnalysis {
       const HloFusionAdaptor& fusion, mlir::MLIRContext* ctx);
 
   // Returns a graph of HLO instructions tiled with the given tile parameters.
+  // The provided tile parameters must satisfy the analysis's constraints.
+  // By default, `ComputetiledHloInstructions` performs a check that the
+  // constraints are satisfied by the chosen tiled parameters. Setting
+  // `constraints_are_known_satisfied` to true bypasses this check.
   absl::StatusOr<TiledHloComputation> ComputeTiledHloInstructions(
-      const std::vector<int64_t>& tile_parameters) const;
+      const std::vector<int64_t>& tile_parameters,
+      bool constraints_are_known_satisfied = false) const;
 
   // Returns the tiled root instruction.
   const SymbolicTiledHloInstruction* GetRoot() const {
@@ -70,6 +76,23 @@ class SymbolicTileAnalysis {
     return symbolic_tiled_hlo_instructions_;
   }
 
+  // Returns the constraints for the parameters of the symbolic tiled HLO
+  // computation. This is the union of the constraints of all the symbolic tiles
+  // encountered throughout the computation.
+  const SymbolicTile::ConstraintMap& GetConstraints() const {
+    return constraints_;
+  }
+
+  // Returns true if a list of tile parameters satisfies the symbolic tile
+  // analysis's constraints.
+  //
+  // Returns false if the constraints are not satisfied but can be evaluated
+  // correctly. Returns an error if the constraints cannot be evaluated
+  // correctly. This is typically the case if too few tile parameters are
+  // provided to fully reduce the constraint expressions to constants.
+  absl::StatusOr<bool> ParametersSatisfyConstraints(
+      const std::vector<int64_t>& tile_parameters) const;
+
   // Return the underlying MLIRContext.
   mlir::MLIRContext* GetMLIRContext() const { return context_; };
 
@@ -81,15 +104,20 @@ class SymbolicTileAnalysis {
  private:
   SymbolicTileAnalysis(std::vector<std::unique_ptr<SymbolicTiledHloInstruction>>
                            symbolic_tiled_hlo_instructions,
+                       SymbolicTile::ConstraintMap constraints,
                        mlir::MLIRContext* context)
       : symbolic_tiled_hlo_instructions_(
             std::move(symbolic_tiled_hlo_instructions)),
+        constraints_(std::move(constraints)),
         context_(context) {}
 
   // The tiled HLO instructions in def-before-use order.
   std::vector<std::unique_ptr<SymbolicTiledHloInstruction>>
       symbolic_tiled_hlo_instructions_;
 
+  // See the documentation of GetConstraints().
+  SymbolicTile::ConstraintMap constraints_;
+
   mlir::MLIRContext* context_;
 };