Adding 6-bit loads.

Author: lialan

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

@@ -901,20 +901,20 @@ def AMDGPU_GatherToLDSOp :
 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:$dst)> {
+    Results<(outs MFMAInTypes:$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.
 
     Operands:
     * `$src`: LDS memref to read from.
     * `$srcIndices`: indices into `$src` to read from for this thread.
-    * `$dst`: target register this transpose load instruction will write to.
+    * `$result`: target register this transpose load instruction will write to.
 
     Note: Lowering is only supported on gfx950 and up.
   }];
   let assemblyFormat = [{
-    $src `[` $srcIndices `]` attr-dict `:` type($src) `->` type($dst)
+    $src `[` $srcIndices `]` attr-dict `:` type($src) `->` type($result)
   }];
   let hasVerifier = 1;
 }

mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp

@@ -1128,6 +1128,10 @@ struct TransposeLoadOpLowering
       rewriter.replaceOpWithNewOp<ROCDL::ds_read_tr4_b64>(
           op, op.getDst().getType(), srcPtr);
       break;
+    case 6:
+      rewriter.replaceOpWithNewOp<ROCDL::ds_read_tr6_b64>(
+          op, op.getDst().getType(), srcPtr);
+      break;
     case 8:
       rewriter.replaceOpWithNewOp<ROCDL::ds_read_tr8_b64>(
           op, op.getDst().getType(), srcPtr);

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

@@ -536,10 +536,26 @@ LogicalResult TransposeLoadOp::verify() {
     return emitOpError("destination type must be a vector type");
   size_t transferSize =
       transferType.getNumElements() * transferType.getElementTypeBitWidth();
-  if (transferSize != 64)
-    return emitOpError("Transferring type size must be 64 bits");
-
-  return success();
+  size_t elementTypeSize = srcType.getElementType().getIntOrFloatBitWidth();
+
+  // ElementSize -> LoadSize
+  const std::map<int, int> KValidLoadSizeMap = {
+      {4, 64},
+      {6, 96},
+      {8, 64},
+      {16, 64},
+  };
+
+  auto validLoadSize = KValidLoadSizeMap.find(elementTypeSize);
+  if (validLoadSize == KValidLoadSizeMap.end())
+    return emitOpError("Unsupported element type size for transpose load: ")
+           << elementTypeSize << " bits";
+  if (transferSize != validLoadSize->second)
+    return emitOpError("Transferring type size must be ")
+           << validLoadSize->second
+           << " bits for element type size "
+
+           return success();
 }
 
 #include "mlir/Dialect/AMDGPU/IR/AMDGPUEnums.cpp.inc"