[mlir][linalg][python] Reject unsigned pooling on non-integer element types in Python

Author: Men-cotton

mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py

@@ -532,9 +532,9 @@ def _binary_max_signed(self, lhs: Value, rhs: Value) -> Value:
         raise NotImplementedError("Unsupported 'max' operands: {lhs}, {rhs}")
 
     def _binary_max_unsigned(self, lhs: Value, rhs: Value) -> Value:
-        if _is_floating_point_type(lhs.type):
-            return arith.MaximumFOp(lhs, rhs).result
-        if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
+        if (
+            _is_integer_type(lhs.type) and not _is_bool_type(lhs.type)
+        ) or _is_index_type(lhs.type):
             return arith.MaxUIOp(lhs, rhs).result
         raise NotImplementedError("Unsupported 'max_unsigned' operands: {lhs}, {rhs}")
 
@@ -546,9 +546,9 @@ def _binary_min_signed(self, lhs: Value, rhs: Value) -> Value:
         raise NotImplementedError("Unsupported 'min' operands: {lhs}, {rhs}")
 
     def _binary_min_unsigned(self, lhs: Value, rhs: Value) -> Value:
-        if _is_floating_point_type(lhs.type):
-            return arith.MinimumFOp(lhs, rhs).result
-        if _is_integer_type(lhs.type) or _is_index_type(lhs.type):
+        if (
+            _is_integer_type(lhs.type) and not _is_bool_type(lhs.type)
+        ) or _is_index_type(lhs.type):
             return arith.MinUIOp(lhs, rhs).result
         raise NotImplementedError("Unsupported 'min_unsigned' operands: {lhs}, {rhs}")
 
@@ -634,6 +634,12 @@ def _is_index_type(t: Type) -> bool:
     return IndexType.isinstance(t)
 
 
+def _is_bool_type(t: Type) -> bool:
+    if not IntegerType.isinstance(t):
+        return False
+    return IntegerType(t).width == 1
+
+
 def _get_floating_point_width(t: Type) -> int:
     # TODO: Create a FloatType in the Python API and implement the switch
     # there.

mlir/test/python/dialects/linalg/opdsl/emit_pooling.py

@@ -150,3 +150,51 @@ def test_f32f32_min_pooling(input, shape, init_result):
 
 
 print(module)
+
+with Context() as ctx, Location.unknown():
+    module = Module.create()
+    with InsertionPoint(module.body):
+        f32 = F32Type.get()
+        bool_t = IntegerType.get_signless(1)
+
+        # CHECK: bool_max_unsigned_error: Unsupported 'max_unsigned' operands
+        @func.FuncOp.from_py_func(
+            RankedTensorType.get((1, 4, 16, 1), f32),
+            RankedTensorType.get((2, 2), f32),
+            RankedTensorType.get((1, 2, 4, 1), bool_t),
+        )
+        def test_bool_i1_max_unsigned_pooling_error(input, shape, init_result):
+            try:
+                pooling_poly(
+                    input,
+                    shape,
+                    outs=[init_result],
+                    reduce=BinaryFn.max_unsigned,
+                    cast=TypeFn.cast_unsigned,
+                    strides=[2, 4],
+                    dilations=[1, 2],
+                )
+            except NotImplementedError as e:
+                print(f"bool_max_unsigned_error: {e}")
+            return init_result
+
+        # CHECK: float_max_unsigned_error: Unsupported 'max_unsigned' operands
+        @func.FuncOp.from_py_func(
+            RankedTensorType.get((1, 4, 16, 1), f32),
+            RankedTensorType.get((2, 2), f32),
+            RankedTensorType.get((1, 2, 4, 1), f32),
+        )
+        def test_f32f32_max_unsigned_pooling_error(input, shape, init_result):
+            try:
+                pooling_poly(
+                    input,
+                    shape,
+                    outs=[init_result],
+                    reduce=BinaryFn.max_unsigned,
+                    cast=TypeFn.cast_unsigned,
+                    strides=[2, 4],
+                    dilations=[1, 2],
+                )
+            except NotImplementedError as e:
+                print(f"float_max_unsigned_error: {e}")
+            return init_result