[compiler] Integrate xDSL kernel compilation into IREE

.github/workflows/build.yaml

@@ -10,7 +10,7 @@ on:
   pull_request:
 
 jobs:
-  build-compiler:
+  build:
     runs-on: ubuntu-22.04
 
     steps:
@@ -40,9 +40,25 @@ jobs:
           sudo apt-get update
           sudo apt-get install lld clang
 
-      - name: Configure
+      - name: Install Python dependencies
+        run: |
+          pip install -r requirements.txt
+
+      - name: Install bender
+        uses: baptiste0928/cargo-install@v3
+        with:
+          crate: bender
+          version: '~0.28.0'
+
+      - name: Install Pulp Toolchain
         run: |
-          cmake -G Ninja -Bquidditch-compiler-build \
+          mkdir ${{github.workspace}}/pulp-toolchain
+          wget -qO- https://github.com/pulp-platform/llvm-project/releases/download/0.12.0/riscv32-pulp-llvm-ubuntu2004-0.12.0.tar.gz \
+          | tar --strip-components=1 -xzv -C ${{github.workspace}}/pulp-toolchain
+
+      - name: Configure Compiler
+        run: |
+          cmake -G Ninja -B ${{github.workspace}}/codegen/build \
           -DCMAKE_BUILD_TYPE=Release \
           -DCMAKE_C_COMPILER=clang \
           -DCMAKE_CXX_COMPILER=clang++ \
@@ -51,89 +67,28 @@ jobs:
           -DIREE_ENABLE_ASSERTIONS=ON \
           -DIREE_ENABLE_LLD=ON \
           -DIREE_ENABLE_THIN_ARCHIVES=ON \
-          -DIREE_HAL_DRIVER_DEFAULTS=OFF \
-          -DIREE_TARGET_BACKEND_DEFAULTS=OFF \
-          -DIREE_TARGET_BACKEND_LLVM_CPU=ON \
-          -DPython3_ROOT_DIR="$pythonLocation" \
-          -DPython3_FIND_STRATEGY=LOCATION \
+          -DPULP_CLANG_PATH=${{github.workspace}}/pulp-toolchain/bin/clang \
           -S ${{github.workspace}}/codegen
 
-      - name: Build
-        run: cmake --build quidditch-compiler-build --target iree-compile
-
-      # TODO: Test?
-
-      - name: Remove object files prior to upload
-        working-directory: ${{github.workspace}}/quidditch-compiler-build
-        run: |
-          find . -name "*.o" -type f -delete
-
-      - name: Tar build directory
-        working-directory: ${{github.workspace}}
-        run: |
-          tar -cvf quidditch-compiler-build.tar quidditch-compiler-build
-
-      - name: Upload iree-compile
-        uses: actions/upload-artifact@v4
-        with:
-          name: quidditch-compiler-build-dir
-          path: quidditch-compiler-build.tar
-
-  build-runtime:
-    runs-on: ubuntu-22.04
-    needs: [ build-compiler ]
-
-    steps:
-      - uses: actions/checkout@v4
-        with:
-          submodules: recursive
-          path: 'Quidditch'
+      - name: Build and Test Compiler
+        run: cmake --build ${{github.workspace}}/codegen/build --target check-quidditch
 
       - name: Download Quidditch Toolchain
         run: |
           mkdir ./toolchain
           docker run --rm ghcr.io/opencompl/quidditch/toolchain:main tar -cC /opt/quidditch-toolchain . \
           | tar -xC ./toolchain
 
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: 3.11
-
-      - name: Install Python dependencies
-        run: python -m pip install -r Quidditch/runtime/requirements.txt
-
-      - name: Install minimum required cmake and ninja
-        uses: lukka/get-cmake@latest
-        with:
-          cmakeVersion: "~3.21.0"
-
-      - name: Install bender
-        uses: baptiste0928/cargo-install@v3
-        with:
-          crate: bender
-          version: '~0.28.0'
-
-      - name: Download iree-compile
-        uses: actions/download-artifact@v4
-        with:
-          name: quidditch-compiler-build-dir
-
-      - name: Untar iree-compile
-        run: |
-          tar -xf quidditch-compiler-build.tar
-
-      - name: Configure build
+      - name: Configure Runtime
         run: |
           cmake -GNinja -Bquidditch-runtime-build \
-          -DQUIDDITCH_CODEGEN_BUILD_DIR=${{github.workspace}}/quidditch-compiler-build \
           -DCMAKE_TOOLCHAIN_FILE=${{github.workspace}}/toolchain/ToolchainFile.cmake \
-          -S ${{github.workspace}}/Quidditch/runtime
+          -S ${{github.workspace}}/runtime
 
-      - name: Build
+      - name: Build Runtime
         run: cmake --build quidditch-runtime-build
 
-      - name: Test
+      - name: Test Runtime
         working-directory: ${{github.workspace}}/quidditch-runtime-build
         # TODO: This should run a proper test suite once we are no longer using verilator.
         run: ctest --extra-verbose -j$(nproc) -R HelloWorld

.gitmodules

@@ -4,3 +4,6 @@
 [submodule "snitch_cluster"]
 	path = snitch_cluster
 	url = https://github.com/pulp-platform/snitch_cluster
+[submodule "xdsl"]
+	path = xdsl
+	url = https://github.com/xdslproject/xdsl

codegen/CMakeLists.txt

@@ -4,3 +4,8 @@ project(QuidditchCompiler LANGUAGES CXX C)
 set(IREE_SOURCE_DIR "${CMAKE_CURRENT_LIST_DIR}/../iree" CACHE STRING "IREE source code path")
 
 add_subdirectory(iree-configuration)
+
+enable_testing()
+include(CTest)
+
+add_subdirectory(tests)

codegen/compiler/src/Quidditch/CMakeLists.txt

@@ -5,15 +5,27 @@ iree_compiler_register_plugin(
         ::Quidditch
 )
 
+iree_tablegen_library(
+        NAME
+        PassesIncGen
+        TD_FILE
+        "Passes.td"
+        OUTS
+        --gen-pass-decls Passes.h.inc
+)
+
 iree_cc_library(
         NAME
         Quidditch
         HDRS
-
+        "Passes.h"
+        "Passes.h.inc"
         SRCS
         "QuidditchTarget.cpp"
+        "HoistHALOpsToFunc.cpp"
+        "FilterForxDSL.cpp"
         DEPS
-        IREELLVMCPUTargetDeps
+        ::PassesIncGen
         IREELinalgTransformDialect
         LLVMAnalysis
         LLVMBitReader
@@ -22,6 +34,10 @@ iree_cc_library(
         LLVMLinker
         LLVMSupport
         LLVMTargetParser
+        LLVMRISCVAsmParser
+        LLVMRISCVCodeGen
+        LLVMRISCVDesc
+        LLVMRISCVInfo
         MLIRArmNeonDialect
         MLIRArmSMEDialect
         MLIRArmSMEToLLVMIRTranslation
@@ -44,5 +60,6 @@ iree_cc_library(
         iree::compiler::Dialect::Util::IR
         iree::compiler::PluginAPI
         iree::compiler::Utils
+        iree::compiler::plugins::target::LLVMCPU
         PUBLIC
 )

codegen/compiler/src/Quidditch/FilterForxDSL.cpp

@@ -0,0 +1,28 @@
+#include "Passes.h"
+
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+
+namespace quidditch {
+#define GEN_PASS_DEF_FILTERFORXDSLPASS
+#include "Quidditch/Passes.h.inc"
+} // namespace quidditch
+
+using namespace mlir;
+
+namespace {
+class FilterForxDSL
+    : public quidditch::impl::FilterForxDSLPassBase<FilterForxDSL> {
+public:
+  using Base::Base;
+
+protected:
+  void runOnOperation() override;
+};
+} // namespace
+
+void FilterForxDSL::runOnOperation() {
+  getOperation()->walk([&](Operation* op) {
+    if (isa<memref::AssumeAlignmentOp>(op))
+      op->erase();
+  });
+}

codegen/compiler/src/Quidditch/HoistHALOpsToFunc.cpp

@@ -0,0 +1,98 @@
+#include "Passes.h"
+
+#include <iree/compiler/Dialect/HAL/IR/HALDialect.h>
+#include <mlir/Analysis/SliceAnalysis.h>
+#include <mlir/Dialect/Func/IR/FuncOps.h>
+#include <mlir/IR/IRMapping.h>
+
+namespace quidditch {
+#define GEN_PASS_DEF_HOISTHALOPSTOFUNCPASS
+#include "Quidditch/Passes.h.inc"
+} // namespace quidditch
+
+using namespace mlir;
+
+namespace {
+class HoistHALOpsToFunc
+    : public quidditch::impl::HoistHALOpsToFuncPassBase<HoistHALOpsToFunc> {
+public:
+  using Base::Base;
+
+protected:
+  void runOnOperation() override;
+};
+} // namespace
+
+void HoistHALOpsToFunc::runOnOperation() {
+  OpBuilder builder(&getContext());
+  builder.setInsertionPointToEnd(getOperation().getBody());
+
+  Dialect *dialect =
+      getContext()
+          .getLoadedDialect<mlir::iree_compiler::IREE::HAL::HALDialect>();
+  if (!dialect)
+    return;
+
+  FunctionType emptyFunctionType = builder.getFunctionType({}, {});
+  for (auto func : llvm::to_vector(getOperation().getOps<func::FuncOp>())) {
+    // IREE functions all have empty function types with public visibility. Skip
+    // over any other functions.
+    if (func.getFunctionType() != emptyFunctionType || !func.isPublic())
+      continue;
+
+    func->setAttr("xdsl_generated", builder.getUnitAttr());
+
+    // Find all HAL operations that need to be hoisted and any other operations
+    // they depend on.
+    SmallVector<Operation *> halOperations;
+    SetVector<Operation *> toClone;
+    func.getFunctionBody().walk([&](Operation *operation) {
+      if (operation->getDialect() != dialect)
+        return;
+
+      halOperations.push_back(operation);
+      for (Value result : operation->getResults()) {
+        unsigned int index = func.getNumArguments();
+        func.insertArgument(index, result.getType(),
+                            builder.getDictionaryAttr({}), result.getLoc());
+        result.replaceAllUsesWith(func.getArgument(index));
+      }
+
+      // Include all operations that the HAL operation transitively depends on.
+      BackwardSliceOptions options;
+      options.inclusive = true;
+      // This inserts into the SetVector in topological order. As we are going
+      // through HAL operations in-order, it is guaranteed that any operation
+      // already contained in the set appears prior to the HAL operation.
+      mlir::getBackwardSlice(operation, &toClone, options);
+    });
+
+    auto wrapper = builder.create<func::FuncOp>(
+        builder.getUnknownLoc(), (func.getName() + "$iree_to_xdsl").str(),
+        emptyFunctionType);
+
+    OpBuilder::InsertionGuard guard{builder};
+    builder.setInsertionPointToStart(wrapper.addEntryBlock());
+    // Create the clones and persist mapping between clones. This makes sure
+    // all operands are remapped.
+    IRMapping mapping;
+    for (Operation *op : toClone) {
+      if (!mlir::isPure(op)) {
+        op->emitError("Pass does not expect HAL operations to depend on "
+                      "impure operations");
+        signalPassFailure();
+        return;
+      }
+
+      builder.insert(op->clone(mapping));
+    }
+
+    SmallVector<Value> arguments;
+    for (Operation *op : halOperations)
+      for (Value value : op->getResults())
+        arguments.push_back(mapping.lookup(value));
+
+    builder.create<func::CallOp>(wrapper.getLoc(), func, arguments);
+    builder.create<func::ReturnOp>(wrapper.getLoc());
+  }
+}

codegen/compiler/src/Quidditch/Passes.h

@@ -0,0 +1,10 @@
+
+#pragma once
+
+#include <mlir/IR/BuiltinOps.h>
+#include <mlir/Pass/Pass.h>
+
+namespace quidditch {
+#define GEN_PASS_DECL
+#include "Quidditch/Passes.h.inc"
+} // namespace quidditch

codegen/compiler/src/Quidditch/Passes.td

@@ -0,0 +1,31 @@
+#ifndef QUIDDITCH_PASSES
+#define QUIDDITCH_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def HoistHALOpsToFuncPass
+  : Pass<"quidditch-hoist-hal-ops-to-func", "mlir::ModuleOp"> {
+  let description = [{
+    Pass performing HAL operation outlining.
+    For every function in the module this creates another function called "<old-name>$iree_to_xDSL" and moves all HAL
+    operations to that function before calling the original function.
+
+    This allows to pass the original function to xDSL for kernel compilation while compiling the new generated function
+    can be lowered by IREE.
+    This pass can be removed if xDSL were to learn how to lower IREE HAL operations to RISC-V.
+
+    The original function is additionally tagged with a "xdsl_generated" unit attribute.
+  }];
+}
+
+def FilterForxDSLPass
+  : Pass<"quidditch-filter-for-xdsl"> {
+  let description = [{
+    Pass rewriting operations to workaround xDSL operations.
+
+    Right now these are:
+    * Remove `memref.assume_alignment`.
+  }];
+}
+
+#endif