MLIR-based code generation pipeline (daphne-eu#633)

As an alternative way to implement DAPHNE operators, this commit provides an MLIR-based code generation pipeline which progressively lowers the DaphneIR available after parsing the DaphneDSL script to operations in either the same dialect or operations from other dialects. With that, DAPHNE can optionally replace certain kernels by generating code directly, and also perform a hybrid compilation approach where we mix kernel calls with code generation in order to exploit advantages of both, precompiled kernel libraries and code generation.

This includes the following features:
- Documentation (Codegen.md).
- Code generation passes src/compiler/lowering/.
- Kernels to facilitate interoperability between the memref-dialect and the DenseMatrix runtime object.
- CLI arguments to enable code generation (--mlir_codegen and --mlir_codegen --mlir_hybrid_codegen) and explain-passes (--explain mlir_codegen).
- Script-level tests that run daphne with a DaphneDSL script as input in test/api/cli/codegen/.
- llvm-lit-based unit tests of the IR produced from passes/pass pipelines in test/codegen/ using intermediate IR as input (.mlir files).
- The daphne-opt tool, DAPHNE's own version of the mlir-opt tool with its source files in daphne-opt/ and a new build target daphne-opt.

daphne-opt is included since the lit tests require the daphne-opt target. With daphne-opt one can test passes in isolation by providing the input IR for a pass to daphne-opt and the correct flag to run the pass (or
a pass pipeline) on the IR. The daphne-opt tool inherits all the functionality of the mlir-opt tool.

The list of affected operations (only when daphne is executed with the --codegen flag): MatMulOp, MapOp, Ew(Sqrt|Abs|Add|Sub|Mul|Div|Pow|Mod)Op, AllAggSumOp, ConvertDenseMatrixToMemRefOp, ConvertMemRefToDenseMatrixOp.

daphne --mlir_codegen --mlir_hybrid_codegen can be used to disable code generation for certain operators in order to run a hybrid compilation pipeline (currently only the MatMulLoweringPass is affected by this.)

.github/workflows/main.yml

@@ -47,8 +47,7 @@ jobs:
     - name: Testing
       run: |
         mkdir --parents src/api/python/tmp
-        PYTHONPATH="$PYTHONPATH:$PWD/src/" bin/run_tests
-  
+        LD_LIBRARY_PATH=$PWD/lib:$LD_LIBRARY_PATH PATH=$PWD/bin:/usr/lib/llvm-10/bin:$PATH PYTHONPATH="$PYTHONPATH:$PWD/src/:/usr/lib/llvm-10/build/utils/lit/" bin/run_tests
 
     - name: "List generated files"
       run: |
@@ -64,4 +63,4 @@ jobs:
         name: daphne
         path: |
           bin/
-          lib/
+          lib/

.gitignore

@@ -6,6 +6,9 @@ build_*/
 /lib
 /tmp
 
+# runtime dump
+**/*.ll
+
 # documentation build output
 doc_build/
 
@@ -25,7 +28,26 @@ __pycache__/
 .idea/
 .clion.source.upload.marker
 
+# local test/dev scripts
 tmpdaphne.daphne
+*.daphne
+*.mlir
+*.log
+
+# tags file
+tags
+tags.lock
+tags.temp
+
+# clangd cache
+.cache/
+
+# gdb
+.gdb_history
+
+# compile commands
+compile_commands.json
+
 
 # release scripts output
 /artifacts
@@ -36,3 +58,7 @@ profiler/
 precompiled-dependencies/
 /cmake*/
 /data
+
+# Allow .daphne and .mlir files in test/
+!test/**/*.mlir
+!test/**/*.daphne

CMakeLists.txt

@@ -182,4 +182,5 @@ add_subdirectory(src/util)
 
 add_dependencies(CompilerUtils MLIRDaphneTransformsIncGen)
 
+add_subdirectory(daphne-opt)
 add_subdirectory(test)

UserConfig.json

@@ -3,6 +3,7 @@
     "use_vectorized_exec": false,
     "use_obj_ref_mgnt": true,
     "cuda_fuse_any": false,
+    "use_mlir_codegen": false,
     "vectorized_single_queue": false,
     "debug_llvm": false,
     "explain_kernels": false,
@@ -14,6 +15,7 @@
     "explain_type_adaptation": false,
     "explain_vectorized": false,
     "explain_obj_ref_mgnt": false,
+    "explain_mlir_codegen": false,
     "taskPartitioningScheme": "STATIC",
     "numberOfThreads": -1,
     "minimumTaskSize": 1,

containers/daphne.Dockerfile

@@ -64,7 +64,7 @@ LABEL "org.opencontainers.image.version"="$TIMESTAMP"
 LABEL "org.opencontainers.image.created"="${CREATION_DATE}"
 LABEL "org.opencontainers.image.revision"="${GIT_HASH}"
 RUN apt-get -qq -y update && apt-get -y upgrade && apt-get -y --no-install-recommends install  \
-    libtinfo6 libssl1.1 zlib1g python3-numpy python3-pandas \
+    libtinfo6 libssl1.1 zlib1g python3-numpy python3-pandas\
     && apt-get clean && rm -rf /var/lib/apt/lists/*
 COPY --from=daphne-build $DAPHNE_DIR/bin/* /usr/local/bin
 COPY --from=daphne-build $DAPHNE_DIR/lib/* /usr/local/lib

daphne-opt/CMakeLists.txt

@@ -0,0 +1,45 @@
+# Copyright 2023 The DAPHNE Consortium
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
+set(LIBS
+        ${dialect_libs}
+        ${conversion_libs}
+
+        MLIRDaphne
+        MLIRAnalysis
+        MLIRCallInterfaces
+        MLIRCastInterfaces
+        MLIRExecutionEngine
+        MLIRIR
+        # MLIRLLVMCommonConversion
+        MLIRLLVMToLLVMIRTranslation
+        # MLIRMemRefDialect
+        # MLIRLLVMDialect
+        MLIRParser
+        MLIRPass
+        MLIRSideEffectInterfaces
+        MLIRSupport
+        MLIRTargetLLVMIRExport
+        MLIRTransforms
+        MLIROptLib
+        )
+add_llvm_executable(daphne-opt daphne-opt.cpp)
+set_target_properties(daphne-opt PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}/bin)
+
+llvm_update_compile_flags(daphne-opt)
+target_link_libraries(daphne-opt PRIVATE ${LIBS})
+
+mlir_check_all_link_libraries(daphne-opt)

daphne-opt/daphne-opt.cpp

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2023 The DAPHNE Consortium
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "daphne-opt.h"
+
+#include <mlir/Dialect/LLVMIR/LLVMDialect.h>
+
+#include "ir/daphneir/Passes.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/InitLLVM.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/ToolOutputFile.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/Passes.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/InitAllDialects.h"
+#include "mlir/InitAllPasses.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Support/FileUtilities.h"
+#include "mlir/Tools/mlir-opt/MlirOptMain.h"
+
+int main(int argc, char **argv) {
+    mlir::registerAllPasses();
+    // NOTE: One can also register standalone passes here.
+    mlir::daphne::registerDaphnePasses();
+
+    mlir::DialectRegistry registry;
+    registry.insert<mlir::daphne::DaphneDialect, mlir::arith::ArithDialect,
+                    mlir::func::FuncDialect, mlir::scf::SCFDialect,
+                    mlir::LLVM::LLVMDialect, mlir::AffineDialect,
+                    mlir::memref::MemRefDialect, mlir::linalg::LinalgDialect,
+                    mlir::math::MathDialect>();
+    // Add the following to include *all* MLIR Core dialects, or selectively
+    // include what you need like above. You only need to register dialects that
+    // will be *parsed* by the tool, not the one generated
+    // registerAllDialects(registry);
+
+    return mlir::asMainReturnCode(mlir::MlirOptMain(
+        argc, argv, "Standalone DAPHNE optimizing compiler driver\n",
+        registry));
+}

daphne-opt/daphne-opt.h

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2023 The DAPHNE Consortium
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef DAPHNEOPT_DAPHNEOP_H
+#define DAPHNEOPT_DAPHNEOP_H
+
+#include "mlir/IR/Dialect.h"
+
+#include "ir/daphneir/Daphne.h"
+
+#endif // DAPHNEOPT_DAPHNEOP_H

doc/Codegen.md

@@ -0,0 +1,100 @@
+# Code Generation with MLIR
+
+This document describes the process of directly generating code with the MLIR
+framework.
+
+## Motivation
+
+DAPHNE provides a kernel for (almost) every DaphneIR operation which reside in
+`src/runtime/local/kernels/`. These are precompiled as a shared library and
+linked during compile-time. Even though these kernels can be highly optimized
+and thus achieve great runtime characteristics, they may not provide a desired
+level of extensibility for custom value types. They may also be lacking
+information only available at compile-time that could enable further
+optimizations. Additionally, through the process of progressively lowering the
+input IR, the code generation pipeline may enable more optimization
+possibilities such as operator or loop fusion.
+
+
+As an alternative way to implement our operators we provide the code generation
+pipeline which progressively lowers the DaphneIR available after parsing the
+DaphneDSL script to operations in either the same dialect or operations from
+other dialects. With that, we can optionally replace certain kernels by
+generating code directly, and also perform a hybrid compilation approach where
+we mix kernel calls with code generation in order to exploit advantages of
+both, precompiled kernel libraries and code generation. Code generation passes
+are found in `src/compiler/lowering/`.
+
+
+## Guidelines
+
+Currently, the code generation pipeline is enabled with the CLI flag
+`--mlir-codegen`. This adds the following passes that perform transformations and
+lowerings:
+
+- [DenseMatrixOptPass](src/compiler/lowering/DaphneOptPass.cpp)
+- [MatMulOpLoweringPass](src/compiler/lowering/MatMulOpLowering.cpp)
+- [AggAllLoweringPass](src/compiler/lowering/AggAllOpLowering.cpp)
+- [MapOpLoweringPass](src/compiler/lowering/MapOpLowering.cpp)
+- InlinerPass
+- [LowerEwOpPass](src/compiler/lowering/EwOpsLowering.cpp)
+- ConvertMathToLLVMPass
+- [ModOpLoweringPass](src/compiler/lowering/ModOpLowering.cpp)
+- Canonicalizer
+- CSE
+- LoopFusion
+- AffineScalarReplacement
+- LowerAffinePass
+
+These passes are added in the `DaphneIrExecutor::buildCodegenPipeline`
+function. The `--mlir-hybrid-codegen` flag disables the `MatMulOpLoweringPass` since the
+kernel implementation vastly outperforms the generated code of this pass.
+
+
+#### Runtime Interoperability
+
+Runtime interoperability with the `DenseMatrix` object is achieved with two
+kernels in `src/runtime/local/kernels/ConvertDenseMatrixToMemRef.h` and
+`src/runtime/local/kernels/ConvertMemRefToDenseMatrix.h` and the corresponding
+DaphneOps `Daphne_ConvertMemRefToDenseMatrix` and
+`Daphne_ConvertDenseMatrixToMemRef`. These kernels define how a MemRef is
+passed to a kernel and how a kernel can return a `StridedMemRefType`.
+
+
+#### Debugging
+
+In order to enable our debug `PrintIRPass` pass, one has to add `--explain
+mlir_code_gen` when running `daphne`. Additionally, it is recommended to use the
+`daphne-opt` tool to test passes in isolation. One just has to provide the
+input IR for a pass to `daphne-opt` and the correct flag to run the pass (or
+multiple passes) on the IR. `daphne-opt` provides all the functionality of the
+`mlir-opt` tool.
+
+`daphne-opt --lower-ew --debug-only=dialect-conversion ew.mlir` performs the
+`LowerEwOpPass` on the input file `ew.mlir` while providing dialect conversion
+debug information.
+
+
+
+#### Testing
+
+To test the generated code, there currently are two different approaches.
+
+End-to-end tests can be found under `test/api/cli/codegen/` and are part of the
+existing Catch2 test-suite with the its own tag, `TAG_CODEGEN`.
+
+Additionally, there are tests that check the generated IR by running the
+`llvm-lit`, `daphne-opt`, and `FileCheck` utilities. These tests reside under
+`test/compiler/lowering/`. They are `.mlir` files containing the input IR of a
+certain pass, or pass pipeline, and the `llvm-lit` directive at the top of the
+file (`RUN:`). In that line we specify how `llvm-lit` executes the test, e.g.,
+`// RUN: daphne-opt --lower-ew %s | FileCheck %s`, means that `daphne-opt` is
+called with the `--lower-ew` flag and the current file as input, the output of
+that, in addition to the file itself, is piped to `FileCheck`. `FileCheck` uses
+the comments in the `.mlir` file to check for certain conditions, e.g., `//
+CHECK-NOT: daphne.ewAdd` looks through the IR and fails if `daphne.ewAdd` can be
+found. These `llvm-lit` tests are all run by the `codegen` testcase in
+`test/codegen/Codegen.cpp`.
+
+
+All codegen tests can be executed by running `bin/run_tests '[codegen]'`.