Commit 9284f9fd authored by Nguyen, Thien Minh's avatar Nguyen, Thien Minh
Browse files

Handle for loop with negative range step 



Affine loop doesn't support negative step, hence need to transform the loop by ourselves.

Signed-off-by: default avatarThien Nguyen <nguyentm@ornl.gov>
parent 751830fa

mlir/parsers/qasm3/tests/test_optimization.cpp

+45 −0
Original line number Diff line number Diff line
@@ -259,6 +259,51 @@ for i in [0:100] { 
  EXPECT_EQ(countSubstring(llvm, "__quantum__qis__cnot"), 6);

}



TEST(qasm3PassManagerTester, checkAffineLoopRevert) {

  // Check loop with negative step:

  const std::string src = R"#(OPENQASM 3;

include "stdgates.inc";

def cnot_ladder() qubit[4]:q {

  h q;

  for i in [0:3] {

    cx q[i], q[i + 1];

  }

}



def cnot_ladder_inv() qubit[4]:q {

  for i in [3:-1:0] {

    cx q[i-1], q[i];

  }

  

  h q;

}



qubit q[4];

double theta = 0.01;

for i in [0:100] {

  cnot_ladder q;

  rz(theta) q[3];

  cnot_ladder_inv q;

}

)#";

  auto llvm =

      qcor::mlir_compile(src, "test_kernel", qcor::OutputType::LLVMIR, false);

  std::cout << "LLVM:\n" << llvm << "\n";

  

  // Get the main kernel section only 

  llvm = llvm.substr(llvm.find("@__internal_mlir_test_kernel"));

  const auto last = llvm.find_first_of("}");

  llvm = llvm.substr(0, last + 1);

  std::cout << "LLVM:\n" << llvm << "\n";

  // Only a single Rz remains (combine all angles)

  // 4 Hadamard before + 3 CX before

  // 4 Hadamard after + 3 CX after

  EXPECT_EQ(countSubstring(llvm, "__quantum__qis"), 15);

  EXPECT_EQ(countSubstring(llvm, "__quantum__qis__rz"), 1);

  EXPECT_EQ(countSubstring(llvm, "__quantum__qis__h"), 8);

  EXPECT_EQ(countSubstring(llvm, "__quantum__qis__cnot"), 6);

}



int main(int argc, char **argv) {

  ::testing::InitGoogleTest(&argc, argv);

  auto ret = RUN_ALL_TESTS();

mlir/parsers/qasm3/utils/expression_handler.cpp

+7 −2
Original line number Diff line number Diff line
@@ -444,8 +444,13 @@ antlrcpp::Any qasm3_expression_generator::visitAdditiveExpression( 
        }



        createOp<mlir::SubFOp>(location, lhs, rhs);

      } else if (lhs.getType().isa<mlir::IntegerType>() &&

                 rhs.getType().isa<mlir::IntegerType>()) {

      } else if ((lhs.getType().isa<mlir::IntegerType>() ||

                  lhs.getType().isa<mlir::IndexType>()) &&

                 (rhs.getType().isa<mlir::IntegerType>() ||

                  rhs.getType().isa<mlir::IndexType>())) {

        if (lhs.getType() != rhs.getType()) {

          rhs = builder.create<mlir::IndexCastOp>(location, lhs.getType(), rhs);

        }

        createOp<mlir::SubIOp>(location, lhs, rhs).result();

      } else {

        printErrorMessage("Could not perform subtraction, incompatible types: ",

mlir/parsers/qasm3/visitor_handlers/loop_stmt_handler.cpp

+42 −11
Original line number Diff line number Diff line
@@ -9,9 +9,16 @@ namespace { 
/// Creates a single affine "for" loop, iterating from lbs to ubs with

/// the given step.

/// to construct the body of the loop and is passed the induction variable.

void affineLoopBuilder(mlir::ValueRange lbs, mlir::ValueRange ubs, int64_t step,

void affineLoopBuilder(mlir::Value lbs_val, mlir::Value ubs_val, int64_t step,

                       std::function<void(mlir::Value)> bodyBuilderFn,

                       mlir::OpBuilder &builder, mlir::Location &loc) {

  // Note: Affine for loop only accepts **positive** step:

  // The stride, represented by step, is a positive constant integer which

  // defaults to “1” if not present.

  assert(step != 0);

  if (step > 0) {

    mlir::ValueRange lbs(lbs_val);

    mlir::ValueRange ubs(ubs_val);

    // Create the actual loop

    builder.create<mlir::AffineForOp>(

        loc, lbs, builder.getMultiDimIdentityMap(lbs.size()), ubs,
@@ -22,6 +29,30 @@ void affineLoopBuilder(mlir::ValueRange lbs, mlir::ValueRange ubs, int64_t step, 
          bodyBuilderFn(iv);

          nestedBuilder.create<mlir::AffineYieldOp>(nestedLoc);

        });

  } else {

    // Negative step:

    // a -> b step c (minus)

    // -a -> -b step c (plus) and minus the loop var

    mlir::Value minus_one = builder.create<mlir::ConstantOp>(

        loc, mlir::IntegerAttr::get(lbs_val.getType(), -1));

    lbs_val = builder.create<mlir::MulIOp>(loc, lbs_val, minus_one).result();

    ubs_val = builder.create<mlir::MulIOp>(loc, ubs_val, minus_one).result();

    mlir::ValueRange lbs(lbs_val);

    mlir::ValueRange ubs(ubs_val);

    builder.create<mlir::AffineForOp>(

        loc, lbs, builder.getMultiDimIdentityMap(lbs.size()), ubs,

        builder.getMultiDimIdentityMap(ubs.size()), -step, llvm::None,

        [&](mlir::OpBuilder &nestedBuilder, mlir::Location nestedLoc,

            mlir::Value iv, mlir::ValueRange itrArgs) {

          mlir::OpBuilder::InsertionGuard guard(nestedBuilder);

          mlir::Value minus_one_idx = nestedBuilder.create<mlir::ConstantOp>(

              nestedLoc, mlir::IntegerAttr::get(iv.getType(), -1));

          bodyBuilderFn(

              nestedBuilder.create<mlir::MulIOp>(nestedLoc, iv, minus_one_idx)

                  .result());

          nestedBuilder.create<mlir::AffineYieldOp>(nestedLoc);

        });

  }

}

} // namespace

namespace qcor {