Fix issue 1064: bug in bridge when lambda argument contains a loop. (#…

…1069) This is a workaround added to the `for` loop handling. A more general solution in the bridge should be implemented.
NVIDIA · Jan 4, 2024 · a8c2788 · a8c2788
1 parent 602caa1
commit a8c2788
Show file tree

Hide file tree

Showing 5 changed files with 134 additions and 33 deletions.
diff --git a/lib/Frontend/nvqpp/ConvertStmt.cpp b/lib/Frontend/nvqpp/ConvertStmt.cpp
@@ -322,36 +322,32 @@ bool QuakeBridgeVisitor::TraverseCompoundStmt(clang::CompoundStmt *stmt,
                                               DataRecursionQueue *q) {
   auto loc = toLocation(stmt->getSourceRange());
   SymbolTableScope var_scope(symbolTable);
+  auto traverseAndCheck = [&](clang::Stmt *cs) {
+    LLVM_DEBUG(llvm::dbgs() << "[[[\n"; cs->dump());
+    if (!TraverseStmt(cs))
+      reportClangError(cs, mangler, "statement not supported in qpu kernel");
+    LLVM_DEBUG({
+      if (!typeStack.empty()) {
+        llvm::dbgs() << "\n\nERROR: type stack has garbage after stmt:\n";
+        for (auto t : llvm::reverse(typeStack))
+          t.dump();
+        typeStack.clear();
+      }
+      llvm::dbgs() << "]]]\n";
+    });
+  };
   if (skipCompoundScope) {
     skipCompoundScope = false;
-    for (auto *cs : stmt->body()) {
-      LLVM_DEBUG(llvm::dbgs() << "[[[\n"; cs->dump());
-      if (TraverseStmt(static_cast<clang::Stmt *>(cs))) {
-        LLVM_DEBUG({
-          if (!typeStack.empty()) {
-            llvm::dbgs() << "\n\nERROR: type stack has garbage after stmt:\n";
-            for (auto t : llvm::reverse(typeStack))
-              t.dump();
-            typeStack.clear();
-          }
-          llvm::dbgs() << "]]]\n";
-        });
-      } else {
-        reportClangError(cs, mangler, "statement not supported in qpu kernel");
-      }
-    }
+    for (auto *cs : stmt->body())
+      traverseAndCheck(static_cast<clang::Stmt *>(cs));
     return true;
   }
-  bool result = true;
   builder.create<cc::ScopeOp>(loc, [&](OpBuilder &builder, Location loc) {
     for (auto *cs : stmt->body())
-      if (!TraverseStmt(static_cast<clang::Stmt *>(cs))) {
-        result = false;
-        return;
-      }
+      traverseAndCheck(static_cast<clang::Stmt *>(cs));
     builder.create<cc::ContinueOp>(loc);
   });
-  return result;
+  return true;
 }
 
 // Shared implementation for lowering of `do while` and `while` loops.
@@ -464,6 +460,7 @@ bool QuakeBridgeVisitor::TraverseForStmt(clang::ForStmt *x,
   bool result = true;
   auto loc = toLocation(x);
   auto *cond = x->getCond();
+  const auto initialValueDepth = valueStack.size();
   auto whileBuilder = [&](OpBuilder &builder, Location loc, Region &region) {
     if (!result)
       return;
@@ -523,6 +520,18 @@ bool QuakeBridgeVisitor::TraverseForStmt(clang::ForStmt *x,
     builder.create<cc::LoopOp>(loc, ValueRange{}, postCondition, whileBuilder,
                                bodyBuilder);
   }
+  const auto finalValueDepth = valueStack.size();
+  if (finalValueDepth > initialValueDepth) {
+    [[maybe_unused]] auto vals =
+        lastValues(finalValueDepth - initialValueDepth);
+    LLVM_DEBUG({
+      llvm::dbgs() << "Garbage left after loop:\n";
+      for (auto v : vals)
+        v.dump();
+      llvm::dbgs() << "\n";
+    });
+  }
+
   return result;
 }
 

diff --git a/test/AST-Quake/compute_action-3.cpp b/test/AST-Quake/compute_action-3.cpp
@@ -0,0 +1,85 @@
+/*******************************************************************************
+ * Copyright (c) 2022 - 2023 NVIDIA Corporation & Affiliates.                  *
+ * All rights reserved.                                                        *
+ *                                                                             *
+ * This source code and the accompanying materials are made available under    *
+ * the terms of the Apache License 2.0 which accompanies this distribution.    *
+ ******************************************************************************/
+
+// clang-format off
+// RUN: cudaq-quake %s | cudaq-opt --lambda-lifting --canonicalize --apply-op-specialization | FileCheck %s
+// clang-format on
+
+#include <cudaq.h>
+
+struct not_entry {
+  auto operator()() __qpu__ {
+    cudaq::qubit p;
+    cudaq::qubit q;
+    cudaq::compute_action([&]() {},
+                          [&]() {
+                            for (size_t i = 0; i < 1; ++i) {
+                            }
+                          });
+  }
+};
+
+struct entry {
+  auto operator()() __qpu__ {
+    cudaq::qubit p;
+    cudaq::qubit q;
+    cudaq::compute_action([&]() { x(p); },
+                          [&]() {
+                            for (size_t i = 0; i < 1; ++i) {
+                              y(q);
+                            }
+                          });
+  }
+};
+
+int not_main() {
+  auto counts = cudaq::sample(not_entry{});
+  counts.dump();
+  return 0;
+}
+
+int main() {
+  auto counts = cudaq::sample(entry{});
+  counts.dump();
+  return 0;
+}
+
+// clang-format off
+// CHECK-DAG:   func.func private @__nvqpp__lifted.lambda.0.adj() {
+// CHECK-DAG:   func.func private @__nvqpp__lifted.lambda.2.adj(
+
+// CHECK-LABEL:   func.func @__nvqpp__mlirgen__not_entry()
+// CHECK:           %[[VAL_0:.*]] = arith.constant 0 : i64
+// CHECK:           %[[VAL_1:.*]] = arith.constant 1 : i64
+// CHECK:           %[[VAL_2:.*]] = quake.alloca !quake.ref
+// CHECK:           %[[VAL_3:.*]] = quake.alloca !quake.ref
+// CHECK:           call @__nvqpp__lifted.lambda.0() : () -> ()
+// CHECK:           call @__nvqpp__lifted.lambda.1(%[[VAL_0]], %[[VAL_1]]) : (i64, i64) -> ()
+// CHECK:           call @__nvqpp__lifted.lambda.0.adj() : () -> ()
+// CHECK:           return
+// CHECK:         }
+
+// CHECK-LABEL:   func.func @__nvqpp__mlirgen__entry()
+// CHECK:           %[[VAL_0:.*]] = arith.constant 0 : i64
+// CHECK:           %[[VAL_1:.*]] = arith.constant 1 : i64
+// CHECK:           %[[VAL_2:.*]] = quake.alloca !quake.ref
+// CHECK:           %[[VAL_3:.*]] = quake.alloca !quake.ref
+// CHECK:           call @__nvqpp__lifted.lambda.2(%[[VAL_2]]) : (!quake.ref) -> ()
+// CHECK:           call @__nvqpp__lifted.lambda.3(%[[VAL_0]], %[[VAL_1]], %[[VAL_3]]) : (i64, i64, !quake.ref) -> ()
+// CHECK:           call @__nvqpp__lifted.lambda.2.adj(%[[VAL_2]]) : (!quake.ref) -> ()
+// CHECK:           return
+// CHECK:         }
+
+// CHECK-DAG:   func.func private @__nvqpp__callable.thunk.lambda.2(
+// CHECK-DAG:   func.func private @__nvqpp__lifted.lambda.2(
+// CHECK-DAG:   func.func private @__nvqpp__callable.thunk.lambda.1(
+// CHECK-DAG:   func.func private @__nvqpp__lifted.lambda.1(
+// CHECK-DAG:   func.func private @__nvqpp__callable.thunk.lambda.0(
+// CHECK-DAG:   func.func private @__nvqpp__lifted.lambda.0() {
+// CHECK-DAG:   func.func private @__nvqpp__callable.thunk.lambda.3(
+// CHECK-DAG:   func.func private @__nvqpp__lifted.lambda.3(
diff --git a/test/AST-error/inaccessible_variable.cpp b/test/AST-error/inaccessible_variable.cpp
@@ -20,12 +20,14 @@ struct Kernel {
 int main() {
   int i;
   auto f = [&]() __qpu__ { // expected-remark{{An inaccessible symbol}}
+    // `i` is a host variable, cannot be used in this kernel.
+    // expected-error@+2 {{statement not supported in qpu kernel}}
+    // expected-error@+1 {{symbol is not accessible in this kernel}}
     cudaq::qvector q(i);
-    // expected-error@-1 {{symbol is not accessible in this kernel}}
-    // expected-error@-2 {{statement not supported in qpu kernel}}
+    // declaration of `q` failed, so it's not available either.
+    // expected-error@+2 {{symbol is not accessible in this kernel}}
+    // expected-error@+1 {{statement not supported in qpu kernel}}
     mz(q);
-    // expected-error@-1 {{symbol is not accessible in this kernel}}
-    // expected-error@-2 {{statement not supported in qpu kernel}}
   };
   Kernel{}(f);
   return 0;

diff --git a/test/AST-error/statements.cpp b/test/AST-error/statements.cpp
@@ -16,8 +16,9 @@ void bar();
 
 struct S2 {
   auto operator()() __qpu__ {
-    try { // expected-error   {{try statement is not yet supported}}
-          // expected-error@-1{{statement not supported in qpu kernel}}
+    // expected-error@+2{{statement not supported in qpu kernel}}
+    // expected-error@+1{{try statement is not yet supported}}
+    try {
       foo();
     } catch (int) {
       bar();
@@ -27,17 +28,19 @@ struct S2 {
 
 struct S4 {
   auto operator()() __qpu__ {
-    goto label; // expected-error   {{goto statement}}
-                // expected-error@-1{{statement not supported in qpu kernel}}
+    // expected-error@+2{{statement not supported in qpu kernel}}
+    // expected-error@+1{{goto statement}}
+    goto label;
   label:
     foo();
   }
 };
 
 struct S5 {
   auto operator()(int sp) __qpu__ {
-    switch (sp) { // expected-error   {{switch statement}}
-                  // expected-error@-1{{statement not supported in qpu kernel}}
+    // expected-error@+2{{statement not supported in qpu kernel}}
+    // expected-error@+1{{switch statement}}
+    switch (sp) {
     case 1:
       foo();
     default:

diff --git a/test/AST-error/vector.cpp b/test/AST-error/vector.cpp
@@ -16,12 +16,14 @@ struct VectorVectorReturner {
   std::vector<std::vector<double>>
   operator()(std::vector<std::vector<int>> theta) __qpu__ {
     std::vector<std::vector<double>> result;
-    // expected-error@+1{{statement not supported in qpu kernel}}
     for (std::size_t i = 0, N = theta.size(); i < N; ++i) {
       auto &v = theta[i];
+      // expected-error@+2{{not supported in qpu}}
       // expected-error@+1{{vector dereference}}
       auto &r = result[i];
+      // expected-error@+1{{not supported in qpu}}
       for (std::size_t j = 0, M = v.size(); j < M; ++j)
+        // expected-error@+1{{symbol is not accessible}}
         r[j] = v[j];
     }
     // expected-error@+1{{C++ ctor (not-default)}}