Add expand_shape to GPU encoding pipeline

Signed-off-by: Alan Li <[email protected]>

compiler/src/iree/compiler/Codegen/Common/EncodingUtils.h

@@ -20,6 +20,9 @@ struct MaterializeEncodingInfo {
   SmallVector<int64_t> innerTileSizes;
   SmallVector<int64_t> outerDimsPerm;
   unsigned srcRank = 0;
+  // Metadata for a generalized expand_shape + transpose
+  SmallVector<int64_t> innerTileShapes;
+  SmallVector<int64_t> permutation;
 };
 
 using MaterializeEncodingFn =

compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel

@@ -59,8 +59,8 @@ iree_compiler_cc_library(
         "GPUDistributionPatterns.cpp",
         "GPUGeneralizeNamedOps.cpp",
         "GPULowerToUKernels.cpp",
-        "GPUMultiBuffering.cpp",
         "GPUMaterializeEncoding.cpp",
+        "GPUMultiBuffering.cpp",
         "GPUNestedLayoutDistributionPatterns.cpp",
         "GPUPatterns.cpp",
         "GPUPipelining.cpp",

compiler/src/iree/compiler/Codegen/Common/GPU/GPUMaterializeEncoding.cpp

@@ -6,10 +6,16 @@
 
 #include "iree/compiler/Codegen/Common/EncodingUtils.h"
 #include "iree/compiler/Codegen/Common/GPU/Passes.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
 #include "mlir/Dialect/MemRef/Transforms/Transforms.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/Transforms/Transforms.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
+#include "mlir/IR/Operation.h"
+
 #define DEBUG_TYPE "iree-codegen-gpu-materialize-encoding"
 
 namespace mlir::iree_compiler {
@@ -30,16 +36,49 @@ static std::optional<TileMxNxK> getIntrinsicSize(TypeRange elementTypes) {
 
 // TODO: Query the value from GPU attributes.
 // TODO: Define a struct with meaningful name for the pair.
-std::optional<std::pair<int64_t, int64_t>>
-getIntrinsicVectorSize(TypeRange elementTypes, int64_t roleIdx) {
+SmallVector<int64_t> getIntrinsicVectorSize(TypeRange elementTypes,
+                                            int64_t roleIdx) {
   Type lhs = elementTypes[0];
   Type rhs = elementTypes[1];
   Type out = elementTypes[2];
   if (lhs.isF32() && rhs.isF32() && out.isF32()) {
-    if (roleIdx == 0 || roleIdx == 1) return std::make_pair(1, 1);
-    if (roleIdx == 2) return std::make_pair(4, 1);
+    if (roleIdx == 0 || roleIdx == 1) {
+      return {1, 1};
+    }
+    if (roleIdx == 2) {
+      return {4, 1};
+    }
+  }
+  return {};
+}
+
+SmallVector<int64_t> getTransposePermutation(int64_t roleIdx) {
+  switch (roleIdx) {
+  case 0: // A
+  case 1: // B
+    // OuterTileX x InnerTileX x OuterTileY x InnerTileY
+    // -> OuterTileY x OuterTileX x InnerTileY x InnerTileX
+    return {2, 0, 3, 1};
+  case 2: // C
+    // ACC:
+    // OuterTileX x InnerTileX x OuterTileY x InnerTileY
+    // -> OuterTileX x OuterTileY x InnerTileX x InnerTileY
+    return {0, 2, 1, 3};
+  default:
+    llvm_unreachable("unexpected roleIdx");
+  }
+}
+
+SmallVector<int64_t> getReverseTransposePermutation(int64_t roleIdx) {
+  switch (roleIdx) {
+  case 0: // A
+  case 1: // B
+    return {1, 2, 0, 3};
+  case 2: // C
+    return {0, 2, 1, 3};
+  default:
+    llvm_unreachable("unexpected roleIdx");
   }
-  return std::nullopt;
 }
 
 // TODO(hanchung): Pass an ExecutableTargetAttr attribute for the target
@@ -88,7 +127,17 @@ materializeEncodingForTarget(RankedTensorType tensorType) {
   // Map the matmul TileMxNxK to an actual tile shape for the tensor at hand,
   // based on its operand index in the matmul.
   auto rank = tensorType.getRank();
-  return getEncodingInfoForMatmul(encoding, rank, enumeratedTileMxNxK[0]);
+
+  auto encodingInfo =
+      getEncodingInfoForMatmul(encoding, rank, enumeratedTileMxNxK[0]);
+
+  // insert inner tile shapes and permutation info
+  auto roleIdx = encoding.getOperandIndex().getInt();
+  auto intrinsicVectorSizes = getIntrinsicVectorSize(elementTypes, roleIdx);
+  auto permutation = getTransposePermutation(roleIdx);
+  encodingInfo.innerTileShapes = intrinsicVectorSizes;
+  encodingInfo.permutation = permutation;
+  return encodingInfo;
 }
 
 namespace {
@@ -114,6 +163,7 @@ struct GPUSetEncodingOpLoweringConversion
         getTypeConverter());
     MaterializeEncodingFn materializeEncodingFn =
         converter->getMaterializeEncodingFn();
+
     auto packOp = lowerSetEncodingOpToPackOp(
         rewriter, encodingOp, adaptor.getSource(), materializeEncodingFn,
         this->materializeEncodingValueFn);
@@ -136,6 +186,9 @@ struct GPUSetEncodingOpLoweringConversion
                                          "unhandled result encoding");
     }
     SmallVector<int64_t> innerTiles = maybeEncodingInfo->innerTileSizes;
+    SmallVector<int64_t> intrinsicVectorShape =
+        maybeEncodingInfo->innerTileShapes;
+    SmallVector<int64_t> transposePermutation = maybeEncodingInfo->permutation;
 
     // TODO(hanchung): Add a util to the encoding attribute, so we don't need
     // the map_to_vector method here.
@@ -144,26 +197,26 @@ struct GPUSetEncodingOpLoweringConversion
     auto elemTypes = llvm::map_to_vector(
         encoding.getElementTypes().getValue(),
         [](Attribute a) { return cast<TypeAttr>(a).getValue(); });
+    auto loc = encodingOp.getLoc();
+
     std::optional<TileMxNxK> intrinsicShape = getIntrinsicSize(elemTypes);
-    std::optional<std::pair<int64_t, int64_t>> intrinsicVectorShape =
-        getIntrinsicVectorSize(elemTypes, roleIdx);
-    if (!intrinsicShape || !intrinsicVectorShape) {
+    if (!intrinsicShape || intrinsicVectorShape.empty()) {
       return failure();
     }
 
     SmallVector<int64_t> targetShape; // for unrolling
-    switch(roleIdx) {
-      case 0:
-        targetShape = {intrinsicShape->M, intrinsicShape->K};
-        break;
-      case 1:
-        targetShape = {intrinsicShape->N, intrinsicShape->K};
-        break;
-      case 2:
-        targetShape = {intrinsicShape->M, intrinsicShape->N};
-        break;
-      default:
-        return failure();
+    switch (roleIdx) {
+    case 0: // A
+      targetShape = {intrinsicShape->M, intrinsicShape->K};
+      break;
+    case 1: // B
+      targetShape = {intrinsicShape->N, intrinsicShape->K};
+      break;
+    case 2: // C
+      targetShape = {intrinsicShape->M, intrinsicShape->N};
+      break;
+    default:
+      return failure();
     }
 
     assert(innerTiles.size() == targetShape.size());
@@ -175,15 +228,57 @@ struct GPUSetEncodingOpLoweringConversion
       assert(packedShape == targetShape);
     }
 
-    // TODO(lialan): create expand_shape. Take LHS as an example:
-    // 16x4xf32 -> 16x1x4x1. Because the vector size used in the intrinsic is
-    // 1x1.
-    // For C-Matrix (i.e., ACC), it is 16x16xf32 -> 4x4x16x1xf32. Because the
-    // vector size is 4x1.
+    // Check that the dimensions of the matrix can be divided by the tile shape,
+    // if not then bail out.
+    auto sourceShape = packOp->getResult().getType().getShape();
+    assert(sourceShape.size() == 4);
+    // inner most two dimensions must be divisible.
+    if (sourceShape[2] % innerTiles[0] != 0 ||
+        sourceShape[3] % innerTiles[1] != 0) {
+      return failure();
+    }
+
+    // Create expand_shape, to tile the inner most two dimensions.
+    llvm::SmallVector<int64_t> expandShapeShape;
+    assert(intrinsicVectorShape.size() == 2); // TODO: relax this
+    auto iT1 = intrinsicVectorShape[0];
+    auto iT2 = intrinsicVectorShape[1];
+    auto oT1 = sourceShape[2] / iT1;
+    auto oT2 = sourceShape[3] / iT2;
+    expandShapeShape = {sourceShape[0], sourceShape[1], oT1, iT1, oT2, iT2};
+    assert(expandShapeShape.size() == 6);
+    RankedTensorType expandShapeType = RankedTensorType::get(
+        expandShapeShape, encodingOp.getSourceType().getElementType());
 
-    // TODO(lialan): create linalg.transpose op.
-    // LHS: 16x1x4x1 -> 4x16x1x1 (perm = [2, 0, 3, 1])
-    // ACC: 4x4x16x1 -> 4x16x4x1 (perm = [0, 2, 1, 3])
+    SmallVector<ReassociationExprs> reassociationMap(4);
+    reassociationMap[0].push_back(rewriter.getAffineDimExpr(0));
+    reassociationMap[1].push_back(rewriter.getAffineDimExpr(1));
+    reassociationMap[2].push_back(rewriter.getAffineDimExpr(2));
+    reassociationMap[2].push_back(rewriter.getAffineDimExpr(3));
+    reassociationMap[3].push_back(rewriter.getAffineDimExpr(4));
+    reassociationMap[3].push_back(rewriter.getAffineDimExpr(5));
+    auto expandShapeOp = rewriter.create<tensor::ExpandShapeOp>(
+        loc, expandShapeType, packOp->getResult(), reassociationMap);
+
+    // create linalg.transpose on expandShapeShape
+    SmallVector<int64_t> transposeResultDims;
+    transposeResultDims.push_back(expandShapeShape[0]);
+    transposeResultDims.push_back(expandShapeShape[1]);
+    for (int i = 0; i < transposePermutation.size(); i++) {
+      transposeResultDims.push_back(
+          expandShapeShape[2 + transposePermutation[i]]);
+    }
+    SmallVector<int64_t> newTransposePermutation(transposePermutation.size() +
+                                                 2);
+    newTransposePermutation[0] = 0;
+    newTransposePermutation[1] = 1;
+    for (int i = 0; i < transposePermutation.size(); i++) {
+      newTransposePermutation[i + 2] = transposePermutation[i] + 2;
+    }
+    auto emptyTensor = rewriter.create<tensor::EmptyOp>(
+        loc, transposeResultDims, encodingOp.getSourceType().getElementType());
+    auto transposeOp = rewriter.create<linalg::TransposeOp>(
+        loc, expandShapeOp, emptyTensor, newTransposePermutation);
 
     // TODO(hanchung): We want to make the shape consistent, so we need to
     // collpase and expand the shape. This is the shape we materialize for Flow
@@ -195,8 +290,35 @@ struct GPUSetEncodingOpLoweringConversion
     //    LHS: 64 -> 16x4 (innerTiles[0]xinnerTiles[1])
     //    ACC: 256 -> 16x16 (innerTiles[0]xinnerTiles[1])
 
-    // TODO(lialan): Replace the op with the tensor.expand_shape op.
-    rewriter.replaceOp(encodingOp, packOp->getResult());
+    // collapse tiled dimensions into one dim
+    SmallVector<ReassociationIndices> collapseReassoc = {
+        {0}, {1}, {2, 3, 4, 5}};
+    SmallVector<int64_t> collapsedShape = {sourceShape[0], sourceShape[1],
+                                           sourceShape[2] * sourceShape[3]};
+    RankedTensorType revertShapeType = RankedTensorType::get(
+        collapsedShape, encodingOp.getSourceType().getElementType());
+
+    auto collapseShapeOp = rewriter.create<tensor::CollapseShapeOp>(
+        loc, revertShapeType, transposeOp->getResult(0),
+        ArrayRef<ReassociationIndices>(collapseReassoc));
+
+    // expand the collapsed shape to the shape intended by the encoding
+    // Create expand_shape, to tile the inner most two dimensions.
+    assert(innerTiles.size() == 2); // TODO: relax this
+    assert(innerTiles[0] * innerTiles[1] == sourceShape[2] * sourceShape[3]);
+    RankedTensorType expandTileShapeType = RankedTensorType::get(
+        {sourceShape[0], sourceShape[1], innerTiles[0], innerTiles[1]},
+        encodingOp.getSourceType().getElementType());
+    SmallVector<ReassociationExprs> tileAssoc(3);
+    tileAssoc[0].push_back(rewriter.getAffineDimExpr(0));
+    tileAssoc[1].push_back(rewriter.getAffineDimExpr(1));
+    tileAssoc[2].push_back(rewriter.getAffineDimExpr(2));
+    tileAssoc[2].push_back(rewriter.getAffineDimExpr(3));
+
+    auto expandTileShapeOp = rewriter.create<tensor::ExpandShapeOp>(
+        loc, expandTileShapeType, collapseShapeOp, tileAssoc);
+
+    rewriter.replaceOp(encodingOp, expandTileShapeOp);
     return success();
   }
 };

compiler/src/iree/compiler/Codegen/LLVMCPU/test/BUILD.bazel

@@ -28,6 +28,7 @@ iree_lit_test_suite(
             "check_ir_before_llvm_conversion.mlir",
             "check_ir_before_llvm_conversion_not_fail_unbound.mlir",
             "convert_to_llvm.mlir",
+            "data_tile.mlir",
             "emit_vectorization_remarks.mlir",
             "expand_f16_op_to_f32.mlir",
             "hal_executable_constants.mlir",

compiler/src/iree/compiler/Codegen/LLVMCPU/test/CMakeLists.txt

@@ -23,6 +23,7 @@ iree_lit_test_suite(
     "check_ir_before_llvm_conversion.mlir"
     "check_ir_before_llvm_conversion_not_fail_unbound.mlir"
     "convert_to_llvm.mlir"
+    "data_tile.mlir"
     "emit_vectorization_remarks.mlir"
     "expand_f16_op_to_f32.mlir"
     "hal_executable_constants.mlir"

compiler/src/iree/compiler/Codegen/LLVMCPU/test/data_tile.mlir

@@ -0,0 +1,27 @@
+// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-gpu-materialize-device-encoding))" --split-input-file %s | FileCheck %s
+#pipeline_layout = #hal.pipeline.layout<push_constants = 0, sets = [
+  #hal.descriptor_set.layout<0, bindings = [
+    #hal.descriptor_set.binding<0, storage_buffer>,
+    #hal.descriptor_set.binding<1, storage_buffer>
+  ]>
+]>
+func.func @set_encoding() {
+  %c0 = arith.constant 0 : index
+  %0 = hal.interface.binding.subspan layout(#pipeline_layout) set(0) binding(0) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<255x513xf32>>
+  %1 = hal.interface.binding.subspan layout(#pipeline_layout) set(0) binding(1) alignment(64) offset(%c0) : !flow.dispatch.tensor<writeonly:tensor<255x513xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], original_type = tensor<255x513xf32>, user_indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>], round_dims_to = array<i64: 16, 16, 16>>>>
+  %2 = flow.dispatch.tensor.load %0, offsets = [0, 0], sizes = [255, 513], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<255x513xf32>> -> tensor<255x513xf32>
+  %3 = iree_encoding.set_encoding %2 : tensor<255x513xf32> -> tensor<255x513xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], original_type = tensor<255x513xf32>, user_indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>], round_dims_to = array<i64: 16, 16, 16>>>
+  flow.dispatch.tensor.store %3, %1, offsets = [0, 0], sizes = [255, 513], strides = [1, 1] : tensor<255x513xf32, #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], original_type = tensor<255x513xf32>, user_indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>], round_dims_to = array<i64: 16, 16, 16>>> -> !flow.dispatch.tensor<writeonly:tensor<255x513xf32,  #iree_encoding.encoding<operand_index = 0, op_type = matmul, element_types = [f32, f32, f32], original_type = tensor<255x513xf32>, user_indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>], round_dims_to = array<i64: 16, 16, 16>>>>
+  return
+}
+
+
+// CHECK-LABEL: func.func @set_encoding
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<16x129x16x4xf32>
+// CHECK: %[[PACK:.*]] = tensor.pack %2 padding_value(%cst : f32) outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [16, 4] into %[[EMPTY]] : tensor<255x513xf32> -> tensor<16x129x16x4xf32>
+// CHECK: %[[EXPAND:.*]] = tensor.expand_shape %[[PACK]]
+// CHECK: %[[TRANSPOSE:.*]] = linalg.transpose ins(%[[EXPAND]]
+
+// CHECK: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[TRANSPOSE]]
+// CHECK: %[[EXPAND2:.*]] = tensor.expand_shape %[[COLLAPSE]]
+// CHECK: flow.dispatch.tensor.store %[[EXPAND2]]