Revert partial [BACKEND][NFC] Clean up TritonGPU to LLVM type conversion (#5647)

Signed-off-by: Whitney Tsang <[email protected]>

Author: whitneywhtsang

include/triton/Conversion/TritonGPUToLLVM/TypeConverter.h

@@ -18,6 +18,7 @@ class TritonGPUToLLVMTypeConverter : public LLVMTypeConverter {
                                const TargetInfoBase &targetInfo,
                                const DataLayoutAnalysis *analysis = nullptr);
 
+  Type convertTritonPointerType(triton::PointerType type);
   Type convertTritonTensorType(RankedTensorType type,
                                const TargetInfoBase &targetInfo);
   Type convertMemDescType(triton::gpu::MemDescType type,

lib/Conversion/TritonGPUToLLVM/TypeConverter.cpp

@@ -16,8 +16,8 @@ TritonGPUToLLVMTypeConverter::TritonGPUToLLVMTypeConverter(
     MLIRContext *ctx, LowerToLLVMOptions &options,
     const TargetInfoBase &targetInfo, const DataLayoutAnalysis *analysis)
     : LLVMTypeConverter(ctx, options, analysis) {
-  addConversion([ctx](triton::PointerType type) -> std::optional<Type> {
-    return LLVM::LLVMPointerType::get(ctx, type.getAddressSpace());
+  addConversion([&](triton::PointerType type) -> std::optional<Type> {
+    return convertTritonPointerType(type);
   });
   addConversion([ctx](TensorDescType type) -> std::optional<Type> {
     return LLVM::LLVMPointerType::get(ctx, 1);
@@ -36,6 +36,31 @@ TritonGPUToLLVMTypeConverter::TritonGPUToLLVMTypeConverter(
                  mlir::Float8E5M2Type, mlir::Float8E5M2FNUZType>();
 }
 
+Type TritonGPUToLLVMTypeConverter::convertTritonPointerType(
+    triton::PointerType type) {
+  auto ctx = type.getContext();
+  auto pointeeType = type.getPointeeType();
+  if (isa<RankedTensorType>(pointeeType)) {
+    auto rankedTensorType = cast<RankedTensorType>(pointeeType);
+    // struct { offset0, offset1, shape0, shape1, stride0,
+    // stride1, base_ptr};
+    auto eleType = rankedTensorType.getElementType();
+    auto shape = rankedTensorType.getShape();
+    SmallVector<Type, 4> types;
+    // offsets
+    for (size_t i = 0; i < shape.size(); ++i)
+      types.push_back(IntegerType::get(ctx, 32));
+    // shapes, strides
+    for (size_t i = 0; i < 2 * shape.size(); ++i)
+      types.push_back(IntegerType::get(ctx, 64));
+
+    types.push_back(LLVM::LLVMPointerType::get(ctx, type.getAddressSpace()));
+
+    return LLVM::LLVMStructType::getLiteral(ctx, types);
+  }
+  return LLVM::LLVMPointerType::get(ctx, type.getAddressSpace());
+}
+
 Type TritonGPUToLLVMTypeConverter::convertTritonTensorType(
     RankedTensorType type, const TargetInfoBase &targetInfo) {
   auto ctx = type.getContext();