[Handshake] Add pass that locks functions (#3855)

This PR introduces a handshake pass that adds a lock mechanism for functions.
Such a mechanism ensures that there is only one active control token in a function at each point in time.

An additional "RUN" directive/command was added to the integration tests that otherwise require a task-pipelining transformation. These test runs disable task-pipelining in favor of locking but should still yield the same results (albeit with lower throughput)

include/circt/Dialect/Handshake/HandshakePasses.h

@@ -32,6 +32,7 @@ std::unique_ptr<mlir::Pass> createHandshakeRemoveBuffersPass();
 std::unique_ptr<mlir::Pass> createHandshakeAddIDsPass();
 std::unique_ptr<mlir::OperationPass<handshake::FuncOp>>
 createHandshakeInsertBuffersPass();
+std::unique_ptr<mlir::Pass> createHandshakeLockFunctionsPass();
 
 /// Iterates over the handshake::FuncOp's in the program to build an instance
 /// graph. In doing so, we detect whether there are any cycles in this graph, as
@@ -54,6 +55,9 @@ LogicalResult addSinkOps(Region &r, OpBuilder &rewriter);
 LogicalResult addForkOps(Region &r, OpBuilder &rewriter);
 void insertFork(Value result, bool isLazy, OpBuilder &rewriter);
 
+// Adds a locking mechanism around the region.
+LogicalResult lockRegion(Region &r, OpBuilder &rewriter);
+
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "circt/Dialect/Handshake/HandshakePasses.h.inc"

include/circt/Dialect/Handshake/HandshakePasses.td

@@ -87,4 +87,14 @@ def HandshakeInsertBuffers
   ];
 }
 
+def HandshakeLockFunctions : Pass<"handshake-lock-functions", "handshake::FuncOp"> {
+  let summary = "Lock each function to only allow single invocations.";
+  let description = [{
+    This pass adds a locking mechanism to each handshake function. This mechanism
+    ensures that only one control token can be active in a function at each point
+    in time.
+  }];
+  let constructor = "circt::handshake::createHandshakeLockFunctionsPass()";
+}
+
 #endif // CIRCT_DIALECT_HANDSHAKE_HANDSHAKEPASSES_TD

integration_test/Dialect/Handshake/conditional_modification/conditional_modification.mlir

@@ -19,6 +19,14 @@
 // RUN: firtool --format=mlir --verilog --lowering-options=disallowLocalVariables > %t.sv && \
 // RUN: %PYTHON% %S/../cocotb_driver.py --objdir=%T --topLevel=top --pythonModule=conditional_modification --pythonFolder=%S %t.sv 2>&1 | FileCheck %s
 
+// Locking the circt should yield the same result
+// RUN: circt-opt %s --lower-std-to-handshake=disable-task-pipelining \
+// RUN:   --canonicalize='top-down=true region-simplify=true' --handshake-lock-functions \
+// RUN:   --handshake-materialize-forks-sinks --canonicalize \
+// RUN:   --handshake-insert-buffers=strategy=cycles --lower-handshake-to-firrtl | \
+// RUN: firtool --format=mlir --verilog --lowering-options=disallowLocalVariables > %t.sv && \
+// RUN: %PYTHON% %S/../cocotb_driver.py --objdir=%T --topLevel=top --pythonModule=conditional_modification --pythonFolder=%S %t.sv 2>&1 | FileCheck %s
+
 // CHECK:      ** TEST
 // CHECK-NEXT: ********************************
 // CHECK-NEXT: ** conditional_modification.oneInput

integration_test/Dialect/Handshake/multiple_loops/multiple_loops.mlir

@@ -19,6 +19,14 @@
 // RUN: firtool --format=mlir --verilog --lowering-options=disallowLocalVariables > %t.sv && \
 // RUN: %PYTHON% %S/../cocotb_driver.py --objdir=%T --topLevel=top --pythonModule=multiple_loops --pythonFolder=%S %t.sv 2>&1 | FileCheck %s
 
+// Locking the circt should yield the same result
+// RUN: circt-opt %s --lower-std-to-handshake=disable-task-pipelining \
+// RUN:   --canonicalize='top-down=true region-simplify=true' --handshake-lock-functions \
+// RUN:   --handshake-materialize-forks-sinks --canonicalize \
+// RUN:   --handshake-insert-buffers=strategy=cycles --lower-handshake-to-firrtl | \
+// RUN: firtool --format=mlir --verilog --lowering-options=disallowLocalVariables > %t.sv && \
+// RUN: %PYTHON% %S/../cocotb_driver.py --objdir=%T --topLevel=top --pythonModule=multiple_loops --pythonFolder=%S %t.sv 2>&1 | FileCheck %s
+
 // CHECK: ** TEST
 // CHECK: ** TESTS=[[NUM:.*]] PASS=[[NUM]] FAIL=0 SKIP=0
 

integration_test/Dialect/Handshake/nested_diamonds/nested_diamonds.mlir

@@ -20,6 +20,14 @@
 // RUN: firtool --format=mlir --verilog --lowering-options=disallowLocalVariables > %t.sv && \
 // RUN: %PYTHON% %S/../cocotb_driver.py --objdir=%T --topLevel=top --pythonModule=nested_diamonds --pythonFolder=%S %t.sv 2>&1 | FileCheck %s
 
+// Locking the circt should yield the same result
+// RUN: circt-opt %s --lower-std-to-handshake=disable-task-pipelining \
+// RUN:   --canonicalize='top-down=true region-simplify=true' --handshake-lock-functions \
+// RUN:   --handshake-materialize-forks-sinks --canonicalize \
+// RUN:   --handshake-insert-buffers=strategy=cycles --lower-handshake-to-firrtl | \
+// RUN: firtool --format=mlir --verilog --lowering-options=disallowLocalVariables > %t.sv && \
+// RUN: %PYTHON% %S/../cocotb_driver.py --objdir=%T --topLevel=top --pythonModule=nested_diamonds --pythonFolder=%S %t.sv 2>&1 | FileCheck %s
+
 // CHECK: ** TEST
 // CHECK: ** TESTS=[[NUM:.*]] PASS=[[NUM]] FAIL=0 SKIP=0
 

integration_test/Dialect/Handshake/nested_loops/nested_loops.mlir

@@ -19,6 +19,14 @@
 // RUN: firtool --format=mlir --verilog --lowering-options=disallowLocalVariables > %t.sv && \
 // RUN: %PYTHON% %S/../cocotb_driver.py --objdir=%T --topLevel=top --pythonModule=nested_loops --pythonFolder=%S %t.sv 2>&1 | FileCheck %s
 
+// Locking the circt should yield the same result
+// RUN: circt-opt %s --lower-std-to-handshake=disable-task-pipelining \
+// RUN:   --canonicalize='top-down=true region-simplify=true' --handshake-lock-functions \
+// RUN:   --handshake-materialize-forks-sinks --canonicalize \
+// RUN:   --handshake-insert-buffers=strategy=cycles --lower-handshake-to-firrtl | \
+// RUN: firtool --format=mlir --verilog --lowering-options=disallowLocalVariables > %t.sv && \
+// RUN: %PYTHON% %S/../cocotb_driver.py --objdir=%T --topLevel=top --pythonModule=nested_loops --pythonFolder=%S %t.sv 2>&1 | FileCheck %s
+
 // CHECK: ** TEST
 // CHECK: ** TESTS=[[NUM:.*]] PASS=[[NUM]] FAIL=0 SKIP=0
 

integration_test/Dialect/Handshake/task_pipelining/task_pipelining.mlir

@@ -19,6 +19,14 @@
 // RUN: firtool --format=mlir --verilog --lowering-options=disallowLocalVariables > %t.sv && \
 // RUN: %PYTHON% %S/../cocotb_driver.py --objdir=%T --topLevel=top --pythonModule=task_pipelining --pythonFolder=%S %t.sv 2>&1 | FileCheck %s
 
+// Locking the circt should yield the same result
+// RUN: circt-opt %s --lower-std-to-handshake=disable-task-pipelining \
+// RUN:   --canonicalize='top-down=true region-simplify=true' --handshake-lock-functions \
+// RUN:   --handshake-materialize-forks-sinks --canonicalize \
+// RUN:   --handshake-insert-buffers=strategy=cycles --lower-handshake-to-firrtl | \
+// RUN: firtool --format=mlir --verilog --lowering-options=disallowLocalVariables > %t.sv && \
+// RUN: %PYTHON% %S/../cocotb_driver.py --objdir=%T --topLevel=top --pythonModule=task_pipelining --pythonFolder=%S %t.sv 2>&1 | FileCheck %s
+
 // CHECK:      ** TEST
 // CHECK-NEXT: ********************************
 // CHECK-NEXT: ** task_pipelining.oneInput

integration_test/Dialect/Handshake/tp_memory/tp_memory.mlir

@@ -19,6 +19,14 @@
 // RUN: firtool --format=mlir --verilog --lowering-options=disallowLocalVariables > %t.sv && \
 // RUN: %PYTHON% %S/../cocotb_driver.py --objdir=%T --topLevel=top --pythonModule=tp_memory --pythonFolder=%S %t.sv 2>&1 | FileCheck %s
 
+// Locking the circt should yield the same result
+// RUN: circt-opt %s --lower-std-to-handshake=disable-task-pipelining \
+// RUN:   --canonicalize='top-down=true region-simplify=true' --handshake-lock-functions \
+// RUN:   --handshake-materialize-forks-sinks --canonicalize \
+// RUN:   --handshake-insert-buffers=strategy=cycles --lower-handshake-to-firrtl | \
+// RUN: firtool --format=mlir --verilog --lowering-options=disallowLocalVariables > %t.sv && \
+// RUN: %PYTHON% %S/../cocotb_driver.py --objdir=%T --topLevel=top --pythonModule=tp_memory --pythonFolder=%S %t.sv 2>&1 | FileCheck %s
+
 // CHECK:      ** TEST
 // CHECK-NEXT: ********************************
 // CHECK-NEXT: ** tp_memory.oneInput

lib/Dialect/Handshake/Transforms/CMakeLists.txt

@@ -3,12 +3,14 @@ add_circt_dialect_library(CIRCTHandshakeTransforms
   PassHelpers.cpp
   Materialization.cpp
   Buffers.cpp
+  LockFunctions.cpp
 
   DEPENDS
   CIRCTHandshakeTransformsIncGen
 
   LINK_LIBS PUBLIC
   CIRCTHandshake
+  CIRCTSupport
   MLIRIR
   MLIRPass
   MLIRTransformUtils

lib/Dialect/Handshake/Transforms/LockFunctions.cpp

@@ -0,0 +1,86 @@
+//===- LockFunctions.cpp - lock functions pass ------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Contains the definitions of the lock functions pass.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetails.h"
+#include "circt/Dialect/Handshake/HandshakeOps.h"
+#include "circt/Dialect/Handshake/HandshakePasses.h"
+#include "circt/Support/BackedgeBuilder.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Rewrite/FrozenRewritePatternSet.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace circt;
+using namespace handshake;
+using namespace mlir;
+
+LogicalResult handshake::lockRegion(Region &r, OpBuilder &rewriter) {
+  Block *entry = &r.front();
+  Location loc = r.getLoc();
+
+  if (entry->getNumArguments() == 0)
+    return r.getParentOp()->emitError("cannot lock a region without arguments");
+
+  auto *ret = r.front().getTerminator();
+  if (ret->getNumOperands() == 0)
+    return r.getParentOp()->emitError("cannot lock a region without results");
+
+  rewriter.setInsertionPointToStart(entry);
+  BackedgeBuilder bebuilder(rewriter, loc);
+  auto backEdge = bebuilder.get(rewriter.getNoneType());
+
+  auto buff = rewriter.create<handshake::BufferOp>(
+      loc, rewriter.getNoneType(), 1, backEdge,
+      /*bufferType=*/BufferTypeEnum::seq);
+
+  // Dummy value that causes a buffer initialization, but itself does not have a
+  // semantic meaning.
+  buff->setAttr("initValues", rewriter.getI64ArrayAttr({0}));
+
+  SmallVector<Value> inSyncOperands =
+      llvm::to_vector_of<Value>(entry->getArguments());
+  inSyncOperands.push_back(buff);
+  auto sync = rewriter.create<SyncOp>(loc, inSyncOperands);
+
+  // replace all func arg usages with the synced ones
+  // TODO is this UB?
+  for (auto &&[arg, synced] :
+       llvm::drop_end(llvm::zip(inSyncOperands, sync.getResults())))
+    arg.replaceAllUsesExcept(synced, sync);
+
+  rewriter.setInsertionPoint(ret);
+  SmallVector<Value> endJoinOperands = llvm::to_vector(ret->getOperands());
+  // Add the axilirary control signal output to the end-join
+  endJoinOperands.push_back(sync.getResults().back());
+  auto endJoin = rewriter.create<JoinOp>(loc, endJoinOperands);
+
+  backEdge.setValue(endJoin);
+  return success();
+}
+
+namespace {
+
+struct HandshakeLockFunctionsPass
+    : public HandshakeLockFunctionsBase<HandshakeLockFunctionsPass> {
+  void runOnOperation() override {
+    handshake::FuncOp op = getOperation();
+
+    OpBuilder builder(op);
+    if (failed(lockRegion(op.getRegion(), builder)))
+      signalPassFailure();
+  };
+};
+} // namespace
+
+std::unique_ptr<mlir::Pass>
+circt::handshake::createHandshakeLockFunctionsPass() {
+  return std::make_unique<HandshakeLockFunctionsPass>();
+}

test/Dialect/Handshake/errors-lock-functions.mlir

@@ -0,0 +1,15 @@
+// RUN: circt-opt %s --split-input-file -handshake-lock-functions --verify-diagnostics
+
+// expected-error @+1 {{cannot lock a region without arguments}}
+handshake.func @no_arg() -> none {
+  %ctrl = source
+  return %ctrl : none
+}
+
+// -----
+
+// expected-error @+1 {{cannot lock a region without results}}
+handshake.func @no_res(%ctrl: none) {
+  sink %ctrl : none
+  return
+}

test/Dialect/Handshake/test-lock-functions.mlir

@@ -0,0 +1,48 @@
+// RUN: circt-opt --split-input-file -handshake-lock-functions %s | FileCheck %s
+
+// CHECK-LABEL:   handshake.func @single_block(
+// CHECK-SAME:        %[[VAL_0:.*]]: i32, %[[VAL_1:.*]]: i32,
+// CHECK-SAME:        %[[VAL_2:.*]]: none, ...) -> (i32, none)
+// CHECK:           %[[VAL_3:.*]] = buffer [1] seq %[[VAL_4:.*]] {initValues = [0]} : none
+// CHECK:           %[[VAL_5:.*]]:4 = sync %[[VAL_0]], %[[VAL_1]], %[[VAL_2]], %[[VAL_3]] : i32, i32, none, none
+// CHECK:           %[[VAL_6:.*]] = arith.addi %[[VAL_5]]#0, %[[VAL_5]]#1 : i32
+// CHECK:           %[[VAL_4]] = join %[[VAL_6]], %[[VAL_5]]#2, %[[VAL_5]]#3 : i32, none, none
+// CHECK:           return %[[VAL_6]], %[[VAL_5]]#2 : i32, none
+// CHECK:         }
+
+handshake.func @single_block(%arg0: i32, %arg1: i32, %arg2: none, ...) -> (i32, none) {
+  %0 = arith.addi %arg0, %arg1 : i32
+  return %0, %arg2 : i32, none
+}
+
+// -----
+
+// CHECK-LABEL:   handshake.func @triangle(
+// CHECK-SAME:                             %[[VAL_0:.*]]: i32,
+// CHECK-SAME:                             %[[VAL_1:.*]]: i1,
+// CHECK-SAME:                             %[[VAL_2:.*]]: none, ...) -> (i32, none)
+// CHECK:           %[[VAL_3:.*]] = buffer [1] seq %[[VAL_4:.*]] {initValues = [0]} : none
+// CHECK:           %[[VAL_5:.*]]:4 = sync %[[VAL_0]], %[[VAL_1]], %[[VAL_2]], %[[VAL_3]] : i32, i1, none, none
+// CHECK:           %[[VAL_6:.*]]:2 = fork [2] %[[VAL_5]]#1 : i1
+// CHECK:           %[[VAL_7:.*]], %[[VAL_8:.*]] = cond_br %[[VAL_6]]#1, %[[VAL_5]]#0 : i32
+// CHECK:           sink %[[VAL_7]] : i32
+// CHECK:           %[[VAL_9:.*]], %[[VAL_10:.*]] = cond_br %[[VAL_6]]#0, %[[VAL_5]]#2 : none
+// CHECK:           %[[VAL_11:.*]]:2 = fork [2] %[[VAL_9]] : none
+// CHECK:           %[[VAL_12:.*]] = constant %[[VAL_11]]#0 {value = 42 : i32} : i32
+// CHECK:           %[[VAL_13:.*]], %[[VAL_14:.*]] = control_merge %[[VAL_11]]#1, %[[VAL_10]] : none
+// CHECK:           %[[VAL_15:.*]] = mux %[[VAL_14]] {{\[}}%[[VAL_12]], %[[VAL_8]]] : index, i32
+// CHECK:           %[[VAL_4]] = join %[[VAL_15]], %[[VAL_13]], %[[VAL_5]]#3 : i32, none, none
+// CHECK:           return %[[VAL_15]], %[[VAL_13]] : i32, none
+// CHECK:         }
+
+handshake.func @triangle(%arg0: i32, %arg1: i1, %arg2: none, ...) -> (i32, none) {
+  %0:2 = fork [2] %arg1 : i1
+  %trueResult, %falseResult = cond_br %0#1, %arg0 : i32
+  sink %trueResult : i32
+  %trueResult_0, %falseResult_1 = cond_br %0#0, %arg2 : none
+  %1:2 = fork [2] %trueResult_0 : none
+  %2 = constant %1#0 {value = 42 : i32} : i32
+  %result, %index = control_merge %1#1, %falseResult_1 : none
+  %3 = mux %index [%2, %falseResult] : index, i32
+  return %3, %result : i32, none
+}