Rocm jaxlib v0.5.0 warpsize global (#177)

* cherry-picked warp size passing to triton calls, and globally enabled warpsize=64

* Fix.

---------

Co-authored-by: Pavel Emeliyanenko <[email protected]>
(cherry picked from commit f013645)

Author: 2 people

xla/backends/gpu/codegen/emitters/reduction.h

@@ -121,7 +121,7 @@ class ReductionFusion : public EmitterBase {
     return IndexingMap::GetUndefined();
   }
 
-  int64_t WarpSize() const {
+  virtual int64_t WarpSize() const {
     return ::xla::gpu::WarpSize(analysis_.device_info());
   }
 
@@ -198,6 +198,11 @@ class ColumnReductionFusion : public ReductionFusion {
  public:
   explicit ColumnReductionFusion(const HloFusionAnalysis& analysis);
 
+  int64_t WarpSize() const override {
+    // PAE HACK HACK
+    return 32;
+  }
+
  protected:
   llvm::SmallVector<mlir::Value> EmitReduction(
       int group_id, EmitterState& state) const override;
@@ -216,6 +221,11 @@ class SmallColumnReductionFusion : public ReductionFusion {
  public:
   explicit SmallColumnReductionFusion(const HloFusionAnalysis& analysis);
 
+  int64_t WarpSize() const override {
+    // PAE HACK HACK
+    return 32;
+  }
+
  protected:
   llvm::SmallVector<mlir::Value> EmitReduction(
       int group_id, EmitterState& state) const override;

xla/backends/gpu/codegen/emitters/transpose.cc

@@ -75,9 +75,10 @@ using mlir::ValueRange;
 using mlir::func::FuncOp;
 using mlir::func::ReturnOp;
 
-constexpr int kNumRows = 4;
-constexpr int kNumThreadsPerBlock = 128;
-constexpr int kMaxVectorizedBytes = 4;
+constexpr int kTileSize = 32;
+constexpr int kNumRows = 8;
+constexpr int kNumThreadsPerBlock = kNumRows * kTileSize;
+constexpr int kMaxVectorizedBytes = 16;
 
 }  // namespace
 
@@ -87,7 +88,7 @@ TransposeFusion::TransposeFusion(const HloFusionAnalysis& analysis)
       permutation_(transpose_.permutation),
       input_shape_(
           Permute(transpose_.dimensions, InversePermutation(permutation_))),
-      base_block_size_(WarpSize(analysis_.device_info())) {
+      base_block_size_(kTileSize) {
   ConstHloInstructionSet transposes_to_tile;
   int index = 0;
   int64_t shmem_usage = 0;

xla/stream_executor/rocm/rocm_executor.cc

@@ -286,8 +286,7 @@ absl::StatusOr<int64_t> GetMaxRegistersPerBlock(hipDevice_t device) {
 }
 
 absl::StatusOr<int64_t> GetThreadsPerWarp(hipDevice_t device) {
-  // TODO(ROCm): This is almost certainly wrong but tests seem to rely on it.
-  return 32;
+  return GetSimpleAttribute<int64_t>(device, hipDeviceAttributeWarpSize);
 }
 
 absl::Status GetGridLimits(int* x, int* y, int* z, hipDevice_t device) {