[AIROCMLIR-445] Lower linalg.generic convolution into rock (#2252)

The conversion primarily depends on the attribute in linalg.generic to extract the padding, dilation, stride, and op code. We transfer those attribute to `rock.conv`. The conversion works as the following: 

1. `isConv` function reads the stride, dilation, and padding from the op attributes and transform it into the equivalent rock.conv attribute
2. Padding (tensor.padding) is now removed since rock.conv handles that internally.
3. Layout is then transform into the layout expected by the `rock.conv`.
4. Emit attribute for `rock.conv`

For more context, this PR tries to match whatever we are currently doing in the migraphx -> tosa -> rock pipeline.

Author: Mr-Anyone

mlir/include/mlir/Dialect/Rock/IR/Rock.h

@@ -97,6 +97,10 @@ TransformMapAttr getTransformMapAttrChecked(
     llvm::function_ref<mlir::InFlightDiagnostic()> emitError,
     MLIRContext *context, ArrayRef<TransformAttr> ops, AffineMapAttr map,
     DenseI64ArrayAttr upperBounds, DenseI64ArrayAttr lowerBounds);
+
+// Attributes used in the linalg.generic operations so that linalg -> rock is
+// easier to match
+constexpr llvm::StringLiteral linalgConvOpAttrName = "conv_op";
 } // namespace rock
 } // namespace mlir
 #endif // MLIR_ROCKOPS_OPS_H_

mlir/lib/Conversion/LinalgToRock/LinalgToRock.cpp

@@ -14,6 +14,7 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Rock/IR/Rock.h"
+#include "mlir/Dialect/Rock/IR/TransformMapBuilder.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/PatternMatch.h"
@@ -189,7 +190,7 @@ LogicalResult MatmulConverter<LinalgMatOp>::matchAndRewrite(
   }
   MatmulContext context = maybeContext.value();
 
-  // TODO: handle split K attributes as well
+  // TODO: see (AIROCMLIR-696)
   // TODO: handle broadcasting for matrix A and B
   RankedTensorType outputType =
       cast<RankedTensorType>(op.getOutputs()[0].getType());
@@ -259,9 +260,273 @@ LogicalResult ExpandStrideConverter::matchAndRewrite(
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// ConvLinalgConverter: linalg.generic (conv) -> rock.conv
+//===----------------------------------------------------------------------===//
+namespace {
+struct ConvFields {
+  rock::LinalgConvType type;
+  int64_t spatialDim;
+  ArrayAttr padding, stride, dilation;
+  StringAttr perfConfig;
+};
+} // namespace
+
+static int64_t getSpatialDim(rock::LinalgConvType type) {
+  switch (type) {
+  case rock::LinalgConvType::Conv1dNgchGkch:
+    return 1;
+  case rock::LinalgConvType::Conv2dNgchwGkchw:
+    return 2;
+  case rock::LinalgConvType::Conv3dNgchwdGkchwd:
+    return 3;
+  }
+  llvm_unreachable("unknown LinalgConvType");
+}
+
+/// Set filter_layout, input_layout, and output_layout on a rock.conv op.
+/// Layouts match the linalg convention: GKC*, NGC*, NGK*.
+static void setConvLayoutAttrs(OpBuilder &builder, rock::ConvOp cop,
+                               int64_t spatialDim) {
+  auto *ctx = builder.getContext();
+  auto setLayout = [&](StringRef attrName, ArrayRef<StringRef> prefix,
+                       StringRef suffix) {
+    SmallVector<Attribute> layout;
+    for (StringRef dim : prefix)
+      layout.push_back(StringAttr::get(ctx, dim));
+    for (int64_t i = 0; i < spatialDim; ++i)
+      layout.push_back(StringAttr::get(ctx, Twine(i) + suffix));
+    cop->setAttr(attrName, builder.getArrayAttr(layout));
+  };
+
+  // The input layout in the operand of linalg.generic is NGC*, and
+  // the filter layout is GKC*. We have to transfer these attribute
+  // because later on in the pass, ConvToGemm expect them to be attached.
+  setLayout("filter_layout", {"g", "k", "c"}, "");
+  setLayout("input_layout", {"ni", "gi", "ci"}, "i");
+  setLayout("output_layout", {"no", "go", "ko"}, "o");
+}
+
+/// Remove the tensor.pad + tensor.expand_shape pattern emitted by
+/// migraphx-to-linalg, replacing it with just tensor.expand_shape on the
+/// unpadded source. rock.conv handles padding internally.
+///
+/// Expected IR structure:
+///   %padded = tensor.pad %original ...
+///   %expanded = tensor.expand_shape %padded ...
+/// Replaced with:
+///   %expanded = tensor.expand_shape %original ...
+static FailureOr<Value>
+removePaddingFromInput(ConversionPatternRewriter &rewriter,
+                       linalg::GenericOp op, Value in, ArrayAttr padding) {
+  bool hasPadding = llvm::any_of(padding.getValue(), [](Attribute attr) {
+    return cast<IntegerAttr>(attr).getInt() != 0;
+  });
+  if (!hasPadding)
+    return in;
+
+  auto expanded = in.getDefiningOp<tensor::ExpandShapeOp>();
+  auto padded = (expanded != nullptr)
+                    ? expanded->getOperand(0).getDefiningOp<tensor::PadOp>()
+                    : nullptr;
+  // We require padding here to have one use because the code structure emitted
+  // by the MIGraphX -> Linalg have one use. In theory, you don't need this
+  // check, but better be safe than sorry. This goes with expanded as well
+  if (!padded || !padded->hasOneUse()) {
+    op.emitError("unexpected padding code structure");
+    return failure();
+  }
+
+  if (!expanded || !expanded->hasOneUse()) {
+    return op.emitError("unexpected group expansion shape code structure");
+  }
+
+  SmallVector<int64_t> resultShape(expanded.getResultType().getShape());
+  // The tensor.pad operand has no group dimension: [N, G*C, spatial...].
+  // The expanded result has [N, G, C, spatial_padded...]. Take the first 3
+  // dims (N, G, C) from the expanded shape and append the unpadded spatial
+  // dims directly from the pad source starting at position 2.
+  auto padSourceShape =
+      cast<RankedTensorType>(padded.getOperand(0).getType()).getShape();
+  resultShape.resize(3);
+  resultShape.insert(resultShape.begin() + 3, padSourceShape.begin() + 2,
+                     padSourceShape.end());
+
+  RankedTensorType newResultType = RankedTensorType::get(
+      resultShape, padded.getResultType().getElementType());
+  Value result = tensor::ExpandShapeOp::create(
+      rewriter, expanded.getLoc(), newResultType, padded.getOperand(0),
+      expanded.getReassociationIndices());
+  // erase the operations as well
+  rewriter.eraseOp(expanded);
+  rewriter.eraseOp(padded);
+  return result;
+}
+
+namespace {
+struct ConvLinalgConverter final
+    : public OpConversionPattern<linalg::GenericOp> {
+  using OpConversionPattern<linalg::GenericOp>::OpConversionPattern;
+  using OpConversionPattern<linalg::GenericOp>::getTypeConverter;
+  using OpAdaptor = typename OpConversionPattern<linalg::GenericOp>::OpAdaptor;
+
+  LogicalResult
+  matchAndRewrite(linalg::GenericOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+
+private:
+  FailureOr<ConvFields> isConv(ConversionPatternRewriter &rewriter,
+                               linalg::GenericOp op) const;
+};
+} // namespace
+
+FailureOr<ConvFields>
+ConvLinalgConverter::isConv(ConversionPatternRewriter &rewriter,
+                            linalg::GenericOp op) const {
+  auto name =
+      op->getAttrOfType<rock::LinalgConvTypeAttr>(rock::linalgConvOpAttrName);
+  if (!name)
+    return failure();
+  rock::LinalgConvType convType = name.getValue();
+  int64_t spatialDim = getSpatialDim(convType);
+  // Conv1D is broadcasted into Conv2D. To check for error, we
+  // use effectiveDim instead because it one has more stride/dilation
+  // in the expanded dimension
+  int64_t effectiveDim = (spatialDim == 1) ? spatialDim + 1 : spatialDim;
+
+  auto convertToArrayAttr =
+      [&](Attribute arr, ArrayRef<int64_t> dimOneDefaults = {}) -> ArrayAttr {
+    if (!arr || !isa<ArrayAttr>(arr)) {
+      return ArrayAttr{};
+    }
+
+    SmallVector<int64_t, 4> values;
+    llvm::transform(
+        cast<ArrayAttr>(arr).getValue(), std::back_inserter(values),
+        [](Attribute val) { return cast<IntegerAttr>(val).getInt(); });
+    // Conv1D is expanded into Conv2D: append identity defaults for the
+    // extra spatial dimension (stride=1, dilation=1, pad=0).
+    if (spatialDim == 1)
+      values.insert(values.end(), dimOneDefaults.begin(), dimOneDefaults.end());
+    return rewriter.getIndexArrayAttr(values);
+  };
+
+  auto dilation =
+      convertToArrayAttr(op->getAttr("dilation"), /*dimOneDefaults=*/{1});
+  auto stride =
+      convertToArrayAttr(op->getAttr("stride"), /*dimOneDefaults=*/{1});
+  if (!dilation || static_cast<int64_t>(dilation.size()) != effectiveDim) {
+    op.emitError("invalid dilation");
+    return failure();
+  }
+
+  if (!stride || static_cast<int64_t>(stride.size()) != effectiveDim) {
+    op.emitError("invalid stride");
+    return failure();
+  }
+
+  // Input format:  [dim0_low, dim1_low, ..., dim0_high, dim1_high, ...]
+  // Rock  format:  [dim0_low, dim0_high, dim1_low, dim1_high, ...]
+  auto originalPadding = convertToArrayAttr(op->getAttr("pad"));
+  if (!originalPadding) {
+    op.emitError("no padding found");
+    return failure();
+  }
+  int64_t numSpatial = originalPadding.size() / 2;
+  SmallVector<Attribute, 8> interleavedPad;
+  for (int64_t i = 0; i < numSpatial; ++i) {
+    interleavedPad.push_back(originalPadding[i]);
+    interleavedPad.push_back(originalPadding[numSpatial + i]);
+  }
+  // Conv1D is expanded into Conv2D
+  if (spatialDim == 1) {
+    interleavedPad.push_back(rewriter.getIndexAttr(0));
+    interleavedPad.push_back(rewriter.getIndexAttr(0));
+  }
+  auto padding = rewriter.getArrayAttr(interleavedPad);
+  // note that Conv1D is expanded into Conv2D
+  if (effectiveDim * 2 != (int64_t)padding.size()) {
+    op.emitError("invalid number of padding");
+    return failure();
+  }
+
+  StringAttr perfConfig = op->getAttrOfType<StringAttr>("perf_config");
+  return ConvFields{convType, spatialDim, padding,
+                    stride,   dilation,   perfConfig};
+}
+
+LogicalResult ConvLinalgConverter::matchAndRewrite(
+    linalg::GenericOp op, OpAdaptor adaptor,
+    ConversionPatternRewriter &rewriter) const {
+  FailureOr<ConvFields> maybeConv = isConv(rewriter, op);
+  if (failed(maybeConv))
+    return failure();
+
+  ConvFields conv = *maybeConv;
+  Location loc = op.getLoc();
+
+  auto maybeInput =
+      removePaddingFromInput(rewriter, op, op.getOperand(0), conv.padding);
+  if (failed(maybeInput))
+    return failure();
+
+  Value input = *maybeInput;
+  Value filter = op.getOperand(1);
+
+  // Conv1D is expanded into Conv2D: unmerge the single spatial dim
+  // into (spatial, W=1) for filter and input.
+  int64_t effectiveSpatialDim = conv.spatialDim;
+  if (conv.spatialDim == 1) {
+    effectiveSpatialDim = 2;
+    auto filterShape = cast<RankedTensorType>(filter.getType()).getShape();
+    rock::BottomUpTMBuilder builder(rewriter, {"g", "k", "c", "0"}, filterShape,
+                                    loc);
+    builder.passThrough({"gf", "kf", "cf"}, {0, 1, 2}, {"g", "k", "c"});
+    builder.unmerge({"0f", "1f"}, {3, 4}, "0", {filterShape[3], 1});
+    filter = rock::TransformOp::create(rewriter, loc, filter, builder.get());
+
+    auto inputShape = cast<RankedTensorType>(input.getType()).getShape();
+    rock::BottomUpTMBuilder b(rewriter, {"n", "g", "c", "0"}, inputShape, loc);
+    b.passThrough({"nu", "gu", "cu"}, {0, 1, 2}, {"n", "g", "c"});
+    b.unmerge({"0u", "1u"}, {3, 4}, "0", {inputShape[3], 1});
+    input = rock::TransformOp::create(rewriter, loc, input, b.get());
+  }
+
+  RankedTensorType linalgResultType =
+      cast<RankedTensorType>(op.getResult(0).getType());
+  SmallVector<int64_t> rockShape(linalgResultType.getShape());
+  if (conv.spatialDim == 1)
+    rockShape.push_back(1);
+  RankedTensorType rockResultType =
+      RankedTensorType::get(rockShape, linalgResultType.getElementType());
+  Value output =
+      bufferization::AllocTensorOp::create(rewriter, loc, rockResultType, {});
+  auto cop = rock::ConvOp::create(rewriter, loc, rockResultType, filter, input,
+                                  output, /*features=*/nullptr,
+                                  /*blockSize=*/nullptr, /*gridSize=*/nullptr,
+                                  conv.padding, conv.stride, conv.dilation,
+                                  /*params=*/nullptr);
+  // TODO: add splitk see (AIROCMLIR-696)
+  if (conv.perfConfig)
+    cop->setAttr("perf_config", conv.perfConfig);
+  setConvLayoutAttrs(rewriter, cop, effectiveSpatialDim);
+
+  Value result = cop.getResult();
+  if (conv.spatialDim == 1) {
+    auto shape = cast<RankedTensorType>(result.getType()).getShape();
+    rock::BottomUpTMBuilder b(rewriter, {"n", "g", "k", "0", "1"}, shape, loc);
+    b.passThrough({"no", "go", "ko"}, {0, 1, 2}, {"n", "g", "k"});
+    b.merge("0o", 3, {"0", "1"});
+    result = rock::TransformOp::create(rewriter, loc, result, b.get());
+  }
+
+  rewriter.replaceOp(op, result);
+  return success();
+}
+
 void mlir::rock::populateLinalgToRockConversionPattern(
     RewritePatternSet &pattern, MLIRContext *context) {
   pattern.add<MatmulConverter<linalg::BatchMatmulOp>,
               MatmulConverter<linalg::MatmulOp>, ExpandStrideConverter,
-              MatmulConverter<linalg::GenericOp>>(context);
+              MatmulConverter<linalg::GenericOp>, ConvLinalgConverter>(context);
 }

mlir/lib/Conversion/LinalgToRock/LinalgToRockPass.cpp

@@ -59,8 +59,15 @@ static void populateLinalgToRockDialectConversion(ConversionTarget &target) {
           return false;
         }
 
-        return linalg::isElementwise(linalgOp) || isa<linalg::GenericOp>(op) ||
-               isa<linalg::YieldOp>(op);
+        // Convolution linalg.generic has reduction iteration type. It is not
+        // a legal operation in that case
+        linalg::GenericOp castedOp = dyn_cast<linalg::GenericOp>(op);
+        if (castedOp && castedOp->hasAttr(rock::linalgConvOpAttrName)) {
+          return false;
+        }
+
+        return linalg::isElementwise(linalgOp) || isa<linalg::YieldOp>(op) ||
+               castedOp;
       });
 }