Add backwards data convolution op to MIGraphX dialect (#1946)

* Add initial lowering from MIGraphX -> TOSA

* Initial changes to adding multiple kernels

* Partial implementation of TosaToRock lowering

* Remove old check disabling split k and bwd convs

* Remove TosaToRock conversions from this pass

* Remove unnecessary code changes

* Minor fix

* Add checks to LIT test

* Clang-format

* Fix Copilot review comments

* Add back splitK logic

* Update comment

* Remove additional line

* Attend to review comments

* Add new line to new LIT test

* Fix newline

* Update out_pad and add some more tests

* Remove bwd_weight

* Attend to minor review comments

* Change wording of op definition

* Clang format MIGraphX.td

---------

Co-authored-by: Umang Yadav <[email protected]>

Author: justinrosner

mlir/include/mlir/Dialect/MIGraphX/IR/MIGraphX.td

@@ -340,6 +340,18 @@ def MIGraphX_ConvolutionOp :
   }];
 }
 
+def MIGraphX_ConvolutionBwdDataOp
+    : MIGraphX_ConvOpBase<
+          "backwards_data_convolution", [F32, F16, BF16], [F32, F16, BF16]> {
+  let summary = "Backwards data convolution";
+  let description = [{
+    The `migraphx.backwards_data_convolution` op computes a transposed
+    convolution op which effectively reverses a standard convolution's
+    spatial transformation. It is an upsampling technique that increases
+    the height and width of an input feature map. 
+  }];
+}
+
 def MIGraphX_BatchNormOp :
     MIGraphX_Op<"batch_norm_inference">,
     Arguments<(ins AnyMIXRShaped:$input,

mlir/lib/Conversion/MIGraphXToTosa/MIGraphXToTosa.cpp

@@ -245,6 +245,7 @@ LogicalResult ConvConverter<ConvType>::matchAndRewrite(
   auto outputTy = cast<MIXRShapedType>(results[0].getType());
   Type outElementTy = outputTy.getElementType();
   Type newOutElementTy = getTypeConverter()->convertType(outElementTy);
+  bool isBwdDataConvOp = isa<migraphx::ConvolutionBwdDataOp>(op);
 
   if (outElementTy.isUnsignedInteger())
     return op.emitError("No support for unsigned convolution.\n");
@@ -270,8 +271,8 @@ LogicalResult ConvConverter<ConvType>::matchAndRewrite(
   newShape.push_back(outShape[1]);
   Type newOutTy = RankedTensorType::get(newShape, newOutElementTy);
 
-  // There is no tosa.conv1d, so instead we'll add a dummy x1 dimension
-  // to the input tensors, and make a tosa.conv2d.
+  // There is no tosa.conv1d or tosa.transpose_conv1d, so instead we'll add a
+  // dummy x1 dimension to the input tensors, and make a tosa.conv2d.
   auto expandTo2D = [&rewriter, loc](mlir::Value value) {
     ArrayRef<int64_t> origShape = cast<ShapedType>(value.getType()).getShape();
     SmallVector<int64_t> expShape(origShape.drop_back());
@@ -283,13 +284,14 @@ LogicalResult ConvConverter<ConvType>::matchAndRewrite(
     return reshaped;
   };
 
-  // Construct a new Conv2DOp.
+  // Construct a new Conv2DOp/TransposeConv2DOp.
   Operation *cop;
   Type new1DOutTy;
   Value inputZp, weightZp;
   switch (dims) {
   case 1:
-    // Expand to do a conv2d, because there's no conv1d op.
+    // Expand to do a conv2d/transpose_conv2d, because there's no 1d version of
+    // the ops.
     newShape.insert(std::prev(newShape.end()), 1);
     new1DOutTy = RankedTensorType::get(newShape, newOutElementTy);
     input = expandTo2D(input);
@@ -299,31 +301,57 @@ LogicalResult ConvConverter<ConvType>::matchAndRewrite(
     weightZp =
         tosa::createZeroPointTensor(rewriter, loc, filter.getType(), 0).value();
 
-    cop = rewriter.create<tosa::Conv2DOp>(
-        loc, new1DOutTy,
-        ValueRange{
-            input, filter,
-            getZeroTensor(loc, newOutElementTy,
-                          cast<ShapedType>(filter.getType()).getShape()[0],
-                          rewriter),
-            inputZp, weightZp});
+    if (isBwdDataConvOp) {
+      cop = rewriter.create<tosa::TransposeConv2DOp>(
+          loc, new1DOutTy,
+          ValueRange{
+              input, filter,
+              getZeroTensor(loc, newOutElementTy,
+                            cast<ShapedType>(filter.getType()).getShape()[0],
+                            rewriter),
+              inputZp, weightZp});
+    } else {
+      cop = rewriter.create<tosa::Conv2DOp>(
+          loc, new1DOutTy,
+          ValueRange{
+              input, filter,
+              getZeroTensor(loc, newOutElementTy,
+                            cast<ShapedType>(filter.getType()).getShape()[0],
+                            rewriter),
+              inputZp, weightZp});
+    }
     break;
 
   case 2:
     inputZp =
         tosa::createZeroPointTensor(rewriter, loc, input.getType(), 0).value();
     weightZp =
         tosa::createZeroPointTensor(rewriter, loc, filter.getType(), 0).value();
-    cop = rewriter.create<tosa::Conv2DOp>(
-        loc, newOutTy,
-        ValueRange{
-            input, filter,
-            getZeroTensor(loc, newOutElementTy,
-                          cast<ShapedType>(filter.getType()).getShape()[0],
-                          rewriter),
-            inputZp, weightZp});
+    if (isBwdDataConvOp) {
+      cop = rewriter.create<tosa::TransposeConv2DOp>(
+          loc, newOutTy,
+          ValueRange{
+              input, filter,
+              getZeroTensor(loc, newOutElementTy,
+                            cast<ShapedType>(filter.getType()).getShape()[0],
+                            rewriter),
+              inputZp, weightZp});
+    } else {
+      cop = rewriter.create<tosa::Conv2DOp>(
+          loc, newOutTy,
+          ValueRange{
+              input, filter,
+              getZeroTensor(loc, newOutElementTy,
+                            cast<ShapedType>(filter.getType()).getShape()[0],
+                            rewriter),
+              inputZp, weightZp});
+    }
     break;
   case 3:
+    if (isBwdDataConvOp)
+      return op->emitError("Only 1-D and 2-D backwards convolution ops are "
+                           "supported");
+
     inputZp =
         tosa::createZeroPointTensor(rewriter, loc, input.getType(), 0).value();
     weightZp =
@@ -361,8 +389,6 @@ LogicalResult ConvConverter<ConvType>::matchAndRewrite(
     dilations.push_back(dyn_cast<IntegerAttr>(dilationAttr[i]).getInt());
   }
 
-  int64_t group = op.getGroup();
-
   // Determine the accumulation type based on the output type.
   Type accType;
   if (isa<FloatType>(elementTy)) {
@@ -386,11 +412,31 @@ LogicalResult ConvConverter<ConvType>::matchAndRewrite(
     pads.push_back(0);
   }
 
+  // Set attributes common to both forwards and backwards conv
   cop->setAttr("dilation", rewriter.getDenseI64ArrayAttr(dilations));
   cop->setAttr("stride", rewriter.getDenseI64ArrayAttr(strides));
-  cop->setAttr("pad", rewriter.getDenseI64ArrayAttr(pads));
-  cop->setAttr("group", rewriter.getI64IntegerAttr(group));
   cop->setAttr("acc_type", TypeAttr::get(accType));
+  int64_t group = op.getGroup();
+  cop->setAttr("group", rewriter.getI64IntegerAttr(group));
+
+  // Set padding for forwards and backwards convolution. Note: the padding here
+  // applies to input padding (which transpose.conv2D does not inherently
+  // support). TransposeConv2D will still require an output pad attribute, so we
+  // can just set that to zeros
+  if (isBwdDataConvOp) {
+    SmallVector<int64_t> zeroPads(pads.size(), 0);
+    cop->setAttr("out_pad", rewriter.getDenseI64ArrayAttr(zeroPads));
+  }
+  cop->setAttr("pad", rewriter.getDenseI64ArrayAttr(pads));
+
+  // For both types of backwards convolution, we will be using
+  // tosa.transpose_conv2d, so we are going to add a conv_kind attribute so
+  // that we can distinguish between the two types in TosaToRock.
+  // TODO: We will need to add conv_kind = "bwd_weight" when we eventually
+  // add support for bwd_weight ops in MIGraphX.
+  if (isa<migraphx::ConvolutionBwdDataOp>(op)) {
+    cop->setAttr("conv_kind", rewriter.getStringAttr("bwd_data"));
+  }
 
   // Convert optional attributes
   if (auto attr = (*op).template getAttrOfType<StringAttr>("perf_config"))
@@ -1499,7 +1545,8 @@ LogicalResult MHALLaunchConverter::matchAndRewrite(
 
 void migraphx::populateMIGraphXToTosaConversionPatterns(
     RewritePatternSet &patterns, TypeConverter &typeConverter) {
-  patterns.add<ConvConverter<ConvolutionOp>, ConvConverter<QuantConvolutionOp>,
+  patterns.add<ConvConverter<ConvolutionBwdDataOp>,
+               ConvConverter<ConvolutionOp>, ConvConverter<QuantConvolutionOp>,
                DotConverter<DotOp>, DotConverter<QuantDotOp>,
                BroadcastConverter, MultiBroadcastConverter, TransposeConverter,
                ReshapeConverter, SliceConverter, ReduceMeanConverter,

mlir/test/Conversion/MIGraphXToTosa/migraphx-to-tosa-bwd-failure.mlir

@@ -0,0 +1,24 @@
+// This test checks that we emit an error when trying to convert 3-D
+// backwards convolution ops.
+
+// RUN: not rocmlir-opt -split-input-file --migraphx-to-tosa %s 2>&1 | FileCheck %s
+
+module {
+    func.func @bwd_data_conv_3d(
+        %arg0: !migraphx.shaped<1x16x4x4x4xf32, 1024x64x16x4x1>,
+        %arg1: !migraphx.shaped<16x16x1x1x1xf32, 16x1x1x1x1>,
+        %arg2: !migraphx.shaped<1x16x4x4x4xf32, 1024x64x16x4x1>
+    ) -> !migraphx.shaped<1x16x4x4x4xf32, 1024x64x16x4x1> {
+        // CHECK: Only 1-D and 2-D backwards convolution ops are supported
+        %0 = migraphx.backwards_data_convolution %arg1, %arg0 {
+            dilation = [1, 1, 1],
+            group = 1 : i64,
+            padding = [0, 0, 0, 0, 0, 0],
+            padding_mode = 0 : i64,
+            stride = [1, 1, 1],
+            kernelId = 0 : i64
+        } : <16x16x1x1x1xf32, 16x1x1x1x1>, <1x16x4x4x4xf32, 1024x64x16x4x1> -> <1x16x4x4x4xf32, 1024x64x16x4x1>
+        return %0 : !migraphx.shaped<1x16x4x4x4xf32, 1024x64x16x4x1>
+    }
+}
+

mlir/test/Conversion/MIGraphXToTosa/migraphx-to-tosa.mlir

@@ -206,6 +206,39 @@ func.func @quant_conv2d_float8(%arg0: !migraphx.shaped<1x16x4x4xf8E5M2, 256x16x4
   return %0 : !migraphx.shaped<1x16x4x4xf32, 256x16x4x1>
 }
 
+// CHECK-LABEL: @bwd_data_conv
+func.func @bwd_data_conv(%arg0: !migraphx.shaped<1x16x4x4xf32, 256x16x4x1>, %arg1: !migraphx.shaped<16x16x1x1xf32, 16x1x1x1>, %arg2: !migraphx.shaped<1x16x4x4xf32, 256x16x4x1>) -> !migraphx.shaped<1x16x4x4xf32, 256x16x4x1> {
+  // CHECK: tosa.transpose_conv2d
+  // CHECK-SAME: {acc_type = f32, conv_kind = "bwd_data", dilation = array<i64: 1, 1>, group = 1 : i64, out_pad = array<i64: 0, 0, 0, 0>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} : (tensor<16x1x1x16xf32>, tensor<1x4x4x16xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x4x4x16xf32> 
+  %0 = migraphx.backwards_data_convolution %arg1, %arg0 {dilation = [1, 1], group = 1 : i64, padding = [0, 0, 0, 0], padding_mode = 0 : i64, stride = [1, 1], kernelId = 0 : i64} : <16x16x1x1xf32, 16x1x1x1>, <1x16x4x4xf32, 256x16x4x1> -> <1x16x4x4xf32, 256x16x4x1>  
+  return %0 : !migraphx.shaped<1x16x4x4xf32, 256x16x4x1>  
+}
+
+// CHECK-LABEL: @bwd_data_conv_attributes
+func.func @bwd_data_conv_attributes(%arg0: !migraphx.shaped<1x16x4x4xf32, 256x16x4x1>, %arg1: !migraphx.shaped<16x16x1x1xf32, 16x1x1x1>, %arg2: !migraphx.shaped<1x16x4x4xf32, 256x16x4x1>) -> !migraphx.shaped<1x16x4x4xf32, 256x16x4x1> {
+  // CHECK: tosa.transpose_conv2d
+  // CHECK-SAME: {acc_type = f32, conv_kind = "bwd_data", dilation = array<i64: 2, 2>, group = 2 : i64, out_pad = array<i64: 0, 0, 0, 0>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} : (tensor<16x1x1x16xf32>, tensor<1x4x4x16xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x4x4x16xf32> 
+  %0 = migraphx.backwards_data_convolution %arg1, %arg0 {dilation = [2, 2], group = 2 : i64, padding = [0, 0, 0, 0], padding_mode = 0 : i64, stride = [1, 1], kernelId = 0 : i64} : <16x16x1x1xf32, 16x1x1x1>, <1x16x4x4xf32, 256x16x4x1> -> <1x16x4x4xf32, 256x16x4x1>  
+  return %0 : !migraphx.shaped<1x16x4x4xf32, 256x16x4x1>  
+}
+
+// CHECK-LABEL: @bwd_data_conv_stride
+func.func @bwd_data_conv_stride(%arg0: !migraphx.shaped<1x32x3x3xf32, 288x9x3x1>, %arg1: !migraphx.shaped<32x16x4x4xf32, 256x16x4x1>, %arg2: !migraphx.shaped<1x32x9x9xf32, 2592x81x9x1>) -> !migraphx.shaped<1x32x9x9xf32, 2592x81x9x1> {
+  // CHECK: tosa.transpose_conv2d
+  // CHECK-SAME: {acc_type = f32, conv_kind = "bwd_data", dilation = array<i64: 1, 1>, group = 1 : i64, out_pad = array<i64: 0, 0, 0, 0>, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 2, 2>}
+  %0 = migraphx.backwards_data_convolution %arg1, %arg0 {dilation = [1, 1], group = 1 : i64, padding = [0, 0, 0, 0], padding_mode = 0 : i64, stride = [2, 2], kernelId = 0 : i64} : <32x16x4x4xf32, 256x16x4x1>, <1x32x3x3xf32, 288x9x3x1> -> <1x32x9x9xf32, 2592x81x9x1>
+  return %0 : !migraphx.shaped<1x32x9x9xf32, 2592x81x9x1>
+}
+
+// CHECK-LABEL: @bwd_data_conv1d
+func.func @bwd_data_conv1d(%arg0: !migraphx.shaped<1x64x224xf32, 0x1x0>, %arg1: !migraphx.shaped<1x3x224xf32, 672x224x1>, %arg2: !migraphx.shaped<64x3x1xf32, 3x1x1>) -> !migraphx.shaped<1x64x224xf32, 14336x224x1> {
+  // CHECK: tosa.transpose_conv2d
+  // CHECK-SAME: {acc_type = f32, conv_kind = "bwd_data", dilation = array<i64: 1, 1>, group = 1 : i64, out_pad = array<i64: 0, 0, 0, 0>, pad = array<i64: 3, 3, 0, 0>, stride = array<i64: 1, 1>} : (tensor<1x224x1x3xf32>, tensor<64x1x1x3xf32>, tensor<64xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x224x1x64xf32>
+  %0 = migraphx.backwards_data_convolution %arg1, %arg2 {dilation = [1], group = 1 : i64, padding = [3, 3], padding_mode = 0 : i64, stride = [1]} : <1x3x224xf32, 672x224x1>, <64x3x1xf32, 3x1x1> -> <1x64x224xf32, 14336x224x1>
+  %1 = migraphx.add %0, %arg0 : <1x64x224xf32, 14336x224x1>, <1x64x224xf32, 0x1x0> -> <1x64x224xf32, 14336x224x1>
+  return %1 : !migraphx.shaped<1x64x224xf32, 14336x224x1>
+}
+
 // -----
 
 // CHECK-LABEL: @dot_f16