[XLA:GPU][NFC] Pad the tile sizes before the tiling propagation.

PiperOrigin-RevId: 891661338

Author: pifon2a

third_party/xla/xla/backends/gpu/codegen/triton/BUILD

@@ -226,6 +226,7 @@ cc_library(
         "//xla/codegen/tiling:tiling_specification",
         "//xla/codegen/tiling/experimental:tile",
         "//xla/codegen/tiling/experimental:tiled_hlo",
+        "//xla/codegen/xtile/codegen:emitter_helpers",
         "//xla/codegen/xtile/codegen:fusion_emitter",
         "//xla/codegen/xtile/ir:xtile",
         "//xla/codegen/xtile/ir/transforms:passes",

third_party/xla/xla/backends/gpu/codegen/triton/xtile_compiler.cc

@@ -92,6 +92,7 @@ limitations under the License.
 #include "xla/codegen/tiling/experimental/tiling_space.h"
 #include "xla/codegen/tiling/symbolic_tile_analysis.h"
 #include "xla/codegen/tiling/tiling_specification.h"
+#include "xla/codegen/xtile/codegen/emitter_helpers.h"
 #include "xla/codegen/xtile/codegen/experimental_fusion_emitter.h"
 #include "xla/codegen/xtile/codegen/fusion_emitter.h"
 #include "xla/codegen/xtile/ir/transforms/passes.h"
@@ -216,8 +217,9 @@ absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> TileAndEmitXTileModule(
           "roots.",
           block_level_parameters.output_tile_sizes.size());
     }
-    tiling_space->AssignTileSizes(
-        block_level_parameters.output_tile_sizes.front());
+    // Triton requires that all block dimensions are a power of 2.
+    tiling_space->AssignTileSizes(xtile::GetPaddedTileSizes(
+        block_level_parameters.output_tile_sizes.front()));
 
     TileAnalysisOrError tiled_computation_or =
         TiledHloComputation::Tile(*fusion_adaptor, std::move(tiling_space));

third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.cc

@@ -30,12 +30,14 @@ limitations under the License.
 #include "absl/types/span.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/bit.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/MathExtras.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
@@ -93,6 +95,12 @@ namespace mm = ::mlir::math;
 namespace {
 using TensorValue = mlir::TypedValue<mlir::RankedTensorType>;
 
+bool ArePowersOfTwo(ArrayRef<int64_t> values) {
+  return llvm::all_of(values, [](int64_t size) {
+    return size > 0 && llvm::has_single_bit(static_cast<uint64_t>(size));
+  });
+}
+
 // Emit a value as Index clamped to [lower, upper].
 Value EmitClampedIndex(mlir::ImplicitLocOpBuilder& b, Value value,
                        int64_t lower, int64_t upper) {
@@ -543,7 +551,8 @@ Value Bitcast(mlir::ImplicitLocOpBuilder& b, Value value, Type type) {
                    tiled_hlo.tile().GetStaticTileSizes());
   ASSIGN_OR_RETURN(SmallVector<int64_t> tile_strides,
                    tiled_hlo.tile().GetStaticTileStrides());
-  auto padded_tile_sizes = GetPaddedTileSizes(tile_sizes);
+  DCHECK(ArePowersOfTwo(tile_sizes)) << "Tile sizes must be a power of 2.";
+
   const Shape& shape = tiled_hlo.hlo()->shape();
   SmallVector<int64_t> original_shape;
   original_shape.assign(shape.dimensions().begin(), shape.dimensions().end());
@@ -554,7 +563,7 @@ Value Bitcast(mlir::ImplicitLocOpBuilder& b, Value value, Type type) {
 
   auto minor_to_major_layout = llvm::to_vector(LayoutUtil::MinorToMajor(shape));
 
-  return TileInfo(offsets, tile_strides, original_shape, padded_tile_sizes,
+  return TileInfo(offsets, tile_strides, original_shape, tile_sizes,
                   minor_to_major_layout, storage_type);
 }