[Mlir-commits] [mlir] [mlir][CAPI][python] expose the python bindings for linalg::isaContractionOpInterface and linalg::inferContractionDims (PR #134935)

[Jakub Kuderski]

================
@@ -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]
+        body_block = func_op.regions[0].blocks[0]
+        fill_op = body_block.operations[1]
+        matmul_op = body_block.operations[2]
+
+        assert not linalg.isa_contraction_op(fill_op)
----------------
kuhar wrote:

> but there are already both in this file?

There are a few asserts towards the top of the file that deal with strict type checking AFAICT. The overwhelming majority of the code uses CHECK lines. I think for consistency, we should at minimum have a check that the function executed and then use CHECK lines for IR objects that can be printed and handle the rest via asserts (e.g., that the module parsed or that something is not `None`). Or alternatively, move this to a new file like `mlir/test/python/dialects/linalg/utils.py` where we test all the helper functions with asserts.

https://github.com/llvm/llvm-project/pull/134935