[XLA] Implement a TiledHloSchedule that transposes the iteration order over the non-contracting dimensions of a dot.

bchetioui · Google-ML-Automation · commit 89fc91309957 · 2025-10-16T07:08:55.000-07:00
A concrete use case when such a schedule is useful is when we have a matrix
multiplication such that a chunk of shape `(block_m, k)` of the left-hand
side argument fully fits into L2. The transposed iteration order will step
through the `n` dimension first, allowing to hit L2 cache more often when
loading tiles of the left-hand side.

This schedule is intentionally restricted at the moment in order to unblock
launching the generic Triton emitter for GEMMs.

PiperOrigin-RevId: 820214481
diff --git a/xla/codegen/tiling/BUILD b/xla/codegen/tiling/BUILD
@@ -118,12 +118,18 @@ cc_library(
     srcs = ["tiled_hlo_schedule.cc"],
     hdrs = ["tiled_hlo_schedule.h"],
     deps = [
+        ":tiling_specification",
         "//xla:util",
         "//xla/hlo/analysis:indexing_analysis",
+        "//xla/hlo/ir:hlo",
         "//xla/service/gpu/model/experimental:symbolic_expr",
+        "//xla/tsl/platform:errors",
+        "//xla/tsl/platform:statusor",
+        "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/types:span",
         "@llvm-project//llvm:Support",
@@ -135,15 +141,23 @@ xla_cc_test(
     name = "tiled_hlo_schedule_test",
     srcs = ["tiled_hlo_schedule_test.cc"],
     deps = [
+        ":symbolic_tile_analysis",
         ":tiled_hlo_schedule",
+        ":tiling_specification",
         "//xla/hlo/analysis:indexing_analysis",
         "//xla/hlo/analysis:interval",
+        "//xla/hlo/ir:hlo",
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+        "//xla/hlo/testlib:verified_hlo_module",
         "//xla/service/gpu/model/experimental:symbolic_expr",
+        "//xla/tsl/lib/core:status_test_util",
         "//xla/tsl/platform:statusor",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/strings:string_view",
         "@com_google_googletest//:gtest_main",
+        "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",
     ],
 )
diff --git a/xla/codegen/tiling/tiled_hlo_schedule.cc b/xla/codegen/tiling/tiled_hlo_schedule.cc
@@ -18,25 +18,37 @@ limitations under the License.
 #include <cstdint>
 #include <vector>
 
+#include "absl/algorithm/container.h"
 #include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
 #include "llvm/ADT/STLExtras.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/MLIRContext.h"
+#include "xla/codegen/tiling/tiling_specification.h"
 #include "xla/hlo/analysis/indexing_analysis.h"
 #include "xla/hlo/analysis/indexing_map.h"
+#include "xla/hlo/ir/hlo_casting_utils.h"
+#include "xla/hlo/ir/hlo_instructions.h"
+#include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/service/gpu/model/experimental/symbolic_expr.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
 #include "xla/util.h"
 
 namespace xla {
 
-absl::StatusOr<IndexingMap> MajorToMinorTiledHloSchedule::Schedule(
-    const IndexingMap& tile_offsets_indexing, IterationSpace iteration_space,
-    gpu::SymbolicExprContext* symbolic_expr_context) const {
-  mlir::MLIRContext* mlir_context = symbolic_expr_context->GetMLIRContext();
+namespace {
+
+// Helper to validate that an iteration space is compatible with a tile offsets
+// indexing map.
+absl::Status ValidateIterationSpace(const IterationSpace& iteration_space,
+                                    const IndexingMap& tile_offsets_indexing) {
   if (iteration_space.size() != tile_offsets_indexing.GetDimVarsCount()) {
     return absl::InvalidArgumentError(absl::StrFormat(
         "Expected iteration space to have exactly as many dimensions as there "
@@ -45,6 +57,30 @@ absl::StatusOr<IndexingMap> MajorToMinorTiledHloSchedule::Schedule(
         iteration_space.size(), tile_offsets_indexing.GetDimVarsCount()));
   }
 
+  std::vector<int64_t> iteration_space_dims;
+  iteration_space_dims.reserve(iteration_space.size());
+
+  for (const auto& [dim_id, dim_size] : iteration_space) {
+    if (dim_id >= tile_offsets_indexing.GetDimVarsCount() || dim_id < 0) {
+      return absl::InvalidArgumentError(absl::StrFormat(
+          "Dimension id %d is out of bounds for tile offsets indexing map with "
+          "%d dimensions. This can happen if ",
+          dim_id, tile_offsets_indexing.GetDimVarsCount()));
+    }
+
+    if (absl::c_linear_search(iteration_space_dims, dim_id)) {
+      return absl::InvalidArgumentError(absl::StrFormat(
+          "Iteration space contains multiple dimensions with id %d.", dim_id));
+    }
+    iteration_space_dims.push_back(dim_id);
+  }
+  return absl::OkStatus();
+}
+
+absl::StatusOr<IndexingMap> MajorToMinorScheduleImpl(
+    const IndexingMap& tile_offsets_indexing, IterationSpace iteration_space,
+    gpu::SymbolicExprContext* symbolic_expr_context) {
+  mlir::MLIRContext* mlir_context = symbolic_expr_context->GetMLIRContext();
   mlir::AffineExpr program_id = mlir::getAffineDimExpr(0, mlir_context);
 
   std::vector<int64_t> iteration_space_sizes;
@@ -60,12 +96,6 @@ absl::StatusOr<IndexingMap> MajorToMinorTiledHloSchedule::Schedule(
   for (auto [dim_info, tile_expr] : llvm::zip(
            iteration_space, DelinearizeIndex(iteration_space_sizes, program_id,
                                              symbolic_expr_context))) {
-    if (dim_info.dimension_id >= tile_exprs.size()) {
-      return absl::InvalidArgumentError(absl::StrFormat(
-          "Dimension id %d is out of bounds for tile offsets indexing map with "
-          "%d dimensions. This can happen if ",
-          dim_info.dimension_id, tile_exprs.size()));
-    }
     tile_exprs[dim_info.dimension_id] = tile_expr;
   }
   std::vector<IndexingMap::Variable> dim_vars{
@@ -81,5 +111,116 @@ absl::StatusOr<IndexingMap> MajorToMinorTiledHloSchedule::Schedule(
   scheduled_indexing.RemoveUnusedSymbols();
   return scheduled_indexing;
 }
+}  // namespace
+
+absl::StatusOr<IndexingMap> MajorToMinorTiledHloSchedule::Schedule(
+    const IndexingMap& tile_offsets_indexing, IterationSpace iteration_space,
+    gpu::SymbolicExprContext* ctx) const {
+  TF_RETURN_IF_ERROR(
+      ValidateIterationSpace(iteration_space, tile_offsets_indexing));
+  return MajorToMinorScheduleImpl(tile_offsets_indexing, iteration_space, ctx);
+}
+
+absl::StatusOr<TransposedDotTiledHloSchedule>
+TransposedDotTiledHloSchedule::Create(
+    const TilingSpecification& tiling_specification) {
+  const TilingSpecification::ParameterMapping& parameter_mapping =
+      tiling_specification.parameter_mapping();
+  CHECK(!parameter_mapping.empty());
+  const HloDotInstruction* dot =
+      ::xla::DynCast<HloDotInstruction>(parameter_mapping.front().instruction);
+  if (dot == nullptr) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("TransposedDotTiledHloSchedule expects its root to be a "
+                     "dot instruction "
+                     "but got ",
+                     parameter_mapping.front().instruction->ToString()));
+  }
+  if (absl::c_any_of(absl::MakeSpan(parameter_mapping).subspan(1),
+                     [](const auto& param) {
+                       return param.instruction->opcode() == HloOpcode::kDot;
+                     })) {
+    return absl::InvalidArgumentError(
+        "TransposedDotTiledHloSchedule is only supported for "
+        "TilingSpecifications specifying tiling for a single dot "
+        "instruction.");
+  }
+
+  int64_t num_lhs_non_contracting_dims =
+      dot->operand(0)->shape().dimensions().size() -
+      dot->dot_dimension_numbers().lhs_contracting_dimensions().size() -
+      dot->dot_dimension_numbers().lhs_batch_dimensions().size();
+
+  int64_t num_rhs_non_contracting_dims =
+      dot->operand(1)->shape().dimensions().size() -
+      dot->dot_dimension_numbers().rhs_contracting_dimensions().size() -
+      dot->dot_dimension_numbers().rhs_batch_dimensions().size();
+
+  constexpr absl::string_view kErrorFormat =
+      "TransposedDotTiledHloSchedule is only supported for dot instructions "
+      "with a single non-contracting dimension, but got %d non-contracting "
+      "dimensions on the %s operand of %s.";
+
+  if (num_lhs_non_contracting_dims != 1) {
+    return absl::InvalidArgumentError(absl::StrFormat(
+        kErrorFormat, num_lhs_non_contracting_dims, "lhs", dot->ToString()));
+  }
+
+  if (num_rhs_non_contracting_dims != 1) {
+    return absl::InvalidArgumentError(absl::StrFormat(
+        kErrorFormat, num_rhs_non_contracting_dims, "rhs", dot->ToString()));
+  }
+
+  // The shape of the dot's output is now known to always be of the form
+  // [..., m, n]. This is because batch dimensions precede non-contracting
+  // dimensions, the lhs non-contracting dimensions precede the rhs
+  // non-contracting dimensions, and there is exactly one such dimension on
+  // each side.
+  //
+  // Figure out the parameter index of the m and n dimensions within the op.
+  int64_t m_local_parameter_index =
+      parameter_mapping.front().num_tiling_parameters - 2;
+  int64_t n_local_parameter_index =
+      parameter_mapping.front().num_tiling_parameters - 1;
+
+  // Using the local parameter index, we can compute the global parameter index
+  // (i.e. the parameter index within the sequence of all tiling parameters).
+  TF_ASSIGN_OR_RETURN(int64_t m_dim_id, tiling_specification.ParameterIndex(
+                                            dot, m_local_parameter_index));
+  TF_ASSIGN_OR_RETURN(int64_t n_dim_id, tiling_specification.ParameterIndex(
+                                            dot, n_local_parameter_index));
+
+  return TransposedDotTiledHloSchedule(tiling_specification, m_dim_id,
+                                       n_dim_id);
+}
+
+absl::StatusOr<IndexingMap> TransposedDotTiledHloSchedule::Schedule(
+    const IndexingMap& tile_offsets_indexing, IterationSpace iteration_space,
+    gpu::SymbolicExprContext* ctx) const {
+  CHECK_EQ(iteration_space.size(), tiling_specification_.num_parameters());
+  TF_RETURN_IF_ERROR(
+      ValidateIterationSpace(iteration_space, tile_offsets_indexing));
+
+  DimensionInfo m_dim_info = iteration_space[m_dim_id_];
+  DimensionInfo n_dim_info = iteration_space[n_dim_id_];
+
+  if (m_dim_info.dimension_id != m_dim_id_) {
+    return absl::InvalidArgumentError(absl::StrFormat(
+        "Expected dimension at offset %d to have id %d but got %d.", m_dim_id_,
+        m_dim_id_, m_dim_info.dimension_id));
+  }
+  if (n_dim_info.dimension_id != n_dim_id_) {
+    return absl::InvalidArgumentError(absl::StrFormat(
+        "Expected dimension at offset %d to have id %d but got %d.", n_dim_id_,
+        n_dim_id_, n_dim_info.dimension_id));
+  }
+
+  std::vector<DimensionInfo> transposed_iteration_space(iteration_space.begin(),
+                                                        iteration_space.end());
+  transposed_iteration_space[m_dim_id_] = n_dim_info;
+  transposed_iteration_space[n_dim_id_] = m_dim_info;
+  return MajorToMinorScheduleImpl(tile_offsets_indexing,
+                                  transposed_iteration_space, ctx);
+}
 
 }  // namespace xla
diff --git a/xla/codegen/tiling/tiled_hlo_schedule.h b/xla/codegen/tiling/tiled_hlo_schedule.h
@@ -20,6 +20,7 @@ limitations under the License.
 
 #include "absl/status/statusor.h"
 #include "absl/types/span.h"
+#include "xla/codegen/tiling/tiling_specification.h"
 #include "xla/hlo/analysis/indexing_map.h"
 #include "xla/service/gpu/model/experimental/symbolic_expr.h"
 
@@ -68,7 +69,7 @@ class TiledHloSchedule {
   //     themselves);
   virtual absl::StatusOr<IndexingMap> Schedule(
       const IndexingMap& tile_offsets_indexing, IterationSpace iteration_space,
-      gpu::SymbolicExprContext* symbolic_expr_context) const = 0;
+      gpu::SymbolicExprContext* ctx) const = 0;
 };
 
 // The indexing map returned by this schedule iterates over the iteration space
@@ -79,7 +80,45 @@ class MajorToMinorTiledHloSchedule : public TiledHloSchedule {
  public:
   absl::StatusOr<IndexingMap> Schedule(
       const IndexingMap& tile_offsets_indexing, IterationSpace iteration_space,
-      gpu::SymbolicExprContext* symbolic_expr_context) const override;
+      gpu::SymbolicExprContext* ctx) const override;
+};
+
+// Given a `TilingSpecification` where some of the output tile sizes are
+// provided by a `dot` operation with one left-hand-side and one
+// right-hand-side non-contracting dimensions, this schedule transposes the
+// iteration pattern over these output dimensions.
+//
+// This schedule is only constructible when the underlying `TilingSpecification`
+// contains a single `dot` node.
+//
+// TODO(b/417977182): this is implemented as a very bespoke pattern to unblock
+// the launch of the generic emitter. We probably will want to subsume this with
+// a more flexible approach for user-specified transposed schedules (that don't
+// rely on the "dot" instruction being at the root).
+class TransposedDotTiledHloSchedule : public TiledHloSchedule {
+ public:
+  absl::StatusOr<IndexingMap> Schedule(
+      const IndexingMap& tile_offsets_indexing, IterationSpace iteration_space,
+      gpu::SymbolicExprContext* ctx) const override;
+
+  static absl::StatusOr<TransposedDotTiledHloSchedule> Create(
+      const TilingSpecification& tiling_specification);
+
+ private:
+  TransposedDotTiledHloSchedule(const TilingSpecification& tiling_specification,
+                                int64_t m_dim_id, int64_t n_dim_id)
+      : tiling_specification_(tiling_specification),
+        m_dim_id_(m_dim_id),
+        n_dim_id_(n_dim_id) {}
+
+  // The `TilingSpecification` used to construct this schedule.
+  TilingSpecification tiling_specification_;
+  // The index of the `m` dimension within the parameter mapping of the
+  // `TilingSpecification`.
+  int64_t m_dim_id_;
+  // The index of the `n` dimension within the parameter mapping of the
+  // `TilingSpecification`.
+  int64_t n_dim_id_;
 };
 
 // TODO(b/417977182): implement the `PlanarSnakeTiledHloSchedule` schedule.
diff --git a/xla/codegen/tiling/tiled_hlo_schedule_test.cc b/xla/codegen/tiling/tiled_hlo_schedule_test.cc
