small fixes (#172)

Author: makslevental

mlir/extras/dialects/ext/arith.py

@@ -132,7 +132,7 @@ def index_cast(
 ) -> Value:
     if loc is None:
         loc = get_user_code_loc()
-    assert bool(to) != bool(out), "either `to` or `out` must be set"
+    assert bool(to) != bool(out), "either `to` or `out` must be set but not both"
     res_type = out or to
     if res_type is None:
         res_type = IndexType.get()

mlir/extras/dialects/ext/scf.py

@@ -288,9 +288,10 @@ def __init__(self, operands, num_reductions, *, loc=None, ip=None):
             self.regions[i].blocks.append(operands[i].type, operands[i].type)
 
 
-def reduce_(*operands, num_reductions=1):
-    loc = get_user_code_loc()
-    return ReduceOp(operands, num_reductions, loc=loc)
+def reduce_(*operands, num_reductions=1, loc=None, ip=None):
+    if loc is None:
+        loc = get_user_code_loc()
+    return ReduceOp(operands, num_reductions, loc=loc, ip=ip)
 
 
 reduce = region_op(reduce_, terminator=lambda xs: reduce_return(*xs))

tests/test_scf.py

@@ -2400,11 +2400,11 @@ def res1(lhs: T.index(), rhs: T.index()):
             return lhs + rhs
 
         @another_reduce(res1)
-        def res1(lhs: T.index(), rhs: T.index()):
+        def res2(lhs: T.index(), rhs: T.index()):
             return lhs + rhs
 
-        @another_reduce(res1)
-        def res2(lhs: T.index(), rhs: T.index()):
+        @another_reduce(res2)
+        def res3(lhs: T.index(), rhs: T.index()):
             return lhs + rhs
 
     ctx.module.operation.verify()