[NFC] Prepare utilities for warp specialization lowering on AMD (#8968)

Exposes additional utilities needed for warp specialization lowering on
AMD.

Author: PMylon

include/triton/Conversion/TritonGPUToLLVM/Utility.h

@@ -468,6 +468,10 @@ Value mxfpScaleBf16(RewriterBase &rewriter, Location loc, Value v, Value scale,
 // Hardware Indices
 // -----------------------------------------------------------------------
 
+// If an operation is contained within a warp specialize region, this returns
+// the warp ID offset of that warpgroup.
+std::optional<int> getWarpGroupStartWarpId(Block *block);
+
 // If an operation is contained within a warp specialize region, this returns
 // the thread ID offset of that warpgroup.
 std::optional<int> getWarpGroupStartThreadId(Block *block);

lib/Conversion/TritonGPUToLLVM/Utility.cpp

@@ -301,7 +301,7 @@ applyLinearLayout(Location loc, RewriterBase &rewriter,
   return outIndices;
 }
 
-std::optional<int> getWarpGroupStartThreadId(Block *block) {
+std::optional<int> getWarpGroupStartWarpId(Block *block) {
   using namespace triton::gpu;
 
   // Look for an enclosing `ttg.warp_specialize` op.
@@ -317,9 +317,19 @@ std::optional<int> getWarpGroupStartThreadId(Block *block) {
   std::optional<ArrayRef<int32_t>> startIds = ws.getWarpGroupStartIds();
   assert(startIds && "cannot get warp group ID before warp group allocation");
   int32_t warpStartId = (*startIds)[idx];
-  int threadsPerWarp =
-      TritonGPUDialect::getThreadsPerWarp(ws->getParentOfType<ModuleOp>());
-  return warpStartId * threadsPerWarp;
+  return warpStartId;
+}
+
+std::optional<int> getWarpGroupStartThreadId(Block *block) {
+  using namespace triton::gpu;
+
+  std::optional<int> warpStartId = getWarpGroupStartWarpId(block);
+  if (!warpStartId)
+    return {};
+
+  int threadsPerWarp = TritonGPUDialect::getThreadsPerWarp(
+      block->getParentOp()->getParentOfType<ModuleOp>());
+  return *warpStartId * threadsPerWarp;
 }
 
 Value getThreadId(OpBuilder &rewriter, Location loc) {

third_party/amd/include/TritonAMDGPUToLLVM/TypeConverter.h

@@ -0,0 +1,43 @@
+#ifndef TRITON_CONVERSION_TRITONAMDGPU_TO_LLVM_TYPECONVERTER_H
+#define TRITON_CONVERSION_TRITONAMDGPU_TO_LLVM_TYPECONVERTER_H
+
+#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "triton/Conversion/MLIRTypes.h"
+#include "triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h"
+#include "triton/Conversion/TritonGPUToLLVM/TypeConverter.h"
+#include "triton/Dialect/Triton/IR/Types.h"
+#include "triton/Dialect/TritonGPU/IR/Types.h"
+
+using namespace mlir;
+using namespace mlir::triton;
+
+class TritonAMDGPUToLLVMTypeConverter : public TritonGPUToLLVMTypeConverter {
+public:
+  TritonAMDGPUToLLVMTypeConverter(MLIRContext *ctx,
+                                  const LowerToLLVMOptions &options,
+                                  const TargetInfoBase &targetInfo,
+                                  const DataLayoutAnalysis *analysis = nullptr)
+      : TritonGPUToLLVMTypeConverter(ctx, options, targetInfo, analysis) {
+    addConversion([&](TensorDescType type) -> std::optional<Type> {
+      return convertTensorDescType(type);
+    });
+  }
+
+  Type convertTensorDescType(triton::TensorDescType type) {
+    auto ctx = type.getContext();
+    auto blockType = type.getBlockType();
+    auto shape = blockType.getShape();
+
+    // Determine the number of dwords based on tensor dimensions
+    // 2D tensors: group0 (4) + group1 (8) = 12 dwords
+    // 3D-5D tensors: group0 (4) + group1 (8) + group2 (4) + group3 (4) = 20
+    // dwords
+    int numDwords = (shape.size() > 2) ? (4 + 8 + 4 + 4) : (4 + 8);
+
+    auto types = SmallVector<Type>(numDwords, IntegerType::get(ctx, 32));
+    return LLVM::LLVMStructType::getLiteral(ctx, types);
+  }
+};
+
+#endif

third_party/amd/lib/TritonAMDGPUToLLVM/TritonGPUToLLVM.cpp

@@ -4,6 +4,7 @@
 #include "PatternTritonGPUOpToLLVM.h"
 #include "TargetInfo.h"
 #include "TritonAMDGPUToLLVM/MembarUtility.h"
+#include "TritonAMDGPUToLLVM/TypeConverter.h"
 #include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
 #include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
@@ -63,34 +64,6 @@ class TritonLLVMConversionTarget : public ConversionTarget {
   }
 };
 
-class TritonAMDGPUToLLVMTypeConverter : public TritonGPUToLLVMTypeConverter {
-public:
-  TritonAMDGPUToLLVMTypeConverter(MLIRContext *ctx,
-                                  const LowerToLLVMOptions &options,
-                                  const TargetInfoBase &targetInfo,
-                                  const DataLayoutAnalysis *analysis = nullptr)
-      : TritonGPUToLLVMTypeConverter(ctx, options, targetInfo, analysis) {
-    addConversion([&](TensorDescType type) -> std::optional<Type> {
-      return convertTensorDescType(type);
-    });
-  }
-
-  Type convertTensorDescType(triton::TensorDescType type) {
-    auto ctx = type.getContext();
-    auto blockType = type.getBlockType();
-    auto shape = blockType.getShape();
-
-    // Determine the number of dwords based on tensor dimensions
-    // 2D tensors: group0 (4) + group1 (8) = 12 dwords
-    // 3D-5D tensors: group0 (4) + group1 (8) + group2 (4) + group3 (4) = 20
-    // dwords
-    int numDwords = (shape.size() > 2) ? (4 + 8 + 4 + 4) : (4 + 8);
-
-    auto types = SmallVector<Type>(numDwords, IntegerType::get(ctx, 32));
-    return LLVM::LLVMStructType::getLiteral(ctx, types);
-  }
-};
-
 struct ConvertTritonAMDGPUToLLVM
     : public triton::impl::ConvertTritonAMDGPUToLLVMBase<
           ConvertTritonAMDGPUToLLVM> {