[XLA:GPU] doc updates to experimental tiling

PiperOrigin-RevId: 817064107

Author: metaflow

xla/service/gpu/model/experimental/BUILD

@@ -31,6 +31,7 @@ cc_library(
         "//xla:xla_data_proto_cc",
         "//xla/codegen/tiling:constraint_expression",
         "//xla/hlo/analysis:indexing_analysis",
+        "//xla/hlo/analysis:interval",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/utils:hlo_traversal",
         "@com_google_absl//absl/container:flat_hash_map",

xla/service/gpu/model/experimental/symbolic_tile.h

@@ -30,51 +30,69 @@ namespace xla::gpu::experimental {
 
 class TilingSpace;
 
-// A map from tile IDs, sizes and runtime variables to tile's offsets, sizes,
-// strides and upper bounds. Offsets-sizes-strides define what slice to extract,
-// upper bounds define masking, i.e. if the tile attempts to extract elements
-// with the indices outside of the bounds, the tile will be masked.
+// Tiling of a single dimension.
 //
-// (tile IDs) [tile sizes] {runtime variables} ->
-//     offsets [offsets_]  sizes [sizes_] strides [strides_]
-//     upper bounds [upper_bounds_]
+// Offsets, sizes and strides define the slice of the tensor dimension. Upper
+// bounds set the range [0, upper_bound), values outside of this range are
+// masked.
 //
-// tile IDs correspond to the dimension variables of the affine expressions;
-// tile sizes and RT vars correspond to the symbol variables.
+// Expressions for offset, size, stride and upper bound are AffineExpr
+// functions. The TilingSpace keeps track of all dimensions and symbols we use
+// in the expressions and allows to create a consistent mapping from dimensions
+// and runtime variables to affine expression dimensions and symbols.
 //
-// The masking condition of the upper bound can be written as:
-// dimension_index < upper_bounds[i](tile IDs)
+// N.B.! not all of the symbols that the TilingSpace defines are used in
+// every expression. That depends on the position of the instruction in
+// the graph and the traversal path that we took.
+// Number of dimensions equals to the number of dimensions of the instruction
+// output, parallel dimensions the corresponding root instruction are followed
+// by sequential dimensions.
 //
-// In most of the cases, the upper bounds will coincide with the shape of the
-// tensor from which the tile is extracted.
-//
-// One example when upper bound does not match the shape is a reshape:
-// output = s32[2, 17] reshape (s32[34] input)
-//
-// If we propagate the `output` tile with the ts0 == 1,
-//
-// (tid0, tid1)[ts1] -> offsets [tid0, tid1 * ts1] sizes [1, ts1] strides [1, 1]
-//              upper bounds [2, 17]
-//
-// to the `input` we will get a stricter upper bound
-//
-// (tid0, tid1)[ts1] -> offsets [17 * tid0 + tid1 * ts1] sizes [ts1] strides [1]
-//              upper bounds [17 * tid0]
+// Symbols are:
+//  - tile sizes of all dimensions, followed by
+//  - runtime variables.
 struct DimTile {
   bool operator==(const DimTile& other) const;
 
   mlir::AffineExpr offset;
   mlir::AffineExpr size;
   mlir::AffineExpr stride;
+  // The masking condition of the upper bound can be written as:
+  // dimension_index < upper_bounds(tile IDs)[tile sizes]{runtime variables}
+  //
+  // In most of the cases, the upper bounds will coincide with the shape of the
+  // tensor from which the tile is extracted. One example when upper bound does
+  // not match the shape is a reshape:
+  //
+  // output = s32[2, 17] reshape (s32[34] input)
+  //
+  // If we propagate the `output` SymbolicTile with the tile size of first
+  // dimension equal to 1
+  //
+  // (tid0, tid1)[ts1] -> offsets [tid0, tid1 * ts1]
+  //                      sizes [1, ts1]
+  //                      strides [1, 1]
+  //                      upper bounds [2, 17]
+  //
+  // then for to the `input` instruction we will get a stricter upper bound
+  //
+  // (tid0, tid1)[ts1] -> offsets [17 * tid0 + tid1 * ts1]
+  //                      sizes [ts1]
+  //                      strides [1]
+  //                      upper bounds [17 * tid0]
   mlir::AffineExpr upper_bound;
 };
+
 template <typename H>
 H AbslHashValue(H h, const DimTile& dim_tile) {
   llvm::hash_code dim_tile_hash = llvm::hash_combine(
       dim_tile.offset, dim_tile.size, dim_tile.stride, dim_tile.upper_bound);
   return H::combine(std::move(h), static_cast<size_t>(dim_tile_hash));
 }
 
+// SymbolicTile is a collection of tilings for every dimension of output tensor
+// of an HLO instruction. SymbolicTiledHloInstruction associates a SymbolicTile
+// with an HLO instruction.
 class SymbolicTile {
  public:
   SymbolicTile(const TilingSpace& tiling_space,
@@ -120,8 +138,9 @@ H AbslHashValue(H h, const SymbolicTile& symbolic_tile) {
   return h;
 }
 
-// Returns a DimTile that covers the entire dimension, i.e.
-// offset 0, size = next_power_of_2(dim_size), stride 1, upper_bound = dim_size.
+// Returns a DimTile that covers the entire dimension with a single power of 2
+// sized tile, i.e. offset 0, size = next_power_of_2(dim_size), stride 1,
+// upper_bound = dim_size.
 DimTile GetFullDimTile(int64_t dim_size, mlir::MLIRContext* ctx);
 
 // Returns a DimTile that covers the entire dimension, i.e.

xla/service/gpu/model/experimental/tiling_space.cc

@@ -27,7 +27,7 @@ limitations under the License.
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "mlir/IR/MLIRContext.h"
-#include "xla/hlo/analysis/indexing_map.h"
+#include "xla/hlo/analysis/interval.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
@@ -65,6 +65,23 @@ void TilingSpace::AppendRTVar(const HloInstruction* hlo, int64_t operand_id,
   hlo_to_rt_var_[std::make_pair(hlo, operand_id)] = &rt_vars_.back();
 }
 
+void TilingSpace::ProcessInstruction(const HloInstruction& hlo) {
+  switch (hlo.opcode()) {
+    case HloOpcode::kDot:
+      ProcessDot(hlo);
+      break;
+    case HloOpcode::kReduce:
+      ProcessReduce(hlo);
+      break;
+    case HloOpcode::kDynamicSlice:
+      ProcessDynamicSlice(hlo);
+      break;
+    default:
+      // TODO(goncharov): should have a explicit list of supported instructions?
+      break;
+  }
+}
+
 // Add dot contraction dimensions in the order of contracting dimensions.
 void TilingSpace::ProcessDot(const HloInstruction& hlo) {
   auto dot = Cast<HloDotInstruction>(&hlo);
@@ -143,15 +160,17 @@ const TilingSpace::DimensionInfo& TilingSpace::GetDimensionInfo(
     const HloInstruction& hlo, int64_t dim_position) const {
   auto it = hlo_to_dimension_.find(std::make_pair(&hlo, dim_position));
   CHECK(it != hlo_to_dimension_.end())
-      << "Dimension not found: " << hlo.ToString() << " " << dim_position;
+      << "Dimension not found for " << hlo.ToString() << " dimension "
+      << dim_position;
   return *it->second;
 }
 
 const TilingSpace::RTVarInfo& TilingSpace::GetRTVarInfo(
     const HloInstruction& hlo, int64_t operand_id) const {
   auto it = hlo_to_rt_var_.find(std::make_pair(&hlo, operand_id));
   CHECK(it != hlo_to_rt_var_.end())
-      << "Runtime variable not found: " << hlo.ToString();
+      << "Runtime variable not found for " << hlo.ToString() << " operand "
+      << operand_id;
   return *it->second;
 }
 
@@ -163,8 +182,10 @@ std::unique_ptr<TilingSpace> TilingSpace::Create(const HloFusionAdaptor& fusion,
   for (const HloInstructionAdaptor& root : roots) {
     const Shape& root_shape = root.shape();
     if (!root.shape().IsArray() && root.opcode() != HloOpcode::kReduce) {
-      LOG(FATAL) << "Unsupported root shape: " << root_shape.ToString();
+      LOG(FATAL) << "Unsupported root shape " << root_shape.ToString()
+                 << " for root " << root.instruction().ToString();
     }
+    // TODO(goncharov): why do we only care about the first shape of a tuple?
     absl::Span<const int64_t> dims =
         GetFirstShape(&root.instruction()).dimensions();
     llvm::SmallVector<DimTile> dim_tiles;
@@ -186,19 +207,7 @@ std::unique_ptr<TilingSpace> TilingSpace::Create(const HloFusionAdaptor& fusion,
   // Iterator in reversed post-order (use-before-def).
   auto post_order = fusion.MakeInstructionPostOrder();
   for (auto it = post_order.rbegin(); it != post_order.rend(); ++it) {
-    switch (it->instruction().opcode()) {
-      case HloOpcode::kDot:
-        tiling_space->ProcessDot(it->instruction());
-        break;
-      case HloOpcode::kReduce:
-        tiling_space->ProcessReduce(it->instruction());
-        break;
-      case HloOpcode::kDynamicSlice:
-        tiling_space->ProcessDynamicSlice(it->instruction());
-        break;
-      default:
-        break;
-    }
+    tiling_space->ProcessInstruction(it->instruction());
   }
   return tiling_space;
 }

xla/service/gpu/model/experimental/tiling_space.h

@@ -26,22 +26,26 @@ limitations under the License.
 #include "llvm/ADT/SmallVector.h"
 #include "mlir/IR/MLIRContext.h"
 #include "xla/codegen/tiling/constraint_expression.h"
-#include "xla/hlo/analysis/indexing_map.h"
+#include "xla/hlo/analysis/interval.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/utils/hlo_traversal.h"
 #include "xla/service/gpu/model/experimental/symbolic_tile.h"
 #include "xla/shape.h"
 
 namespace xla::gpu::experimental {
 
-// TilingSpace contains information about all parallel and sequential dimensions
-// and runtime variables in a fusion.
-// The parallel dimensions correspond to the dimensions of the outputs of the
-// fusion.
-// The sequential dimensions correspond to the contraction/reduction dimensions
-// of the dots/reduces in the fusion.
-// The runtime variables correspond to the offsets of the dynamic slices in the
-// fusion.
+// TilingSpace holds information about all tiling parameters of a fusion.
+//
+// It defines symbolic tiles for the fusion roots as symbolic expressions and
+// constraints of possible tile "variables":
+// * parallel dimensions - output dimensions of the fusion;
+// * sequential dimensions - contraction/reduction dimensions of operations in
+//   the fusion;
+// * runtime variables - for example, offsets of the dynamic slices.
+//
+// This information allows us later to explore the space of all possible tilings
+// and assign concrete tilings for every instruction of the fusion with
+// SymbolicTilePropagation.
 class TilingSpace {
  public:
   TilingSpace() : constraints_(ConstraintExpression::GetAlwaysSatisfied()) {}
@@ -51,29 +55,46 @@ class TilingSpace {
 
   enum class DimensionSemantics { kParallel, kSequential };
   struct DimensionInfo {
-    // Unique ID for the dimension.
+    // Unique ID for the dimension within the tiling space.
     ID id;
     // Size of the dimension.
     int64_t dimension_size;
     // Type of the dimension.
     DimensionSemantics type;
-    // HLO instruction that defines the dimension.
+    // HLO instruction that defines (introduces) the dimension. For example
+    // fusion root instruction defines the parallel dimensions. Dot/reduce
+    // defines the sequential (contraction) dimensions.
     const HloInstruction* hlo;
-    // Index into the ordered list of dimensions of the HLO instruction.
-    // All dimensions in the HLO instruction are described as
+    // Index into the ordered list of dimensions of the HLO instruction `hlo`
+    // that defines the dimension.
+    // All dimensions in the HLO instruction are ordered as
     // [all parallel dims of the output, all reduction/contraction dims].
     //
-    // Example:
-    // [output_dims] dot(lhs, rhs, lhs_contracting_dims, rhs_contracting_dims)
-    // The dimensions are ordered as [output_dims, LHS[lhs_contracting_dims]].
+    // Example, for `[a,b,c] = dot(lhs, rhs, lhs_contracting_dims={d,e}, ...)`.
+    // The ordered list of dimensions is [a,b,c,d,e].
     int64_t dim_position;
   };
 
+  // Information about a runtime variable.
+  // For example:
+  //
+  // off = s32[] parameter(0)
+  // ds = dynamic-slice(tensor, off), ...
+  //
+  // `off = s32[] parameter(0)` instruction (`hlo`) defines the runtime
+  // variable.
+  // User's (dynamic-slice) semantics sets the `bounds` of possible values.
+  //
+  // If the same hlo is used as runtime variable multiple times, there will be
+  // multiple entries in the `rt_vars_` with different IDs.
+  //
+  // RTVarInfo are accessed by (user_hlo, operand_id), in this case it is
+  // (dynamic-slice, 1).
   struct RTVarInfo {
-    // Unique ID for the runtime variable.
+    // Unique ID for the runtime variable within the tiling space.
     ID id;
-    // Feasible bounds of the runtime variable. The values outside of the bounds
-    // will be clamped.
+    // Feasible bounds of the runtime variable.
+    // The values outside of the bounds will be clamped.
     Interval bounds;
     // HLO instruction that defines the runtime variable.
     const HloInstruction* hlo;
@@ -90,9 +111,15 @@ class TilingSpace {
     sink.Append(space.ToString());
   }
 
+  // Returns the dimension info for the given `hlo` and `dim_position`.
+  // `dim_position` is the index into the ordered list of dimensions of the HLO
+  // instruction `hlo` that defines the dimension. The dimension info must
+  // exist.
   const DimensionInfo& GetDimensionInfo(const HloInstruction& hlo,
                                         int64_t dim_position) const;
 
+  // Returns the runtime variable info for `hlo` that uses it and its
+  // `operand_id`. This runtime variable info must exist.
   const RTVarInfo& GetRTVarInfo(const HloInstruction& hlo,
                                 int64_t operand_id) const;
 
@@ -115,6 +142,7 @@ class TilingSpace {
   void ProcessDot(const HloInstruction& hlo);
   void ProcessReduce(const HloInstruction& hlo);
   void ProcessDynamicSlice(const HloInstruction& hlo);
+  void ProcessInstruction(const HloInstruction& hlo);
 
   // Maps from (hlo, dim_position) to the dimension info.
   absl::flat_hash_map<std::pair<const HloInstruction*, int64_t>,
@@ -130,7 +158,9 @@ class TilingSpace {
   // The deque is used to guarantee the pointer stability.
   std::deque<RTVarInfo> rt_vars_;
 
-  // Root instruction of the fusion.
+  // Symbolic tiles for the fusion roots.
+  // For tuple roots, there will be one tile per tuple element. Otherwise,
+  // there will be only one symbolic tile.
   llvm::SmallVector<SymbolicTile, 2> tiled_roots_;
 
   // Constraint expression for the tiling space.