Adding support for 6-bit loadings.

Author: lialan

mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td

@@ -898,10 +898,26 @@ def AMDGPU_GatherToLDSOp :
   let hasVerifier = 1;
 }
 
+def F8Types : AnyTypeOf<[
+  F8E8M0FNU,      // 8 exponent, 0 mantissa
+  F8E5M2,         // 5 exponent, 2 mantissa
+  F8E5M2FNUZ,     // 5 exponent, 2 mantissa
+  F8E4M3,         // 4 exponent, 3 mantissa
+  F8E4M3FN,       // 4 exponent, 3 mantissa
+  F8E4M3B11FNUZ,  // 4 exponent, 3 mantissa (with bias 11)
+  F8E3M4          // 3 exponent, 4 mantissa
+]>;
+def F6Types : AnyTypeOf<[F6E2M3FN, F6E3M2FN]>;
+def TrLoadTypes : AnyTypeOf<[VectorOfLengthAndType<[4], [F16, AnyI<16>]>,
+                             VectorOfLengthAndType<[8], [F8Types, AnyI<8>]>,
+                             VectorOfLengthAndType<[16], [AnyI<4>, F6Types]>,
+                             VectorOfLengthAndType<[3], [I32]>,
+                           ]>;
+
 def AMDGPU_TransposeLoadOp :
     AMDGPU_Op<"transpose_load", [SameVariadicOperandSize]>,
     Arguments<(ins Arg<AnyMemRef, "buffer to transpose load from", [MemRead]>:$src, Variadic<Index>:$srcIndices)>,
-    Results<(outs MFMAInTypes:$result)> {
+    Results<(outs TrLoadTypes:$result)> {
   let summary = "MLIR wrapper for CDNA Transpose Load instructions";
   let description = [{
     The `amdgpu.transpose_load` op is a wrapper around the `ds_read_tr` instructions.

mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp

@@ -1115,30 +1115,37 @@ struct TransposeLoadOpLowering
 
     Location loc = op.getLoc();
     auto srcMemRefType = cast<MemRefType>(op.getSrc().getType());
+    auto resultType = cast<VectorType>(op.getResult().getType());
     Value srcPtr =
         getStridedElementPtr(rewriter, loc, srcMemRefType, adaptor.getSrc(),
                              (adaptor.getSrcIndices()));
-    auto elementTypeSize = cast<VectorType>(op.getDst().getType())
-                               .getElementType()
-                               .getIntOrFloatBitWidth();
 
-    // TODO: support ds_read_tr16_b64 intrinsic.
+    size_t numElements = resultType.getNumElements();
+    size_t elementTypeSize =
+        resultType.getElementType().getIntOrFloatBitWidth();
+
     switch (elementTypeSize) {
     case 4:
-      rewriter.replaceOpWithNewOp<ROCDL::ds_read_tr4_b64>(
-          op, op.getDst().getType(), srcPtr);
+      assert(numElements == 16);
+      rewriter.replaceOpWithNewOp<ROCDL::ds_read_tr4_b64>(op, resultType,
+                                                          srcPtr);
       break;
-    case 6:
-      rewriter.replaceOpWithNewOp<ROCDL::ds_read_tr6_b64>(
-          op, op.getDst().getType(), srcPtr);
+    case 32:
+      // To use ds_read_tr6_b96, the load size is vector<3xi32>.
+      // TODO: support native 6-bit data types.
+      assert(numElements == 3);
+      rewriter.replaceOpWithNewOp<ROCDL::ds_read_tr6_b96>(op, resultType,
+                                                          srcPtr);
       break;
     case 8:
-      rewriter.replaceOpWithNewOp<ROCDL::ds_read_tr8_b64>(
-          op, op.getDst().getType(), srcPtr);
+      assert(numElements == 8);
+      rewriter.replaceOpWithNewOp<ROCDL::ds_read_tr8_b64>(op, resultType,
+                                                          srcPtr);
       break;
     case 16:
-      rewriter.replaceOpWithNewOp<ROCDL::ds_read_tr16_b64>(
-          op, op.getDst().getType(), srcPtr);
+      assert(numElements == 4);
+      rewriter.replaceOpWithNewOp<ROCDL::ds_read_tr16_b64>(op, resultType,
+                                                           srcPtr);
       break;
     default:
       return op.emitOpError("Unsupported element size for transpose load");

mlir/lib/Dialect/AMDGPU/IR/AMDGPUDialect.cpp

@@ -531,31 +531,32 @@ LogicalResult TransposeLoadOp::verify() {
     return emitOpError("source memory address space must be Workgroup");
 
   // TODO: support 6-bit element type vectors.
-  auto transferType = dyn_cast<VectorType>(getDst().getType());
+  auto transferType = dyn_cast<VectorType>(getType());
   if (!transferType)
     return emitOpError("destination type must be a vector type");
   size_t transferSize =
       transferType.getNumElements() * transferType.getElementTypeBitWidth();
   size_t elementTypeSize = srcType.getElementType().getIntOrFloatBitWidth();
 
   // ElementSize -> LoadSize
-  const std::map<int, int> KValidLoadSizeMap = {
+  const std::map<size_t, size_t> KValidLoadSizeMap = {
       {4, 64},
-      {6, 96},
+      {32, 96}, // 6-bit element loads use casted vector<3xi32>
       {8, 64},
       {16, 64},
   };
 
   auto validLoadSize = KValidLoadSizeMap.find(elementTypeSize);
-  if (validLoadSize == KValidLoadSizeMap.end())
+  if (validLoadSize == KValidLoadSizeMap.end()) {
     return emitOpError("Unsupported element type size for transpose load: ")
            << elementTypeSize << " bits";
-  if (transferSize != validLoadSize->second)
+  }
+  if (transferSize != validLoadSize->second) {
     return emitOpError("Transferring type size must be ")
-           << validLoadSize->second
-           << " bits for element type size "
+           << validLoadSize->second << " bits for element type size ";
+  }
 
-           return success();
+  return success();
 }
 
 #include "mlir/Dialect/AMDGPU/IR/AMDGPUEnums.cpp.inc"

mlir/test/Conversion/AMDGPUToROCDL/transpose_load.mlir

@@ -16,3 +16,17 @@ func.func @transpose_load_to_rocdl_8xi8(%idx1 : index, %idx2 : index, %wgmem : m
   %0 = amdgpu.transpose_load %wgmem[%idx1, %idx2] : memref<128x128xi8, #gpu_lds_addrspace> -> vector<8xi8>
   return %0 : vector<8xi8>
 }
+
+// CHECK-LABEL: func @transpose_load_to_rocdl_16xi4
+func.func @transpose_load_to_rocdl_16xi4(%idx1 : index, %idx2 : index, %wgmem : memref<128x16xi4, #gpu_lds_addrspace>) -> vector<16xi4> {
+  // CHECK: rocdl.ds.read.tr4.b64
+  %0 = amdgpu.transpose_load %wgmem[%idx1, %idx2] : memref<128x16xi4, #gpu_lds_addrspace> -> vector<16xi4>
+  return %0 : vector<16xi4>
+}
+
+// CHECK-LABEL: func @transpose_load_to_rocdl_3xi32
+func.func @transpose_load_to_rocdl_3xi32(%idx1 : index, %idx2 : index, %wgmem : memref<128x32xi32, #gpu_lds_addrspace>) -> vector<3xi32> {
+  // CHECK: rocdl.ds.read.tr6.b96
+  %0 = amdgpu.transpose_load %wgmem[%idx1, %idx2] : memref<128x32xi32, #gpu_lds_addrspace> -> vector<3xi32>
+  return %0 : vector<3xi32>
+}