support Zhlt::BackOp in picus mode

risc0 · Dec 19, 2024 · af85011 · af85011
1 parent 8c837e6
commit af85011
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 1 deletion.
diff --git a/zirgen/compiler/picus/picus.cpp b/zirgen/compiler/picus/picus.cpp
@@ -139,7 +139,8 @@ class PicusPrinter {
               PackOp,
               ReturnOp,
               GetGlobalLayoutOp,
-              AliasLayoutOp>([&](auto op) { visitOp(op); })
+              AliasLayoutOp,
+              zirgen::Zhlt::BackOp>([&](auto op) { visitOp(op); })
         .Case<StoreOp, arith::ConstantOp>([](auto) { /* no-op */ })
         .Default([](Operation* op) { llvm::errs() << "unhandled op: " << *op << "\n"; });
   }
@@ -288,6 +289,16 @@ class PicusPrinter {
     }
   }
 
+  void visitOp(zirgen::Zhlt::BackOp back) {
+    auto callee = back.getCallee();
+    size_t distance = back.getDistance().getZExtValue();
+    AnySignal signal = signalize(freshName(), back.getType());
+    if (distance > 0) {
+      declareSignals(signal, /*isInput=*/true);
+    }
+    valuesToSignals.insert({back.getOut(), signal});
+  }
+
   // Constructs a fresh signal structure corresponding to the given type
   AnySignal signalize(std::string prefix, Type type) {
     if (isa<ValType>(type) || isa<RefType>(type)) {

diff --git a/zirgen/compiler/picus/test/back.zir b/zirgen/compiler/picus/test/back.zir
@@ -0,0 +1,13 @@
+// RUN: zirgen %s --emit=picus | FileCheck %s
+
+// CHECK: (prime-number 2013265921)
+// CHECK: (begin-module Count)
+// CHECK: (input x0)
+// CHECK: (end-module)
+
+#[picus]
+component Count(first: Val) {
+  public a : Reg;
+  a := Reg(a@1);
+}
+