[Mlir-commits] [mlir] [mlir][CAPI][python] expose the python bindings for linalg::isaContractionOpInterface and linalg::inferContractionDims (PR #134935)
Maksim Levental
llvmlistbot at llvm.org
Tue Apr 8 16:47:16 PDT 2025
================
@@ -606,3 +606,38 @@ def tensor_pack(src, dst):
# CHECK: return %[[VAL_4]] : tensor<128x128xf32>
# CHECK: }
print(module)
+
+
+ at run
+def test_infer_contraction_dimensions():
+ with Context(), Location.unknown():
+ module = ir.Module.parse(
+ r"""
+ module {
+ func.func @matmul(%arg0: tensor<4x4xf32>, %arg1: tensor<4x4xf32>)
+ -> tensor<4x4xf32> {
+ %cst = arith.constant 0.0 : f32
+ %0 = linalg.fill ins(%cst : f32) outs(%arg0 : tensor<4x4xf32>) -> tensor<4x4xf32>
+ %1 = linalg.matmul ins(%arg0, %arg1 : tensor<4x4xf32>, tensor<4x4xf32>)
+ outs(%0 : tensor<4x4xf32>) -> tensor<4x4xf32>
+ return %1 : tensor<4x4xf32>
+ }
+ }
+ """
+ )
+ func_op = module.body.operations[0]
----------------
makslevental wrote:
you can find some examples at [mlir/test/python/dialects/linalg/ops.py](https://github.com/makslevental/llvm-project/blob/a72616de18c0814ad37b5748d6bdc60b825dd889/mlir/test/python/dialects/linalg/ops.py#L43-L47)
https://github.com/llvm/llvm-project/pull/134935
More information about the Mlir-commits
mailing list