Add general convolution TTIR op. Add pass that legalizes it to Conv2D

include/ttmlir/Dialect/TTIR/IR/CMakeLists.txt

@@ -2,6 +2,12 @@ add_mlir_dialect(TTIROps ttir)
 add_mlir_doc(TTIRBase TTIRDialect src/autogen/md/Dialect/ -gen-dialect-doc)
 add_mlir_doc(TTIROps TTIROp src/autogen/md/Dialect/ -gen-op-doc)
 
+set(LLVM_TARGET_DEFINITIONS TTIROpsAttrs.td)
+mlir_tablegen(TTIROpsAttrs.h.inc -gen-attrdef-decls)
+mlir_tablegen(TTIROpsAttrs.cpp.inc -gen-attrdef-defs)
+add_public_tablegen_target(TTIROpsAttrsIncGen)
+add_dependencies(mlir-headers TTIROpsAttrsIncGen)
+
 set(LLVM_TARGET_DEFINITIONS TTIROpsInterfaces.td)
 mlir_tablegen(TTIROpsInterfaces.h.inc -gen-op-interface-decls)
 mlir_tablegen(TTIROpsInterfaces.cpp.inc -gen-op-interface-defs)

include/ttmlir/Dialect/TTIR/IR/TTIRBase.td

@@ -21,7 +21,7 @@ def TTIR_Dialect : Dialect {
         or dialects that are actually supported by a consuming backend.
     }];
     let cppNamespace = "::mlir::tt::ttir";
-
+    let useDefaultAttributePrinterParser = 1;
     let dependentDialects = [
       "::mlir::arith::ArithDialect",
       "::mlir::func::FuncDialect",

include/ttmlir/Dialect/TTIR/IR/TTIROps.h

@@ -16,6 +16,9 @@
 
 #include "TTIROpsInterfaces.h"
 
+#define GET_ATTRDEF_CLASSES
+#include "ttmlir/Dialect/TTIR/IR/TTIROpsAttrs.h.inc"
+
 #define GET_OP_CLASSES
 #include "ttmlir/Dialect/TTIR/IR/TTIROps.h.inc"
 

include/ttmlir/Dialect/TTIR/IR/TTIROps.td

@@ -7,6 +7,7 @@
 
 include "ttmlir/Dialect/TT/IR/TTOpsTypes.td"
 include "ttmlir/Dialect/TTIR/IR/TTIRBase.td"
+include "ttmlir/Dialect/TTIR/IR/TTIROpsAttrs.td"
 include "ttmlir/Dialect/TTIR/IR/TTIROpsInterfaces.td"
 include "mlir/Dialect/Linalg/IR/LinalgBase.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
@@ -557,6 +558,45 @@ def TTIR_Conv2dOp : TTIR_DPSOp<"conv2d"> {
     let hasVerifier = 1;
 }
 
+def TTIR_ConvolutionOp : TTIR_DPSOp<"convolution"> {
+  let summary = "Generalized convolution op.";
+  let description = [{
+    Applies a convolution of the rhs with the lhs.
+
+    This operation captures convolutions of all dimensionality as well
+    as deconvolution/conv transpose.
+  }];
+
+  let arguments = (ins
+    AnyRankedTensor:$input,
+    AnyRankedTensor:$weight,
+    Optional<AnyRankedTensor>:$bias,
+    AnyRankedTensor:$output,
+    // Default value: one for each of the spatial dimension.
+    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "SmallVector<int64_t>(getConvolutionLayout().getInputSpatialDimensions().size(), 1)">:$window_strides,
+    // Default value: two zeros for each of the spatial dimension.
+    DefaultValuedOptionalAttr<I64ElementsAttr, "SmallVector<int64_t>(getConvolutionLayout().getInputSpatialDimensions().size()*2, 0)">:$padding,
+    // Default value: one for each of the spatial dimension.
+    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "SmallVector<int64_t>(getConvolutionLayout().getInputSpatialDimensions().size(), 1)">:$input_dilation,
+    // Default value: one for each of the spatial dimension.
+    DefaultValuedOptionalAttr<DenseI64ArrayAttr, "SmallVector<int64_t>(getConvolutionLayout().getInputSpatialDimensions().size(), 1)">:$weight_dilation,
+    // Default value: false for each of the spatial dimension.
+    DefaultValuedOptionalAttr<DenseBoolArrayAttr, "SmallVector<bool>(getConvolutionLayout().getInputSpatialDimensions().size(), false)">:$window_reversal,
+    TTIR_ConvolutionLayoutAttr:$convolution_layout,
+    ConfinedAttr<I64Attr, [IntPositive]>:$feature_group_count,
+    ConfinedAttr<I64Attr, [IntPositive]>:$batch_group_count,
+    TT_OperandConstraintArrayAttr:$operand_constraints
+  );
+
+  let results = (outs AnyRankedTensor);
+  let hasVerifier = 1;
+
+  let extraClassDeclaration = [{
+      MutableOperandRange getDpsInitsMutable() { return getOutputMutable(); }
+    }];
+}
+
+
 def TTIR_MaxPool2dOp : TTIR_DPSOp<"max_pool2d"> {
     let summary = "Applies a 2D max pooling over an input signal composed of several input planes.";
     let description = [{

include/ttmlir/Dialect/TTIR/IR/TTIROpsAttrs.td

@@ -0,0 +1,44 @@
+// SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#ifndef TTMLIR_TTIR_ATTRS_TD
+#define TTMLIR_TTIR_ATTRS_TD
+
+include "mlir/IR/AttrTypeBase.td"
+include "ttmlir/Dialect/TTIR/IR/TTIRBase.td"
+
+def TTIR_ConvolutionLayoutAttr : AttrDef<TTIR_Dialect, "ConvolutionLayout", [], "::mlir::Attribute"> {
+  let mnemonic = "convolution_layout";
+  let summary = "Structure of dimension information for convolution op";
+  let description = [{
+    Holds the layout information for the input activation, weights, and output.
+  }];
+  let parameters = (ins
+    "int64_t":$inputBatchDimension,
+    "int64_t":$inputFeatureDimension,
+    ArrayRefParameter<"int64_t">:$inputSpatialDimensions,
+
+    "int64_t":$kernelOutputFeatureDimension,
+    "int64_t":$kernelInputFeatureDimension,
+    ArrayRefParameter<"int64_t">:$kernelSpatialDimensions,
+
+    "int64_t":$outputBatchDimension,
+    "int64_t":$outputFeatureDimension,
+    ArrayRefParameter<"int64_t">:$outputSpatialDimensions
+  );
+
+  let assemblyFormat = [{
+        `input_batch` `=` $inputBatchDimension `,`
+        `input_feature` `=` $inputFeatureDimension`,`
+        `input_spatial_dimensions` `=` custom<DimensionList>($inputSpatialDimensions) `,`
+        `kernel_output_feature` `=` $kernelOutputFeatureDimension `,`
+        `kernel_input_feature` `=` $kernelInputFeatureDimension `,`
+        `kernel_spatial_dimensions` `=` custom<DimensionList>($kernelSpatialDimensions) `,`
+        `output_batch` `=` $outputBatchDimension `,`
+        `output_feature` `=` $outputFeatureDimension `,`
+        `output_spatial_dimensions` `=` custom<DimensionList>($outputSpatialDimensions)
+  }];
+}
+
+#endif // TTMLIR_TTIR_ATTRS_TD

include/ttmlir/Dialect/TTIR/Transforms/Passes.td

@@ -119,4 +119,11 @@ def TTIRLoadSystemDesc: Pass<"ttir-load-system-desc", "::mlir::ModuleOp"> {
     ];
 }
 
+def TTIRConvolutionToConv2d: Pass<"ttir-convolution-to-conv2d", "::mlir::ModuleOp"> {
+  let summary = "Convert eligible convolution ops to conv2d ops.";
+  let description = [{
+    This pass converts eligible convolution ops to conv2d ops.
+  }];
+}
+
 #endif

lib/Conversion/StableHLOToTTIR/StableHLOToTTIRPatterns.cpp

@@ -3,6 +3,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 #include <algorithm>
+#include <type_traits>
 #include <vector>
 
 #include "mlir/IR/Builders.h"
@@ -373,46 +374,28 @@ class StableHLOToTTIRConvolutionOpConversionPattern
     tensor::EmptyOp outputTensor = rewriter.create<tensor::EmptyOp>(
         srcOp.getLoc(), outputType.getShape(), outputType.getElementType());
 
-    std::vector<int64_t> strides =
-        adaptor.getWindowStrides().value_or(ArrayRef<int64_t>({1, 1})).vec();
-    IntegerAttr stride_height_attr =
-        rewriter.getSI32IntegerAttr(static_cast<int32_t>(strides[0]));
-    IntegerAttr stride_width_attr =
-        rewriter.getSI32IntegerAttr(static_cast<int32_t>(strides[1]));
-
-    std::vector<int64_t> dilation =
-        adaptor.getLhsDilation().value_or(ArrayRef<int64_t>({1, 1})).vec();
-
-    IntegerAttr dilation_height_attr =
-        rewriter.getSI32IntegerAttr(static_cast<int32_t>(dilation[0]));
-    IntegerAttr dilation_width_attr =
-        rewriter.getSI32IntegerAttr(static_cast<int32_t>(dilation[1]));
-
-    IntegerAttr groups_attr = rewriter.getSI32IntegerAttr(
-        static_cast<int32_t>(adaptor.getFeatureGroupCount()));
-
-    std::vector<int32_t> padding;
-    if (!adaptor.getPadding().has_value()) {
-      padding = {0, 0, 0, 0};
-    } else {
-      for (auto iter = adaptor.getPadding()->value_begin<int64_t>();
-           iter < adaptor.getPadding()->value_end<int64_t>(); iter++) {
-        padding.push_back(static_cast<int32_t>(*iter));
-      }
-    }
-
-    rewriter.replaceOpWithNewOp<mlir::tt::ttir::Conv2dOp>(
+    auto dimNums = adaptor.getDimensionNumbers();
+    rewriter.replaceOpWithNewOp<mlir::tt::ttir::ConvolutionOp>(
         srcOp, outputType, adaptor.getLhs(), adaptor.getRhs(),
-        mlir::Value(nullptr), outputTensor, stride_height_attr,
-        stride_width_attr, dilation_height_attr, dilation_width_attr,
-        groups_attr, rewriter.getSI32IntegerAttr(padding[0]),
-        rewriter.getSI32IntegerAttr(padding[1]),
-        rewriter.getSI32IntegerAttr(padding[2]),
-        rewriter.getSI32IntegerAttr(padding[3]),
+        mlir::Value(nullptr), outputTensor, adaptor.getWindowStridesAttr(),
+        adaptor.getPaddingAttr(), adaptor.getLhsDilationAttr(),
+        adaptor.getRhsDilationAttr(), adaptor.getWindowReversalAttr(),
+        mlir::tt::ttir::ConvolutionLayoutAttr::get(
+            getContext(), dimNums.getInputBatchDimension(),
+            dimNums.getInputFeatureDimension(),
+            dimNums.getInputSpatialDimensions(),
+            dimNums.getKernelOutputFeatureDimension(),
+            dimNums.getKernelInputFeatureDimension(),
+            dimNums.getKernelSpatialDimensions(),
+            dimNums.getOutputBatchDimension(),
+            dimNums.getOutputFeatureDimension(),
+            dimNums.getOutputSpatialDimensions()),
+        adaptor.getFeatureGroupCountAttr(), adaptor.getBatchGroupCountAttr(),
         rewriter.getArrayAttr(
             SmallVector<Attribute>(adaptor.getOperands().size() + 1,
                                    rewriter.getAttr<OperandConstraintAttr>(
                                        OperandConstraint::AnyDeviceTile))));
+
     return success();
   }
 };

lib/Dialect/TTIR/IR/TTIRDialect.cpp

@@ -8,6 +8,11 @@
 #include "mlir/InitAllDialects.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "ttmlir/Dialect/TTIR/IR/TTIROps.h"
+#include "llvm/ADT/StringSet.h"
+#include "llvm/ADT/TypeSwitch.h"
+
+#define GET_ATTRDEF_CLASSES
+#include "ttmlir/Dialect/TTIR/IR/TTIROpsAttrs.cpp.inc"
 
 using namespace mlir;
 using namespace mlir::tt::ttir;
@@ -59,4 +64,8 @@ void TTIRDialect::initialize() {
 #include "ttmlir/Dialect/TTIR/IR/TTIROps.cpp.inc"
       >();
   addInterfaces<TTIRInlinerInterface>();
+  addAttributes<
+#define GET_ATTRDEF_LIST
+#include "ttmlir/Dialect/TTIR/IR/TTIROpsAttrs.cpp.inc"
+      >();
 }

lib/Dialect/TTIR/IR/TTIROps.cpp

@@ -62,6 +62,54 @@ ::mlir::LogicalResult mlir::tt::ttir::Conv2dOp::verify() {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// ConvolutionOp
+//===----------------------------------------------------------------------===//
+
+::mlir::LogicalResult mlir::tt::ttir::ConvolutionOp::verify() {
+  if (getConvolutionLayout().getInputSpatialDimensions().size() !=
+      getConvolutionLayout().getOutputSpatialDimensions().size()) {
+    return emitOpError("Convolution input, output, and kernel must have the "
+                       "same number of spatial dimensions");
+  }
+  if (getConvolutionLayout().getInputSpatialDimensions().size() !=
+      getConvolutionLayout().getKernelSpatialDimensions().size()) {
+    return emitOpError("Convolution input, output, and kernel must have the "
+                       "same number of spatial dimensions");
+  }
+
+  // Subtract 2 from the rank as to not count batch and feature dimension
+  if (getInput().getType().getRank() - 2 !=
+      (int64_t)getConvolutionLayout().getInputSpatialDimensions().size()) {
+    return emitOpError("Input tensor must have the same number of spatial "
+                       "dimensions as specified in the ConvolutionLayout");
+  }
+
+  if (getWeight().getType().getRank() - 2 !=
+      (int64_t)getConvolutionLayout().getKernelSpatialDimensions().size()) {
+    return emitOpError("Weight tensor must have the same number of spatial "
+                       "dimensions as specified in the ConvolutionLayout");
+  }
+
+  std::optional<::mlir::RankedTensorType> biasType =
+      getBias().getImpl() ? std::make_optional(getBias().getType())
+                          : std::nullopt;
+
+  if (biasType.has_value()) {
+    if (biasType->getRank() != 4) {
+      return emitOpError("Bias must be a 4D tensor");
+    }
+  }
+
+  if (getWindowStrides().size() !=
+      getConvolutionLayout().getInputSpatialDimensions().size()) {
+    return emitOpError("Window strides must have the same number of elements "
+                       "as the spatial dimensions of the input tensor");
+  }
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // MaxPool2dOp
 //===----------------------------------------------------------------------===//