[Experimental][Buffers] A simple pass for inserting a single buffer at a specific channel

experimental/include/experimental/Transforms/HandshakePlaceBuffersCustom.h

@@ -0,0 +1,37 @@
+//===- HandshakePlaceBuffersCustom.h - Place buffers in DFG -----*- C++ -*-===//
+//
+// Dynamatic is 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares the --handshake-place-buffers-custom pass
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef EXPERIMENTAL_TRANSFORMS_BUFFERPLACEMENT_PLACEBUFFERS_CUSTOM_H
+#define EXPERIMENTAL_TRANSFORMS_BUFFERPLACEMENT_PLACEBUFFERS_CUSTOM_H
+
+#include "dynamatic/Support/DynamaticPass.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "llvm/ADT/StringRef.h"
+#include <string>
+
+namespace dynamatic {
+namespace experimental {
+namespace buffer {
+
+std::unique_ptr<dynamatic::DynamaticPass> createHandshakePlaceBuffersCustom(
+    const std::string &pred = "", const unsigned &outid = 1,
+    const unsigned &slots = 1, const std::string &type = "oehb");
+
+#define GEN_PASS_DECL_HANDSHAKEPLACEBUFFERSCUSTOM
+#define GEN_PASS_DEF_HANDSHAKEPLACEBUFFERSCUSTOM
+#include "experimental/Transforms/Passes.h.inc"
+
+} // namespace buffer
+} // namespace experimental
+} // namespace dynamatic
+
+#endif // EXPERIMENTAL_TRANSFORMS_BUFFERPLACEMENT_PLACEBUFFERS_CUSTOM_H

experimental/include/experimental/Transforms/Passes.h

@@ -17,6 +17,7 @@
 #include "dynamatic/Support/LLVM.h"
 #include "experimental/Transforms/HandshakeFixArgNames.h"
 #include "experimental/Transforms/Speculation/HandshakeSpeculation.h"
+#include "experimental/Transforms/HandshakePlaceBuffersCustom.h"
 #include "mlir/Pass/Pass.h"
 
 namespace dynamatic {
@@ -29,4 +30,4 @@ namespace experimental {
 } // namespace experimental
 } // namespace dynamatic
 
-#endif // EXPERIMENTAL_TRANSFORMS_PASSES_H
+#endif // EXPERIMENTAL_TRANSFORMS_PASSES_H

experimental/include/experimental/Transforms/Passes.td

@@ -51,4 +51,19 @@ def HandshakeSpeculation : DynamaticPass<"handshake-speculation"> {
     "operations.">];
 }
 
+def HandshakePlaceBuffersCustom : DynamaticPass<"handshake-placebuffers-custom"> {
+  let summary = "Place buffers on specific channels";
+  let description = [{ Placing a single buffer on a specific output channel of
+  a specfiic unit, this pass is useful for prototyping a custom placing pass
+  externally, e.g., written in python. }];
+  let constructor = "dynamatic::experimental::buffer::createHandshakePlaceBuffersCustom()";
+  let options = [
+  	Option<"pred", "pred", "std::string", "", "the predecessor unit of the channel">,
+  	Option<"outid", "outid", "unsigned", "", "output id of the predecessor, range: from 0 to number of outputs - 1">,
+  	Option<"slots", "slots", "unsigned", "", "num of slots of buffer, range: anything > 0">,
+  	Option<"type", "type", "std::string", "", "type of buffer, can be either oehb or tehb">
+  ];
+}
+
+
 #endif // EXPERIMENTAL_TRANSFORMS_PASSES_TD

experimental/lib/Transforms/CMakeLists.txt

@@ -1,5 +1,6 @@
 add_dynamatic_library(DynamaticExperimentalTransforms
   HandshakeFixArgNames.cpp
+  HandshakePlaceBuffersCustom.cpp
 
   DEPENDS
   DynamaticExperimentalTransformsPassIncGen

experimental/lib/Transforms/HandshakePlaceBuffersCustom.cpp

@@ -0,0 +1,111 @@
+//===- HandshakePlaceBuffersCustom.cpp - Place buffers in DFG ---*- C++ -*-===//
+//
+// Dynamatic is 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Buffer placement pass in Handshake functions, it takes the location (i.e.,
+// the predecessor, and which output channel of it), type (i.e., opaque or
+// transparent), and slots of the buffer that should be placed.
+//
+// This pass facilitates externally prototyping a custom buffer placement
+// analysis, e.g., in Python. This also makes the results of some research
+// artifacts (e.g., Mapbuf) developed in Python easily reproducible in the
+// current Dynamatic framework.
+//
+// Currently, this pass only supports adding units for a IR that is already
+// materized (i.e., each value is produced by exactly one producer, and consumed
+// by exactly one consumer). For future work on unmaterialized IRs, we need to
+// supply the producer--consumer pair as arguments
+//===----------------------------------------------------------------------===//
+
+#include "dynamatic/Analysis/NameAnalysis.h"
+#include "dynamatic/Dialect/Handshake/HandshakeDialect.h"
+#include "dynamatic/Dialect/Handshake/HandshakeOps.h"
+#include "dynamatic/Support/Attribute.h"
+#include "dynamatic/Support/CFG.h"
+#include "dynamatic/Transforms/HandshakeMaterialize.h"
+#include "experimental/Transforms/HandshakePlaceBuffersCustom.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "llvm/ADT/StringRef.h"
+
+using namespace llvm;
+using namespace dynamatic;
+using namespace dynamatic::experimental;
+using namespace dynamatic::experimental::buffer;
+
+namespace {
+
+struct HandshakePlaceBuffersCustomPass
+    : public dynamatic::experimental::buffer::impl::
+          HandshakePlaceBuffersCustomBase<HandshakePlaceBuffersCustomPass> {
+
+  /// Trivial field-by-field constructor.
+  HandshakePlaceBuffersCustomPass(const std::string &pred, const int &outid,
+                                  const int &slots, const std::string &type) {
+    this->pred = pred;
+    this->outid = outid;
+    this->slots = slots;
+    this->type = type;
+  }
+
+  /// Called on the MLIR module provided as input.
+  void runDynamaticPass() override {
+    mlir::ModuleOp modOp = getOperation();
+    MLIRContext *ctx = &getContext();
+
+    // Check if the IR is after being materialized, if not, reject the input
+    // IR.
+    if (failed(verifyIRMaterialized(modOp))) {
+      modOp->emitError() << ERR_NON_MATERIALIZED_MOD;
+      return signalPassFailure();
+    }
+
+    OpBuilder builder(ctx);
+    NameAnalysis &namer = getAnalysis<NameAnalysis>();
+    Operation *op = namer.getOp(pred);
+    if (!op) {
+      llvm::errs() << "No operation named \"" << pred << "\" exists\n";
+      return signalPassFailure();
+    }
+    assert(outid <= op->getNumResults() &&
+           "The output id exceeds the number of output ports!");
+    Value channel = op->getResult(outid);
+    // Set the insertion point to be before the original successor of the
+    // channel.
+    Operation *succ = *channel.getUsers().begin();
+    builder.setInsertionPoint(succ);
+    StringAttr slotName = StringAttr::get(ctx, "slots");
+    StringRef bufOpType;
+    if (type == "oehb") {
+      // if the specified type is "oehb", then we add a Opaque buffer with
+      // number of slots = slots.
+      bufOpType = handshake::OEHBOp::getOperationName();
+    } else if (type == "tehb") {
+      // If the specified type is "tehb", then we add a Transparent buffer
+      // with number of slots = slots.
+      bufOpType = handshake::TEHBOp::getOperationName();
+    } else {
+      llvm::errs() << "Unknown buffer type: \"" << type << "\"!\n";
+      return signalPassFailure();
+    }
+    NamedAttribute namedSlots(slotName, builder.getI32IntegerAttr(slots));
+    auto *bufOp =
+        builder.create(channel.getLoc(), StringAttr::get(ctx, bufOpType),
+                       channel, {channel.getType()}, {namedSlots});
+    inheritBB(succ, bufOp);
+    Value bufferRes = bufOp->getResult(0);
+    succ->replaceUsesOfWith(channel, bufferRes);
+  }
+};
+} // namespace
+
+std::unique_ptr<dynamatic::DynamaticPass>
+dynamatic::experimental::buffer::createHandshakePlaceBuffersCustom(
+    const std::string &pred, const unsigned &outid, const unsigned &slots,
+    const std::string &type) {
+  return std::make_unique<HandshakePlaceBuffersCustomPass>(pred, outid, slots,
+                                                           type);
+}