[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