diff --git a/mlir/lib/Bindings/Python/IRModules.cpp b/mlir/lib/Bindings/Python/IRModules.cpp --- a/mlir/lib/Bindings/Python/IRModules.cpp +++ b/mlir/lib/Bindings/Python/IRModules.cpp @@ -153,6 +153,21 @@ return mlirStringRefCreate(s.data(), s.size()); } +template +static bool isPermutation(std::vector permutation) { + llvm::SmallVector 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. //------------------------------------------------------------------------------ @@ -3914,6 +3929,9 @@ "get_permutation", [](std::vector 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); diff --git a/mlir/test/Bindings/Python/ir_affine_map.py b/mlir/test/Bindings/Python/ir_affine_map.py --- a/mlir/test/Bindings/Python/ir_affine_map.py +++ b/mlir/test/Bindings/Python/ir_affine_map.py @@ -73,6 +73,12 @@ # 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: