[mlir][Python] Add checking process before create an AffineMap from a permutation.

[zhanghb97]

Summary

An invalid permutation will trigger a C++ assertion when attempting to create an AffineMap from the permutation.
This patch adds an isPermutation function to check the given permutation before creating the AffineMap.

[mehdi_amini]

mlir/lib/Bindings/Python/IRModules.cpp

157

Typo: Permutation

mlir/lib/Bindings/Python/IRModules.cpp

createCustomDialectWrapper(const std::string &dialectNamespace,
   // Create the custom implementation.
   return (*dialectClass)(std::move(dialectDescriptor));
 }
 
 static MlirStringRef toMlirStringRef(const std::string &s) {
   return mlirStringRefCreate(s.data(), s.size());
 }
 
+template <typename PermutationTy>
+static bool isPermutation(std::vector<PermutationTy> permutation) {
+  llvm::SmallVector<bool, 8> seen(permutation.size(), false);
+  for (auto val : permutation) {
+    if (val < permutation.size()) {
+      if (seen[val])
+        return false;
+      seen[val] = true;
+      continue;
+    }
+    return false;
+  }
+  return true;
+}
+
 //------------------------------------------------------------------------------
 // Collections.
 //------------------------------------------------------------------------------
 
 namespace {
 
 class PyRegionIterator {
 public:
py::class_<PyAffineMap>(m, "AffineMap")
           py::arg("n_dims"), py::arg("n_results"),
           py::arg("context") = py::none(),
           "Gets a minor identity map with the given number of dimensions and "
           "results.")
       .def_static(
           "get_permutation",
           [](std::vector<unsigned> permutation,
              DefaultingPyMlirContext context) {
+            if (!isPermutation(permutation))
+              throw py::cast_error("Invalid permutation when attempting to "
+                                   "create an AffineMap");
             MlirAffineMap affineMap = mlirAffineMapPermutationGet(
                 context->get(), permutation.size(), permutation.data());
             return PyAffineMap(context->getRef(), affineMap);
           },
           py::arg("permutation"), py::arg("context") = py::none(),
           "Gets an affine map that permutes its inputs.")
       .def("get_submap",
            [](PyAffineMap &self, std::vector<intptr_t> &resultPos) {

mlir/test/Bindings/Python/ir_affine_map.py

with Context() as ctx:
 
     try:
       AffineMap.get(1, 1, [None])
     except RuntimeError as e:
       # CHECK: Invalid expression (None?) when attempting to create an AffineMap
       print(e)
 
     try:
+      AffineMap.get_permutation([1, 0, 1])
+    except RuntimeError as e:
+      # CHECK: Invalid permutation when attempting to create an AffineMap
+      print(e)
+
+    try:
       map3.get_submap([42])
     except ValueError as e:
       # CHECK: result position out of bounds
       print(e)
 
     try:
       map3.get_minor_submap(42)
     except ValueError as e: