[mlir][python] Fix MemRefType IsAFunction in Python bindings

MemRefType was using a wrong `isa` function in the bindings code, which could lead to invalid IR being constructed. Also run the verifier in memref dialect tests. Reviewed By: nicolasvasilache Differential Revision: https://reviews.llvm.org/D111784
llvm · Oct 14, 2021 · a04c0b7 · a04c0b7
1 parent e3e1da2
commit a04c0b7
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 2 deletions.
diff --git a/mlir/lib/Bindings/Python/IRTypes.cpp b/mlir/lib/Bindings/Python/IRTypes.cpp
@@ -406,7 +406,7 @@ class PyMemRefLayoutMapList;
 /// Ranked MemRef Type subclass - MemRefType.
 class PyMemRefType : public PyConcreteType<PyMemRefType, PyShapedType> {
 public:
-  static constexpr IsAFunctionTy isaFunction = mlirTypeIsARankedTensor;
+  static constexpr IsAFunctionTy isaFunction = mlirTypeIsAMemRef;
   static constexpr const char *pyClassName = "MemRefType";
   using PyConcreteType::PyConcreteType;
 

diff --git a/mlir/python/mlir/dialects/_memref_ops_ext.py b/mlir/python/mlir/dialects/_memref_ops_ext.py
@@ -33,5 +33,5 @@ def __init__(self,
     memref_resolved = _get_op_result_or_value(memref)
     indices_resolved = [] if indices is None else _get_op_results_or_values(
         indices)
-    return_type = memref_resolved.type
+    return_type = MemRefType(memref_resolved.type).element_type
     super().__init__(return_type, memref, indices_resolved, loc=loc, ip=ip)
diff --git a/mlir/test/python/dialects/memref.py b/mlir/test/python/dialects/memref.py
@@ -71,3 +71,4 @@ def testCustomBuidlers():
     # CHECK: func @f1(%[[ARG0:.*]]: memref<?x?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)
     # CHECK: memref.load %[[ARG0]][%[[ARG1]], %[[ARG2]]]
     print(module)
+    assert module.operation.verify()