[mlir][python] Swap shape and element_type order for MemRefType.

[stellaraccident]

Summary

• Matches how all of the other shaped types are declared.
• No super principled reason fro this ordering beyond that it makes the one that was different be like the rest.
• Also matches ordering of things like ndarray, et al.

[ftynse]

Matches how all of the other shaped types are declared.

Is it also the case in C++? Maybe we should fix that too.

[ftynse]

Hmm, this change actually makes Python API align on C++ API. I don't remember if there was a good reason to swap the order in C.

[stellaraccident]

Rebase and split from parent.

[stellaraccident]

In D94812#2504416, @ftynse wrote:

Hmm, this change actually makes Python API align on C++ API. I don't remember if there was a good reason to swap the order in C.

I suspect there was no reason. Just done by different people and missed in isolated reviews.

[mehdi_amini]

LGTM

mlir/examples/python/linalg_matmul.py

def FuncOp(name: str, func_type: Type) -> Tuple[Operation, Block]:
   }
   op = Operation.create("func", regions=1, attributes=attrs)
   body_region = op.regions[0]
   entry_block = body_region.blocks.append(*func_type.inputs)
   return op, entry_block
 
 
 def build_matmul_buffers_func(func_name, m, k, n, dtype):
-  lhs_type = MemRefType.get(dtype, [m, k])
-  rhs_type = MemRefType.get(dtype, [k, n])
-  result_type = MemRefType.get(dtype, [m, n])
+  lhs_type = MemRefType.get([m, k], dtype)
+  rhs_type = MemRefType.get([k, n], dtype)
+  result_type = MemRefType.get([m, n], dtype)
   # TODO: There should be a one-liner for this.
   func_type = FunctionType.get([lhs_type, rhs_type, result_type], [])
   _, entry = FuncOp(func_name, func_type)
   lhs, rhs, result = entry.arguments
   with InsertionPoint(entry):
     op = linalg.MatmulOp([lhs, rhs], [result])
     # TODO: Implement support for SingleBlockImplicitTerminator
     block = op.regions[0].blocks.append()
     with InsertionPoint(block):
         linalg.YieldOp(values=[])
 
     std.ReturnOp([])
 
 
 def build_matmul_tensors_func(func_name, m, k, n, dtype):
-  # TODO: MemRefType and TensorTypes should not have inverted dtype/shapes
-  # from each other.
   lhs_type = RankedTensorType.get([m, k], dtype)
   rhs_type = RankedTensorType.get([k, n], dtype)
   result_type = RankedTensorType.get([m, n], dtype)
   # TODO: There should be a one-liner for this.
   func_type = FunctionType.get([lhs_type, rhs_type], [result_type])
   _, entry = FuncOp(func_name, func_type)
   lhs, rhs = entry.arguments
   with InsertionPoint(entry):

mlir/lib/Bindings/Python/IRModules.cpp

public:
   static constexpr const char *pyClassName = "MemRefType";
   using PyConcreteType::PyConcreteType;
 
   PyMemRefLayoutMapList getLayout();
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
          "get",
-         [](PyType &elementType, std::vector<int64_t> shape,
+         [](std::vector<int64_t> shape, PyType &elementType,
             std::vector<PyAffineMap> layout, unsigned memorySpace,
             DefaultingPyLocation loc) {
            SmallVector<MlirAffineMap> maps;
            maps.reserve(layout.size());
            for (PyAffineMap &map : layout)
              maps.push_back(map);
 
            MlirType t = mlirMemRefTypeGetChecked(elementType, shape.size(),
                                                  shape.data(), maps.size(),
                                                  maps.data(), memorySpace, loc);
            // TODO: Rework error reporting once diagnostic engine is exposed
            // in C API.
            if (mlirTypeIsNull(t)) {
              throw SetPyError(
                  PyExc_ValueError,
                  Twine("invalid '") +
                      py::repr(py::cast(elementType)).cast<std::string>() +
                      "' and expected floating point, integer, vector or "
                      "complex "
                      "type.");
            }
            return PyMemRefType(elementType.getContext(), t);
          },
-         py::arg("element_type"), py::arg("shape"),
+         py::arg("shape"), py::arg("element_type"),
          py::arg("layout") = py::list(), py::arg("memory_space") = 0,
          py::arg("loc") = py::none(), "Create a memref type")
         .def_property_readonly("layout", &PyMemRefType::getLayout,
                                "The list of layout maps of the MemRef type.")
         .def_property_readonly(
             "memory_space",
             [](PyMemRefType &self) -> unsigned {
               return mlirMemRefTypeGetMemorySpace(self);

mlir/test/Bindings/Python/ir_types.py

 run(testUnrankedTensorType)
 
 # CHECK-LABEL: TEST: testMemRefType
 def testMemRefType():
   with Context(), Location.unknown():
     f32 = F32Type.get()
     shape = [2, 3]
     loc = Location.unknown()
-    memref = MemRefType.get(f32, shape, memory_space=2)
+    memref = MemRefType.get(shape, f32, memory_space=2)
     # CHECK: memref type: memref<2x3xf32, 2>
     print("memref type:", memref)
     # CHECK: number of affine layout maps: 0
     print("number of affine layout maps:", len(memref.layout))
     # CHECK: memory space: 2
     print("memory space:", memref.memory_space)
 
     layout = AffineMap.get_permutation([1, 0])
-    memref_layout = MemRefType.get(f32, shape, [layout])
+    memref_layout = MemRefType.get(shape, f32, [layout])
     # CHECK: memref type: memref<2x3xf32, affine_map<(d0, d1) -> (d1, d0)>>
     print("memref type:", memref_layout)
     assert len(memref_layout.layout) == 1
     # CHECK: memref layout: (d0, d1) -> (d1, d0)
     print("memref layout:", memref_layout.layout[0])
     # CHECK: memory space: 0
     print("memory space:", memref_layout.memory_space)
 
     none = NoneType.get()
     try:
-      memref_invalid = MemRefType.get(none, shape)
+      memref_invalid = MemRefType.get(shape, none)
     except ValueError as e:
       # CHECK: invalid 'Type(none)' and expected floating point, integer, vector
       # CHECK: or complex type.
       print(e)
     else:
       print("Exception not produced")
 
 run(testMemRefType)