[AffineToStaticLogic] Generalize loop bound computation. (#3234)

The loop bound computation previously relied on methods that attempted
to retrieve constant bounds, which would assert if the bounds were
non-constant (i.e. AffineMaps). This switches to use the same upstream
helpers that are used when lowering Affine loops to SCF loops, which
generate constant or dynamic Values using the Arithmetic dialect.

This fixes #3127.

lib/Conversion/AffineToStaticLogic/AffineToStaticLogic.cpp

@@ -352,13 +352,9 @@ LogicalResult AffineToStaticLogic::createStaticLogicPipeline(
   auto innerLoop = loopNest.back();
   ImplicitLocOpBuilder builder(outerLoop.getLoc(), outerLoop);
 
-  // Create constants for the loop's lower and upper bounds.
-  int64_t lbValue = innerLoop.getConstantLowerBound();
-  auto lowerBound = builder.create<arith::ConstantOp>(
-      IntegerAttr::get(builder.getIndexType(), lbValue));
-  int64_t ubValue = innerLoop.getConstantUpperBound();
-  auto upperBound = builder.create<arith::ConstantOp>(
-      IntegerAttr::get(builder.getIndexType(), ubValue));
+  // Create Values for the loop's lower and upper bounds.
+  Value lowerBound = lowerAffineLowerBound(innerLoop, builder);
+  Value upperBound = lowerAffineUpperBound(innerLoop, builder);
   int64_t stepValue = innerLoop.getStep();
   auto step = builder.create<arith::ConstantOp>(
       IntegerAttr::get(builder.getIndexType(), stepValue));

test/Conversion/AffineToStaticLogic/loops.mlir

@@ -64,3 +64,37 @@ func.func @dot(%arg0: memref<64xi32>, %arg1: memref<64xi32>) -> i32 {
 
   return %0 : i32
 }
+
+// CHECK-LABEL: func @affine_symbol
+#map0 = affine_map<()[s0] -> (s0 + 1)>
+func.func @affine_symbol(%arg0: i32) -> i32 {
+  %c0_i32 = arith.constant 0 : i32
+  // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+  // CHECK-DAG: %[[CAST:.+]] = arith.index_cast %arg0
+  // CHECK-DAG: %[[UB:.+]] = arith.addi %[[CAST]], %[[C1]]
+  %0 = arith.index_cast %arg0 : i32 to index
+  // CHECK: arith.cmpi ult, %arg1, %[[UB]]
+  %1 = affine.for %arg1 = 1 to #map0()[%0] iter_args(%arg2 = %c0_i32) -> (i32) {
+    %2 = arith.index_cast %arg1 : index to i32
+    %3 = arith.addi %arg2, %2 : i32
+    affine.yield %3 : i32
+  }
+  return %1 : i32
+}
+
+// CHECK-LABEL: func @affine_dimension
+#map1 = affine_map<(d0)[] -> (d0 + 1)>
+func.func @affine_dimension(%arg0: i32) -> i32 {
+  %c0_i32 = arith.constant 0 : i32
+  // CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+  // CHECK-DAG: %[[CAST:.+]] = arith.index_cast %arg0
+  // CHECK-DAG: %[[UB:.+]] = arith.addi %[[CAST]], %[[C1]]
+  %0 = arith.index_cast %arg0 : i32 to index
+  // CHECK: arith.cmpi ult, %arg1, %[[UB]]
+  %1 = affine.for %arg1 = 1 to #map1(%0) iter_args(%arg2 = %c0_i32) -> (i32) {
+    %2 = arith.index_cast %arg1 : index to i32
+    %3 = arith.addi %arg2, %2 : i32
+    affine.yield %3 : i32
+  }
+  return %1 : i32
+}