Skip to content

[MLIR][Linalg] Remove matmul_transpose variants #147961

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 7 commits into
base: main
Choose a base branch
from
+628 −895
Open

[MLIR][Linalg] Remove matmul_transpose variants #147961

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
c18ca87
[MLIR][Linalg] Remove matmul_transpose variants
rengolin Jul 8, 2025
285305f
format
rengolin Jul 10, 2025
6a088d9
Specialization for transpose variants
rengolin Jul 11, 2025
9cf9125
format
rengolin Jul 11, 2025
0523a8d
re-add wonrgfuly removed matmul tests
rengolin Jul 14, 2025
32f3566
Check affine map on classof
rengolin Jul 14, 2025
1218998
Adding new create methods
rengolin Jul 24, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
192 changes: 192 additions & 0 deletions mlir/include/mlir/Dialect/Linalg/IR/Linalg.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,4 +144,196 @@ std::pair<int64_t, int64_t> getFmrFromWinogradConv2DFmr(WinogradConv2DFmr fmr);
#define GET_OP_CLASSES
#include "mlir/Dialect/Linalg/IR/LinalgRelayoutOps.h.inc"

namespace mlir::linalg {

/// 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> getAffineMaps(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 isExpectedAffineMaps(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> getAffineMaps(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 isExpectedAffineMaps(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> getAffineMaps(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 isExpectedAffineMaps(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> getAffineMaps(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 isExpectedAffineMaps(Attribute attr);

static bool classof(Operation *op);
};

} // namespace mlir::linalg

#endif // MLIR_DIALECT_LINALG_IR_LINALG_H
Loading
Loading