Skip to content

Latest commit

 

History

History
501 lines (395 loc) · 32.9 KB

2.WritingASimplePass.md

File metadata and controls

501 lines (395 loc) · 32.9 KB

Writing a simple MLIR pass

The second chapter of this tutorial describes the implementation of a simple transformation pass in Dynamatic. This pass operates on Handshake-level IR and simplifies merge-like operations to make our dataflow circuits faster and smaller. We will

  1. declare the pass in TableGen, which will automatically generate a lot of boilerplate C++ code at compile-time,
  2. declare a header for the pass that includes some auto-generated code and declares the pass constructor,
  3. implement the pass constructor and its skeleton using some of the auto-generated code,
  4. configure the project to be able to run our pass with dynamatic-opt, the Dynamatic optimizer,
  5. and, finally, implement our circuit transformation.

You can write the entire pass yourself from the code snippets provided in this tutorial. The write-up assumes that no files related to the pass exist initially and walks you through the creation and implementation of those files. However, the full source code for this tutorial is provided in tutorials/include/tutorials/CreatingPasses and tutorials/lib/CreatingPasses for reference. To avoid name clashes while easily matching between the reference code and the code you may choose to write while reading this tutorial, all relevant names will be prefixed by My in the snippets present in this file compared to names used in the reference code. For example, the pass will be named MySimplifyMergeLike in this tutorial whereas it is named SimplifyMergeLike in the reference code.

The project is configured to build all tutorials along the rest of the project. By the end of this chapter, you will be able to run your own pass using Dynamatic's optimizer!

Declaring our pass in TableGen

The first step in the creation of our pass is to declare it inside a TableGen file (with .td extension). TableGen is an LLVM tool whose "purpose is to help a human develop and maintain records of domain-specific information". For our purposes, we can see TableGen as a preprocessor that inputs text files (with the .td extension by convention) containing information on user-defined MLIR entities (e.g., compiler passes, dialect operations, etc.) and outputs automatically-generated boilerplate C++ code that "implements" these entities. TableGen's input format is mostly declarative; it declares the existence of entities and characterizes their properties, but largely does not directly describe how these entities behave. The behavior of TableGen-defined entities must be written in C++, which we will do in the following sections. For this tutorial, we will have TableGen automatically generate the C++ code corresponding to our pass declaration. TableGen will also automatically generate a registration function that will enable the Dynamatic optimizer to register and run our pass.

Inside tutorials/include/tutorials/, which already exists, start by creating a directory named MyCreatingPasses which will contain all declarations for this tutorial. It's conventional to put the declaration of all transformation passes in a sub-directory called Transforms, so create one such directory within MyCreatingPasses. Finally, create a TableGen file named Passes.td inside that last directory. At this point, the filesystem should look like the following.

├── tutorials
│   ├── include 
│       ├── tutorials
│           ├── CreatingPasses # Reference code for this tutorial
│           ├── MyCreatingPasses # The first directory you just created
│               └── Transforms # The second directory you just created
│                   └── Passes.td # The file you just created
│           ├── CMakeLists.txt
│           └── InitAllPasses.h
│   ├── lib 
│   └── test
├── build.sh 
├── README.md
└── ... # Other files/folders at the top level

We will declare our pass inside Passes.td. Copy and paste the following snippet into the file.

//===- Passes.td - Transformation passes definition --------*- tablegen -*-===//
//
// This file contains the definition for all transformation passes in this
// tutorial.
//
//===----------------------------------------------------------------------===//

#ifndef TUTORIALS_MYCREATINGPASSES_TRANSFORMS_PASSES_TD
#define TUTORIALS_MYCREATINGPASSES_TRANSFORMS_PASSES_TD

include "mlir/Pass/PassBase.td"

def MySimplifyMergeLike : Pass< "tutorial-handshake-my-simplify-merge-like", 
                                "mlir::ModuleOp"> {
  let summary = "Simplifies merge-like operations in Handshake functions.";
  let description = [{
    The pass performs two simple transformation steps sequentially in each
    Handshake function present in the input MLIR module. First, it bypasses and
    removes all merge operations (circt::handshake::MergeOp) with a single
    operand from the IR, since they serve no purpose. Second, it downgrades all
    control merge operations (circt::handshake::ControlMergeOp) whose index
    result is unused into simpler merges with the same operands.
  }];
  let constructor = "dynamatic::tutorials::createMySimplifyMergeLikePass()";
}

#endif // TUTORIALS_MYCREATINGPASSES_TRANSFORMS_PASSES_TD

Let's go over the file's content. You may see that it shares syntactic similarity with C/C++. Like all C++ files in the repository, the file starts with a header comment containing some meta information as well as a description of the file's content. Like a header, it contains an include guard (#ifndef <guard>/#define <guard>/#endif) and includes another TableGen file (include "mlir/Pass/PassBase.td", note the lack of a # before the include keyword). The heart of the file is the declaration of the MySimplifyMergeLikePass which inherits from a Pass. The Pass object is given 2 generic arguments between <>.

  1. First is the flag name that will reference the pass in the Dynamatic optimizer tutorial-handshake-my-simplify-merge-like. Note that the actual flag name will be prefixed by a double-dash, so that it's possible to run the pass on some input Handshake-level IR with 
    $ ./bin/dynamatic-opt handshake-input.mlir --tutorial-handshake-my-simplify-merge-like
  2. Second is the MLIR operation that this pass matches on (i.e., the operation type the pass driver will look for in the input to run the pass on). In the vast majority of cases, we want passes to match an mlir::ModuleOp, which is always the top level operation under which everything is nested in our MLIR inputs.

The pass declaration contains some pass members which one must always define (there exists other members, but they are out of the scope of this tutorial). These are:

  • The summary, containing a one-line short description of what the pass does.
  • The description, containing a more detailed description of what the pass does.
  • The constructor, indicating the full qualified name of a function that returns a unique instance of the pass. We will declare and define this function in the next sections of this chapter. Notice that we create the function under the dynamatic::tutorials namespace. Every public member of Dynamatic should live in the dynamatic namespace. As to not pollute the repository's main namespace, everything related to the tutorials is further placed inside the nested tutorials namespace.

We now need to write some CMake configuration code to instruct the build system to automatically generate C++ code that corresponds to this TableGen file, and then compile this generated C++ along the rest of the project. First, create a file named CMakeLists.txt next to Passes.td with the following content.

set(LLVM_TARGET_DEFINITIONS Passes.td)
mlir_tablegen(Passes.h.inc -gen-pass-decls)
add_public_tablegen_target(DynamaticTutorialsMyCreatingPassesIncGen)
add_dependencies(dynamatic-headers DynamaticTutorialsMyCreatingPassesIncGen)

You do not need to understand precisely how this works. It suffices to know that it instructs the build system to create a target named DynamaticTutorialsMyCreatingPassesIncGen that libraries can depend on to get definitions related to Passes.td's content. To get this file included in the build when running $ cmake ..., we must include its parent directory from CMake files higher in the hierarchy. Modify the existing CMakeLists.txt in tutorials/include/tutorials to add the subdirectory we just created.

add_subdirectory(CreatingPasses)
add_subdirectory(MyCreatingPasses)

Similarly, create another CMakeLists.txt in tutorials/include/tutorial/MyCreatingPasses to include add nested subdirectory we created.

add_subdirectory(Transforms)

Everything we just did will eventually automatically generate a C++ header corresponding to Passes.td. It will be created inside the build directory (build/tutorials/include/tutorials/MyCreatingPasses/Transforms/Passes.h.inc) and will contain a lot of boilerplate code that you will rarely ever have to look at. Re-building the project right now would not generate the header because the build system would be able to identify that no part of the framework depends on it yet. We will see how to include parts of this header file inside our own C++ code using preprocessor flags in the next section, after which building the project will result in the header being genereted.

Declaring our pass in C++

Now that we got TableGen to generate the boilerplate code for this pass, we can finally start writing some C++ of our own. Create a header file called MySimplifyMergeLike.h next to Passes.td. We will include the auto-generated pass declaration and declare our pass constructor there using the following code.

//===- MySimplifyMergeLike.h - Simplifies merge-like ops --------*- C++ -*-===//
//
// This file declares the --tutorial-handshake-my-simplify-merge-like pass.
//
//===----------------------------------------------------------------------===//

#ifndef TUTORIALS_MYCREATINGPASSES_TRANSFORMS_MYSIMPLIFYMERGELIKE_H
#define TUTORIALS_MYCREATINGPASSES_TRANSFORMS_MYSIMPLIFYMERGELIKE_H

#include "dynamatic/Support/LLVM.h"
#include "mlir/Pass/Pass.h"

namespace dynamatic {
namespace tutorials {

#define GEN_PASS_DECL_MYSIMPLIFYMERGELIKE
#define GEN_PASS_DEF_MYSIMPLIFYMERGELIKE
#include "tutorials/MyCreatingPasses/Transforms/Passes.h.inc"

std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
createMySimplifyMergeLikePass();

} // namespace tutorials
} // namespace dynamatic

#endif // TUTORIALS_MYCREATINGPASSES_TRANSFORMS_MYSIMPLIFYMERGELIKE_H

Beyond the standard C++ header structure, this file does two important things.

  1. It includes the auto-generated pass declaration code inside the dynamatic::tutorials namespace. 
    #define GEN_PASS_DECL_MYSIMPLIFYMERGELIKE
#define GEN_PASS_DEF_MYSIMPLIFYMERGELIKE
#include "tutorials/MyCreatingPasses/Transforms/Passes.h.inc"
    Notice the preprocessor flag defined just before including the file. It serves the purpose of isolating a single part of the auto-generated header to include in our own header, here the declaration of our pass. The preprocessor's flag name is also auto-generated using the GEN_PASS_[DEF|DECL]_<my_pass_name_in_all_caps> template. If we were to define more passes inside Passes.td, all of them would get a declaration inside "tutorials/MyCreatingPasses/Transforms/Passes.h.inc". This preprocessor flag allows us to pick the single declaration we care about in the context.
  2. It declares our pass's constructor function, whose name we declared inside Passes.td. Do not pay much attention to the constructor's complicated-looking return type at this point, it is in fact trivial to implement this function.

Implementing the skeleton of our pass

We are now ready to start implementing our circuit transformation! We first write down some boilerplate skeleton code and configure CMake to build our implementation.

Inside tutorials/lib/, which already exists, start by creating two nested directories named MyCreatingPasses/Transforms (notice that the file structure is the same as in the tutorials/include/tutorials/ directory). Now, create a C++ source file named MySimplifyMergeLike.cpp inside the nested directory you just created to contain the implementation of our pass. Copy and paste the following code inside the source file.

//===- MySimplifyMergeLike.cpp - Simplifies merge-like ops ------*- C++ -*-===//
//
// Implements the --tutorial-handshake-my-simplify-merge-like pass, which uses a
// simple OpBuilder object to modify the IR within each handshake function.
//
//===----------------------------------------------------------------------===//

#include "tutorials/MyCreatingPasses/Transforms/MySimplifyMergeLike.h"
#include "circt/Dialect/Handshake/HandshakeOps.h"
#include "mlir/IR/MLIRContext.h"

using namespace mlir;
using namespace circt;

namespace {

/// Simple pass driver for our merge-like simplification transformation. At this
/// point it only prints a message to stdout.
struct MySimplifyMergeLikePass
    : public dynamatic::tutorials::impl::MySimplifyMergeLikeBase<
          MySimplifyMergeLikePass> {

  void runOnOperation() override {
    // Get the MLIR context for the current operation being transformed
    MLIRContext *ctx = &getContext();
    // Get the operation being transformed (the top level module)
    ModuleOp mod = getOperation();
    // Print a message on stdout to prove that the pass is running 
    llvm::outs() << "My pass is running!\n";
  };
};
} // namespace

namespace dynamatic {
namespace tutorials {

/// Returns a unique pointer to an operation pass that matches MLIR modules. In
/// our case, this is simply an instance of our unparameterized
/// MySimplifyMergeLikePass driver.
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
createMySimplifyMergeLikePass() {
  return std::make_unique<MySimplifyMergeLikePass>();
}
} // namespace tutorials
} // namespace dynamatic

Let's take a close look at the content of this source file, which for now only contains the skeleton of our pass. At the very bottom, we see the definition of our pass constructor that we declared in MySimplifyMergeLike.h. It simply returns a unique pointer to an instance of a MySimplifyMergeLikePass, which is a struct defined above inside an anonymous namespace. You can view this struct as the driver for our pass, and an instance of MySimplifyMergeLikePass as a particular instance of our pass. Let's break down the struct declaration and definition.

  • The struct declaration is quite verbose, but it will always have the same structure for any pass you implement. 
    struct MySimplifyMergeLikePass
: public dynamatic::tutorials::impl::MySimplifyMergeLikeBase<
      MySimplifyMergeLikePass> {...}
    The name MySimplifyMergeLikePass does not have any particular importance, but it is conventional to use the pass name as declared in the TableGen file (that we created in this section) suffixed by Pass. The struct inherits from MySimplifyMergeLikeBase, which is defined inside the dynamatic::tutorials::impl namespace. You may not remember defining this class anywhere. This is because it is the pass declaration that was auto-generated from TableGen inside "tutorials/MyCreatingPasses/Transforms/Passes.h.inc" and included from MySimplifyMergeLike.h, which the source file then includes. The name MySimplifyMergeLikeBase is auto-generated from the pass name declared in the TableGen file, to which Base is suffixed (it is the base class we inherit from). Finally, the base class is templated using... the derived struct's itself? This may seem counter-intuitive, and you may wonder how this could even compile, but it is in fact a well-known C++ idiom called the curiously recurring template pattern that is used throughout MLIR.
  • The struct overrides a single method named runOnOperation. It is the method that will be called on each mlir::ModuleOp found in the input IR, since we declared our pass (in Passes.td) to match this operation type. Right now, the method just retrieves the current MLIR context and operation it was matched on, and prints a message to standard output. In the next section, we will implement our circuit transformation within this method.

Running our pass

Configuring CMake

We now configure CMake to build this pass along the rest of the project. We have to create a CMakeLists.txt file in each directory we created and modify the one at tutorials/lib. Starting with the latter, just add a line to include the new directory structure in the build.

add_subdirectory(CreatingPasses)
add_subdirectory(MyCreatingPasses)

Similarly, inside lib/MyCreatingPasses/CMakeLists.txt, just write the following to include the Transforms subdirectory, where our pass implemenation lies.

add_subdirectory(Transforms)

Finally, add the following snippet to lib/MyCreatingPasses/Transforms/CMakeLists.txt.

add_dynamatic_library(DynamaticTutorialsMyCreatingPasses
  MySimplifyMergeLike.cpp

  DEPENDS
  DynamaticTutorialsMyCreatingPassesIncGen

  LINK_LIBS PUBLIC
  MLIRIR
  MLIRSupport
  MLIRTransformUtils
)

This CMake file creates a new Dynamatic library called DynamaticTutorialsMyCreatingPasses, which includes our pass implementation (MySimplifyMergeLike.cpp) and depends on DynamaticTutorialsMyCreatingPassesIncGen (the TableGen target we created earlier in tutorials/include/tutorials/CreatingPasses/Transforms/CMakeLists.txt) as well as a couple of standard MLIR targets which are built as part of our software dependencies.

The last CMake step is to add your new dynamatic library to the optimizer by modifying tools/dynamatic-opt/CMakeLists.txt. This will allow the optimizer to include your pass implementation in its binary. Add your library to the list of existing libraries that the dynamatic-opt tool gets linked to as follows.

target_link_libraries(dynamatic-opt
  PRIVATE
  DynamaticTransforms
  DynamaticTutorialsCreatingPasses
  DynamaticTutorialsMyCreatingPasses # your library!

  <... other libraries>
)

Registering our pass

To be able to run a pass, the optimizer needs to register it at compile-time. The tool is already configured to register all tutorial passes by calling the dynamatic::tutorials::registerAllPasses() function located in tutorials/include/tutorials/InitAllPasses.h, so we just have to add our own pass to this function. To do that, first create a file named Passes.h inside tutorials/include/tutorials/MyCreatingPasses/Transforms/, and paste the following into it.

//===- Passes.h - Transformation passes registration ------------*- C++ -*-===//
//
// This file contains declarations to register transformation passes.
//
//===----------------------------------------------------------------------===//

#ifndef TUTORIALS_MYCREATINGAPASSES_TRANSFORMS_PASSES_H
#define TUTORIALS_MYCREATINGAPASSES_TRANSFORMS_PASSES_H

#include "dynamatic/Support/LLVM.h"
#include "mlir/Pass/Pass.h"
#include "tutorials/MyCreatingPasses/Transforms/MySimplifyMergeLike.h"

namespace dynamatic {
namespace tutorials {

namespace MyCreatingPasses {

/// Generate the code for registering passes.
#define GEN_PASS_REGISTRATION
#include "tutorials/MyCreatingPasses/Transforms/Passes.h.inc"

} // namespace MyCreatingPasses
} // namespace tutorials
} // namespace dynamatic

#endif // TUTORIALS_MYCREATINGAPASSES_TRANSFORMS_PASSES_H

Similarly to MySimplifyMergeLike.h, this file includes some auto-generated code from "tutorials/MyCreatingPasses/Transforms/Passes.h.inc". This time, however, the GEN_PASS_REGISTRATION pre-processor flag indicates that the pass registration functions should be included instead of the pass declarations.

Next, open tutorials/include/tutorials/InitAllPasses.h and add the file you just created to the list of include statements.

#include "tutorials/MyCreatingPasses/Transforms/Passes.h"

Finally, inside the file tutorials/CreatingPasses/include/tutorials/CreatingPasses/InitAllPasses.h in the function registerAllPasses, add the following line to register your pass using the auto-generated registerPasses method.

dynamatic::tutorials::MyCreatingPasses::registerPasses();

We created a lot of directories and files in the last two sections, so let's recap what our file system should look like at this point.

├── tutorials
│   ├── include 
│       └── tutorials
│           ├── CreatingPasses # Reference code for this tutorial
│           ├── MyCreatingPasses
│               ├── CMakeLists.txt
│               └── Transforms
│                   ├── CMakeLists.txt 
│                   ├── MySimplifyMergeLike.h 
│                   ├── Passes.td
│                   └── Passes.h # The file you just created
│           ├── CMakeLists.txt
│           └── InitAllPasses.h # Modified just now to register your pass
│   ├── lib
│       ├── CMakeLists.txt # Modified to add_subdirectory(MyCreatingPasses)
│       ├── CreatingPasses # Reference code for this tutorial
│       └── MyCreatingPasses # All created by you
│           ├── CMakeLists.txt # add_subdirectory(Transforms)
│           └── Transforms
│                   ├── CMakeLists.txt # add_dynamatic_library(...)
│                   └── MySimplifyMergeLike.cpp # Pass skeleton
│   └── test
├── build.sh 
├── README.md
└── ... # Other files/folders at the top level

You should now be able to compile your skeleton pass implementation using the repository's build script (./build.sh, from the top-level directory). Once successfully compiled, and to verify that everything works as intended, try to run your pass on the test file located at tutorials/test/creating-passes.mlir using the following command (run from the repository's top level).

$ ./bin/dynamatic-opt tutorials/test/creating-passes.mlir --tutorial-handshake-my-simplify-merge-like

On stdout, you should see printed the message we put into the pass (My pass is running!) as well as the MLIR input. The optimizer's behavior is to print out the transformed IR after going through all passes. Our pass performs no IR modification at this point, so the input IR gets printed unmodified.

Congratulations on successfully building your own pass! It may seem like a long (and somewhat boilerplate) process but, once you are used to it, it takes only 5 to 10 minutes to setup a pass as these steps are mostly the same for all passes you will ever write. Also keep in mind that you usually won't have to do all of what we just did, since most of the time all the basic infrastructure (i.e., the Tablegen file, some of the headers, the CMakeLists.txt files) is already there. In those cases you would just have to declare an additional pass inside a Passes.td file, add a header/source file pair for your new pass, and include your pass's header inside an already existing Passes.h file. We we will do exactly that in the next chapter.

Implementing our transformation

It's finally time to write our circuit transformation! In this section, we will just be modifying MySimplifyMergeLike.cpp. As this tutorial is mostly about the pass creation process rather than MLIR's IR transformation capabilities, we will not go into the details of how to interact with MLIR data-structures. Instead, see the MLIR primer for an introduction to these concepts.

Start by modifying the runOnOperation method inside MySimplifyMergeLikePass to call a helper function that will perform the transformation for each handshake function (circt::handshake::FuncOp) in the current MLIR module.

void runOnOperation() override {
  // Get the MLIR context for the current operation being transformed
  MLIRContext *ctx = &getContext();
  // Get the operation being transformed (the top level module)
  ModuleOp mod = getOperation();

  // Iterate over all Handshake functions in the module
  for (handshake::FuncOp funcOp : mod.getOps<handshake::FuncOp>())
    // Perform the simple transformation individually on each function. In
    // case the transformation fails for at least a function, the pass should
    // be considered failed
    if (failed(performSimplification(funcOp, ctx)))
      return signalPassFailure();
}

We iterate over all handshake functions in the module using mod.getOps<circt::handshake::FuncOp>() and simplify each of them sequentially using the performSimplification function, which we will write next. In case the transformation fails for a function, we tell the optimizer by calling signalPassFailure() and returning. On receiving this signal, the optimizer will stop processing the IR (cancelling any pass that was supposed to run after ours) and return.

Now, create the skeleton of the function that will perform our transformation outside and above of the anonymous namespace that contains MySimplifyMergeLikePass.

/// Performs the simple transformation on the provided Handshake function,
/// deleting merges with a single input and downgrades control merges with an
/// unused index result into simpler merges.
static LogicalResult performSimplification(handshake::FuncOp funcOp,
                                           MLIRContext *ctx) {
  // Create an operation builder to allow us to create and insert new operation
  // inside the function
  OpBuilder builder(ctx);

  return success();
}

The function returns a LogicalResult, which is the conventional MLIR type to indicate success (return success();) or failure (return failure();). At this point, the function just creates an operation builder (OpBuilder) from the passed MLIR context, which will enable us to create/insert/erase operation from the IR.

Now, add the code of the first transformation step (single-input merge erasure) inside the function.

static LogicalResult performSimplification(handshake::FuncOp funcOp,
                                           MLIRContext *ctx) {
  OpBuilder builder(ctx);

  // Erase all merges with a single input
  for (handshake::MergeOp mergeOp :
        llvm::make_early_inc_range(funcOp.getOps<handshake::MergeOp>())) {
    if (mergeOp->getNumOperands() == 1) {
      // Replace all occurences of the merge's single result throughout the IR
      // with the merge's single operand. This is equivalent to bypassing the
      // merge
      mergeOp.getResult().replaceAllUsesWith(mergeOp.getOperand(0));
      // Erase the merge operation, whose result now has no uses
      mergeOp.erase();
    }
  }

  return success();
}

This simply iterates over all circt::handshake::MergeOp inside the function and, if they have a single operand, rewires the circuit to bypass the useless merge before deleting the latter. Note that we wrap the funcOp.getOps<handshake::MergeOp>() iterator inside a call to llvm::make_early_inc_range. This is necessary because we are erasing the current element pointed to by the iterator inside the loop body (by calling mergeOp.erase()), which is normally unsafe. make_early_inc_range solves this by going to find the next iterator element before returning control to the loop body for the current element.

Next, add the code for the second transformation step (index-less control merge downgrading) below the code we just added.

static LogicalResult performSimplification(handshake::FuncOp funcOp,
                                           MLIRContext *ctx) {
  // [First transformation here]

  // Replace control merges with an unused index result into merges
  for (handshake::ControlMergeOp cmergeOp :
       llvm::make_early_inc_range(funcOp.getOps<handshake::ControlMergeOp>())) {

    // Get the control merge's index result (second result).
    // Equivalently, we could have written:
    //  auto indexResult = cmergeOp->getResults()[1];
    // but using getIndex() is more readable and maintainable
    Value indexResult = cmergeOp.getIndex();

    // We can only perform the transformation if the control merge operation's
    // index result is not used throughout the IR
    if (!indexResult.use_empty())
      continue;

    // Now, we create a new merge operation at the same position in the IR as
    // the control merge we are replacing. The merge has the exact same inputs
    // as the control merge
    builder.setInsertionPoint(cmergeOp);
    handshake::MergeOp newMergeOp = builder.create<handshake::MergeOp>(
        cmergeOp.getLoc(), cmergeOp->getOperands());

    // Then, replace the control merge's first result (the selected input) with
    // the single result of the newly created merge operation
    Value mergeRes = newMergeOp.getResult();
    cmergeOp.getResult().replaceAllUsesWith(mergeRes);

    // Finally, we can delete the original control merge, whose results have
    // no uses anymore
    cmergeOp->erase();
  }

  return success();
}

Again, we simply iterate over all circt::handshake::ControlMergeOp and, for those that have no uses to their index result, replace them with simpler merges. To achieve that, we create a new merge (with the same inputs/operands as the control_merge) at the location of the existing control_merge using builder.create<handshake::MergeOp>(...), rewire the circuit appropriately, and erase the now unused control merge. We again use llvm::make_early_inc_range for the same reason as before.

We have now finished implementing our circuit transformation! Rebuild the project and re-run the following to see the transformed IR printed on stdout.

$ ./bin/dynamatic-opt tutorials/test/creating-passes.mlir --tutorial-handshake-my-simplify-merge-like
module {
  handshake.func @eraseSingleInputMerge(%arg0: none, ...) -> none attributes {argNames = ["start"], resNames = ["out0"]} {
    %0 = return %arg0 : none
    end %0 : none
  }
  handshake.func @downgradeIndexLessControlMerge(%arg0: i32, %arg1: i32, %arg2: none, ...) -> i32 attributes {argNames = ["arg0", "arg1", "start"], resNames = ["out0"]} {
    %0 = merge %arg0, %arg1 : i32
    %1 = return %0 : i32
    end %1 : i32
  }
  handshake.func @isMyArgZero(%arg0: i32, %arg1: none, ...) -> i1 attributes {argNames = ["arg0", "start"], resNames = ["out0"]} {
    %0 = constant %arg1 {value = 0 : i32} : i32
    %1 = arith.cmpi eq, %arg0, %0 : i32
    %trueResult, %falseResult = cond_br %1, %arg1 : none
    %2 = merge %trueResult : none
    %3 = constant %2 {value = true} : i1
    %4 = br %3 : i1
    %5 = merge %falseResult : none
    %6 = constant %5 {value = false} : i1
    %7 = br %6 : i1
    %result, %index = control_merge %2, %5 : none, index
    %8 = mux %index [%4, %7] : index, i1
    %9 = return %8 : i1
    end %9 : i1
  }
}

Compared to the input IR, we can see that:

  • eraseSingleInputMerge lost its single-input merge.
  • downgradeIndexLessControlMerge had its control_merge turned into a simpler merge.
  • isMyArgZero lost its two single-input merges at the top of the function, and its two first control_merges were downgraded to merges (the last one wasn't as its index result is used by the mux).

Congratulations! Your dataflow circuits will now be faster and smaller!

Conclusion

In this chapter, we described in details the full process of creating an MLIR pass from scratch and implemented a simple Handshake-level IR transformation as an example. We verified that the pass works as intended using some simple test inputs that we ran through dynamatic-opt.

Unfortunately, it turns out that our pass misses some optimization opportunities that it should ideally be able to catch. Consider our last test function in tutorials/test/creating-passes.mlir. As we observed in the previous section, two of its index-less control_merges got downgraded to merges, which is expected. These merges, however, could further be removed from the IR since they have a single input, but our pass fails to accomplish this. Generally speaking, the problem is that optimizing these initial controL_merges is, according to how we defined our pass, a two-steps process (first downgrading, then erasure). However, our pass performs the merge erasure step before the control merge downgrading step and then never goes back to it. We could simply fix this issue by reversing the order of these steps, or running our pass a second time on the already transformed IR (doing so is usually an indication of bad design). These solutions will work for this particular pass, which only performs two different optimizations, but what if we had a pass that matched and transformed 10 different IR constructs? How would we know in which order to apply the transformations to get the most optimized IR possible in all cases? Would there exist such an order? The answer to our problem is called greedy pattern rewriting, and we will cover it in this tutorial's next chapter.