Update affine.py

Author: makslevental

mlir/python/mlir/dialects/affine.py

@@ -61,18 +61,28 @@ def __init__(
 
         if isinstance(lower_bound, int):
             lower_bound = AffineMap.get_constant(lower_bound)
-        elif isinstance(lower_bound, _ResultValueT):
+        elif isinstance(lower_bound, (Operation, OpView, Value)):
+            if len(lower_bound_operands):
+                raise ValueError(
+                    f"Either a concrete lower bound or an AffineMap in combination "
+                    f"with lower bound operands, but not both, is supported."
+                )
             lower_bound_operands.append(lower_bound)
-            lower_bound = AffineMap.get_constant(1)
+            lower_bound = AffineMap.get_identity(1)
 
         if not isinstance(lower_bound, AffineMap):
             raise ValueError(f"{lower_bound=} must be int | ResultValueT | AffineMap")
 
         if isinstance(upper_bound, int):
             upper_bound = AffineMap.get_constant(upper_bound)
-        elif isinstance(upper_bound, _ResultValueT):
+        elif isinstance(upper_bound, (Operation, OpView, Value)):
+            if len(upper_bound_operands):
+                raise ValueError(
+                    f"Either a concrete upper bound or an AffineMap in combination "
+                    f"with upper bound operands, but not both, is supported."
+                )
             upper_bound_operands.append(upper_bound)
-            upper_bound = AffineMap.get_constant(1)
+            upper_bound = AffineMap.get_identity(1)
 
         if not isinstance(upper_bound, AffineMap):
             raise ValueError(f"{upper_bound=} must be int | ResultValueT | AffineMap")

mlir/test/python/dialects/affine.py

@@ -119,12 +119,13 @@ def testForSugar():
     memref_t = T.memref(10, T.index())
     range = affine.for_
 
+    # CHECK: #[[$ATTR_2:.+]] = affine_map<(d0) -> (d0)>
+
     # CHECK-LABEL:   func.func @range_loop_1(
     # CHECK-SAME:                            %[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: memref<10xindex>) {
-    # CHECK:           affine.for %[[VAL_3:.*]] = 1 to 1 iter_args() -> () {
+    # CHECK:           affine.for %[[VAL_3:.*]] = #[[$ATTR_2]](%[[VAL_0]]) to #[[$ATTR_2]](%[[VAL_1]]) {
     # CHECK:             %[[VAL_4:.*]] = arith.addi %[[VAL_3]], %[[VAL_3]] : index
     # CHECK:             memref.store %[[VAL_4]], %[[VAL_2]]{{\[}}%[[VAL_3]]] : memref<10xindex>
-    # CHECK:             affine.yield
     # CHECK:           }
     # CHECK:           return
     # CHECK:         }
@@ -138,10 +139,9 @@ def range_loop_1(lb, ub, memref_v):
 
     # CHECK-LABEL:   func.func @range_loop_2(
     # CHECK-SAME:                            %[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: memref<10xindex>) {
-    # CHECK:           affine.for %[[VAL_3:.*]] = 1 to 10 iter_args() -> () {
+    # CHECK:           affine.for %[[VAL_3:.*]] = #[[$ATTR_2]](%[[VAL_0]]) to 10 {
     # CHECK:             %[[VAL_4:.*]] = arith.addi %[[VAL_3]], %[[VAL_3]] : index
     # CHECK:             memref.store %[[VAL_4]], %[[VAL_2]]{{\[}}%[[VAL_3]]] : memref<10xindex>
-    # CHECK:             affine.yield
     # CHECK:           }
     # CHECK:           return
     # CHECK:         }
@@ -154,10 +154,9 @@ def range_loop_2(lb, ub, memref_v):
 
     # CHECK-LABEL:   func.func @range_loop_3(
     # CHECK-SAME:                            %[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: memref<10xindex>) {
-    # CHECK:           affine.for %[[VAL_3:.*]] = 0 to 1 iter_args() -> () {
+    # CHECK:           affine.for %[[VAL_3:.*]] = 0 to #[[$ATTR_2]](%[[VAL_1]]) {
     # CHECK:             %[[VAL_4:.*]] = arith.addi %[[VAL_3]], %[[VAL_3]] : index
     # CHECK:             memref.store %[[VAL_4]], %[[VAL_2]]{{\[}}%[[VAL_3]]] : memref<10xindex>
-    # CHECK:             affine.yield
     # CHECK:           }
     # CHECK:           return
     # CHECK:         }