[MLIR][Linalg] Remove matmul_transpose variants (#147961)

Removes the `(batch_)matmul_transpose_{a|b}` variants from OpDSL and
replace it with `matmul affine_maps [...]` whenever appropriate. This is
in line with the
[plan](https://discourse.llvm.org/t/rfc-op-explosion-in-linalg/82863),
and can be done since #104783 merged.

See:
https://discourse.llvm.org/t/deprecate-batch-matmul-transpose-a-b-linalg-operations/87245

Issues investigated:
* pad transform tests that could use `matmul` instead, so change to
that.
* ArmSME test using transpose actually needed it, so changed to `matmul`
+ affine maps.

Arm tests validated by @banach-space (thanks!!).

Author: rengolin

mlir/include/mlir/Dialect/Linalg/IR/Linalg.h

@@ -145,8 +145,7 @@ std::pair<int64_t, int64_t> getFmrFromWinogradConv2DFmr(WinogradConv2DFmr fmr);
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Linalg/IR/LinalgRelayoutOps.h.inc"
 
-namespace mlir {
-namespace linalg {
+namespace mlir::linalg {
 
 /// Returns the outer shape in the packed domain before applying the
 /// transposition.
@@ -155,7 +154,194 @@ template <typename OpTy,
                                       std::is_same_v<OpTy, linalg::UnPackOp>>>
 SmallVector<int64_t> getPackedOuterShapeWithoutTransposition(OpTy packOrUnPack);
 
-} // namespace linalg
-} // namespace mlir
+/// Specialization of `linalg.matmul` op that has a transpose map on A
+class MatmulTransposeAOp : public MatmulOp {
+  /// Create an affine map for a transpose-A matmul. Used only in the builders.
+  static SmallVector<AffineMap> getDefaultIndexingMaps(OpBuilder &builder);
+
+public:
+  using MatmulOp::MatmulOp;
+  static ::mlir::TypeID resolveTypeID() { return TypeID::get<MatmulOp>(); }
+
+  /// Build a transpose A matmul.
+  static void build(OpBuilder &builder, OperationState &result,
+                    ValueRange inputs, ValueRange outputs,
+                    ArrayRef<NamedAttribute> attributes = {});
+
+  static MatmulTransposeAOp create(OpBuilder &builder, Location location,
+                                   ValueRange inputs, ValueRange outputs,
+                                   ArrayRef<NamedAttribute> attributes = {});
+
+  /// Build a transpose A matmul with a specific result type.
+  static void build(OpBuilder &builder, OperationState &result,
+                    TypeRange resultTensorTypes, ValueRange inputs,
+                    ValueRange outputs,
+                    ArrayRef<NamedAttribute> attributes = {});
+
+  static MatmulTransposeAOp create(OpBuilder &builder, Location location,
+                                   TypeRange resultTensorTypes,
+                                   ValueRange inputs, ValueRange outputs,
+                                   ArrayRef<NamedAttribute> attributes = {});
+
+  /// Build a transpose A matmul with a specific result type and a cast type.
+  static void build(OpBuilder &builder, OperationState &result,
+                    TypeRange resultTensorTypes, ValueRange inputs,
+                    ValueRange outputs, Attribute cast,
+                    ArrayRef<NamedAttribute> attributes = {});
+
+  static MatmulTransposeAOp create(OpBuilder &builder, Location location,
+                                   TypeRange resultTensorTypes,
+                                   ValueRange inputs, ValueRange outputs,
+                                   Attribute cast,
+                                   ArrayRef<NamedAttribute> attributes = {});
+
+  /// Checks if the affine map is the expected one for this operation
+  static bool isDefaultIndexingMaps(Attribute attr);
+
+  static bool classof(Operation *op);
+};
+
+/// Specialization of `linalg.matmul` op that has a transpose map on B
+class MatmulTransposeBOp : public MatmulOp {
+  /// Create an affine map for a transpose-B matmul. Used only in the builders.
+  static SmallVector<AffineMap> getDefaultIndexingMaps(OpBuilder &builder);
+
+public:
+  using MatmulOp::MatmulOp;
+  static ::mlir::TypeID resolveTypeID() { return TypeID::get<MatmulOp>(); }
+
+  /// Build a transpose B matmul.
+  static void build(OpBuilder &builder, OperationState &result,
+                    ValueRange inputs, ValueRange outputs,
+                    ArrayRef<NamedAttribute> attributes = {});
+
+  static MatmulTransposeBOp create(OpBuilder &builder, Location location,
+                                   ValueRange inputs, ValueRange outputs,
+                                   ArrayRef<NamedAttribute> attributes = {});
+
+  /// Build a transpose B matmul with a specific result type.
+  static void build(OpBuilder &builder, OperationState &result,
+                    TypeRange resultTensorTypes, ValueRange inputs,
+                    ValueRange outputs,
+                    ArrayRef<NamedAttribute> attributes = {});
+
+  static MatmulTransposeBOp create(OpBuilder &builder, Location location,
+                                   TypeRange resultTensorTypes,
+                                   ValueRange inputs, ValueRange outputs,
+                                   ArrayRef<NamedAttribute> attributes = {});
+
+  /// Build a transpose B matmul with a specific result type and a cast type.
+  static void build(OpBuilder &builder, OperationState &result,
+                    TypeRange resultTensorTypes, ValueRange inputs,
+                    ValueRange outputs, Attribute cast,
+                    ArrayRef<NamedAttribute> attributes = {});
+
+  static MatmulTransposeBOp create(OpBuilder &builder, Location location,
+                                   TypeRange resultTensorTypes,
+                                   ValueRange inputs, ValueRange outputs,
+                                   Attribute cast,
+                                   ArrayRef<NamedAttribute> attributes = {});
+
+  /// Checks if the affine map is the expected one for this operation
+  static bool isDefaultIndexingMaps(Attribute attr);
+
+  static bool classof(Operation *op);
+};
+
+/// Specialization of `linalg.batch_matmul` op that has a transpose map on A
+class BatchMatmulTransposeAOp : public BatchMatmulOp {
+  /// Create an affine map for a transpose-A batch_matmul. Used only in the
+  /// builders.
+  static SmallVector<AffineMap> getDefaultIndexingMaps(OpBuilder &builder);
+
+public:
+  using BatchMatmulOp::BatchMatmulOp;
+  static ::mlir::TypeID resolveTypeID() { return TypeID::get<BatchMatmulOp>(); }
+
+  /// Build a transpose A matmul.
+  static void build(OpBuilder &builder, OperationState &result,
+                    ValueRange inputs, ValueRange outputs,
+                    ArrayRef<NamedAttribute> attributes = {});
+
+  static BatchMatmulTransposeAOp
+  create(OpBuilder &builder, Location location, ValueRange inputs,
+         ValueRange outputs, ArrayRef<NamedAttribute> attributes = {});
+
+  /// Build a transpose A matmul with a specific result type.
+  static void build(OpBuilder &builder, OperationState &result,
+                    TypeRange resultTensorTypes, ValueRange inputs,
+                    ValueRange outputs,
+                    ArrayRef<NamedAttribute> attributes = {});
+
+  static BatchMatmulTransposeAOp
+  create(OpBuilder &builder, Location location, TypeRange resultTensorTypes,
+         ValueRange inputs, ValueRange outputs,
+         ArrayRef<NamedAttribute> attributes = {});
+
+  /// Build a transpose A matmul with a specific result type and a cast type.
+  static void build(OpBuilder &builder, OperationState &result,
+                    TypeRange resultTensorTypes, ValueRange inputs,
+                    ValueRange outputs, Attribute cast,
+                    ArrayRef<NamedAttribute> attributes = {});
+
+  static BatchMatmulTransposeAOp
+  create(OpBuilder &builder, Location location, TypeRange resultTensorTypes,
+         ValueRange inputs, ValueRange outputs, Attribute cast,
+         ArrayRef<NamedAttribute> attributes = {});
+
+  /// Checks if the affine map is the expected one for this operation
+  static bool isDefaultIndexingMaps(Attribute attr);
+
+  static bool classof(Operation *op);
+};
+
+/// Specialization of `linalg.batch_matmul` op that has a transpose map on B
+class BatchMatmulTransposeBOp : public BatchMatmulOp {
+  /// Create an affine map for a transpose-B batch_matmul. Used only in the
+  /// builders.
+  static SmallVector<AffineMap> getDefaultIndexingMaps(OpBuilder &builder);
+
+public:
+  using BatchMatmulOp::BatchMatmulOp;
+  static ::mlir::TypeID resolveTypeID() { return TypeID::get<BatchMatmulOp>(); }
+
+  /// Build a transpose B matmul.
+  static void build(OpBuilder &builder, OperationState &result,
+                    ValueRange inputs, ValueRange outputs,
+                    ArrayRef<NamedAttribute> attributes = {});
+
+  static BatchMatmulTransposeBOp
+  create(OpBuilder &builder, Location location, ValueRange inputs,
+         ValueRange outputs, ArrayRef<NamedAttribute> attributes = {});
+
+  /// Build a transpose B matmul with a specific result type.
+  static void build(OpBuilder &builder, OperationState &result,
+                    TypeRange resultTensorTypes, ValueRange inputs,
+                    ValueRange outputs,
+                    ArrayRef<NamedAttribute> attributes = {});
+
+  static BatchMatmulTransposeBOp
+  create(OpBuilder &builder, Location location, TypeRange resultTensorTypes,
+         ValueRange inputs, ValueRange outputs,
+         ArrayRef<NamedAttribute> attributes = {});
+
+  /// Build a transpose B matmul with a specific result type and a cast type.
+  static void build(OpBuilder &builder, OperationState &result,
+                    TypeRange resultTensorTypes, ValueRange inputs,
+                    ValueRange outputs, Attribute cast,
+                    ArrayRef<NamedAttribute> attributes = {});
+
+  static BatchMatmulTransposeBOp
+  create(OpBuilder &builder, Location location, TypeRange resultTensorTypes,
+         ValueRange inputs, ValueRange outputs, Attribute cast,
+         ArrayRef<NamedAttribute> attributes = {});
+
+  /// Checks if the affine map is the expected one for this operation
+  static bool isDefaultIndexingMaps(Attribute attr);
+
+  static bool classof(Operation *op);
+};
+
+} // namespace mlir::linalg
 
 #endif // MLIR_DIALECT_LINALG_IR_LINALG_H