| # RUN: %PYTHON -m mlir.dialects.linalg.opdsl.dump_oplib --file %s | FileCheck %s | |
| from mlir.dialects.linalg.opdsl.lang import * | |
| # CHECK: --- | |
| # CHECK-LABEL: matmul | |
| # CHECK: implements: | |
| # CHECK-NEXT: - LinalgContractionOpInterface | |
| @linalg_structured_op | |
| def matmul( | |
| A=TensorDef(T, S.M, S.K), | |
| B=TensorDef(T, S.K, S.N), | |
| C=TensorDef(U, S.M, S.N, output=True)): | |
| implements(ContractionOpInterface) | |
| C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n]) |