fix the off-by-one error in the lowering of eval

j2kun · Oct 31, 2023 · 8422826 · 8422826
1 parent bc47670
commit 8422826
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/lib/Conversion/PolyToStandard/PolyToStandard.cpp b/lib/Conversion/PolyToStandard/PolyToStandard.cpp
@@ -149,6 +149,8 @@ struct ConvertEval : public OpConversionPattern<EvalOp> {
     auto lowerBound =
         b.create<arith::ConstantOp>(b.getIndexType(), b.getIndexAttr(1));
     auto numTermsOp = b.create<arith::ConstantOp>(b.getIndexType(),
+                                                  b.getIndexAttr(numTerms));
+    auto upperBound = b.create<arith::ConstantOp>(b.getIndexType(),
                                                   b.getIndexAttr(numTerms + 1));
     auto step = lowerBound;
 
@@ -163,7 +165,7 @@ struct ConvertEval : public OpConversionPattern<EvalOp> {
     auto accum =
         b.create<arith::ConstantOp>(b.getI32Type(), b.getI32IntegerAttr(0));
     auto loop = b.create<scf::ForOp>(
-        lowerBound, numTermsOp, step, accum.getResult(),
+        lowerBound, upperBound, step, accum.getResult(),
         [&](OpBuilder &builder, Location loc, Value loopIndex,
             ValueRange loopState) {
           ImplicitLocOpBuilder b(op.getLoc(), builder);

diff --git a/tests/poly_to_standard.mlir b/tests/poly_to_standard.mlir
@@ -82,10 +82,11 @@ func.func @test_lower_mul(%0 : !poly.poly<10>, %1 : !poly.poly<10>) -> !poly.pol
 // CHECK-LABEL: test_lower_eval
 // CHECK-SAME:  (%[[poly:.*]]: [[T:tensor<10xi32>]], %[[point:.*]]: i32) -> i32 {
 // CHECK:    %[[c1:.*]] = arith.constant 1 : index
+// CHECK:    %[[c10:.*]] = arith.constant 10 : index
 // CHECK:    %[[c11:.*]] = arith.constant 11 : index
 // CHECK:    %[[accum:.*]] = arith.constant 0 : i32
 // CHECK:    %[[loop:.*]] = scf.for %[[iv:.*]] = %[[c1]] to %[[c11]] step %[[c1]] iter_args(%[[iter_arg:.*]] = %[[accum]]) -> (i32) {
-// CHECK:        %[[coeffIndex:.*]] = arith.subi %c11, %[[iv]]
+// CHECK:        %[[coeffIndex:.*]] = arith.subi %[[c10]], %[[iv]]
 // CHECK:        %[[mulOp:.*]] = arith.muli %[[point]], %[[iter_arg]]
 // CHECK:        %[[nextCoeff:.*]] = tensor.extract %[[poly]][%[[coeffIndex]]]
 // CHECK:        %[[next:.*]] = arith.addi %[[mulOp]], %[[nextCoeff]]