Update

Author: lialan

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

@@ -907,17 +907,11 @@ def F8Types : AnyTypeOf<[
   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 TrLoadTypes:$result)> {
+    Results<(outs AnyTypeOf<[AnyVectorOfNonZeroRank]>:$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/Dialect/AMDGPU/IR/AMDGPUDialect.cpp

@@ -530,30 +530,29 @@ LogicalResult TransposeLoadOp::verify() {
   if (!hasWorkgroupMemorySpace(srcType.getMemorySpace()))
     return emitOpError("source memory address space must be Workgroup");
 
-  // TODO: support 6-bit element type vectors.
   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 numElements = transferType.getNumElements();
   size_t elementTypeSize = srcType.getElementType().getIntOrFloatBitWidth();
 
-  // ElementSize -> LoadSize
+  // ElementSize -> NumElements
   const std::map<size_t, size_t> KValidLoadSizeMap = {
-      {4, 64},
-      {32, 96}, // 6-bit element loads use casted vector<3xi32>
-      {8, 64},
-      {16, 64},
+      {4, 16},
+      {32, 3}, // 6-bit element loads use casted vector<3xi32>
+      {8, 8},
+      {16, 4},
   };
 
-  auto validLoadSize = KValidLoadSizeMap.find(elementTypeSize);
-  if (validLoadSize == KValidLoadSizeMap.end()) {
+  auto validNumElems = KValidLoadSizeMap.find(elementTypeSize);
+  if (validNumElems == 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 ";
+  if (numElements != validNumElems->second) {
+    return emitOpError(
+               "Transferring type size mismatch: expected num of elements: ")
+           << validNumElems->second;
   }
 
   return success();