diff --git a/.github/workflows/build_and_test_cmake.yml b/.github/workflows/build_and_test_cmake.yml index a257c48..78d51db 100644 --- a/.github/workflows/build_and_test_cmake.yml +++ b/.github/workflows/build_and_test_cmake.yml @@ -29,6 +29,14 @@ jobs: ./externals/llvm-project key: ${{ runner.os }}-cmake-${{ hashFiles('bazel/import_llvm.bzl') }}-${{ hashFiles('**/CMakeLists.txt') }} + - name: Cache mlir-tutorial build + id: cache-mlir-tutorial + uses: actions/cache@v3 + with: + path: | + ./build + key: ${{ runner.os }}-cmake-${{ hashFiles('bazel/import_llvm.bzl') }}-${{ hashFiles('**/CMakeLists.txt') }} + - name: Git config run: | git config --global --add safe.directory ${GITHUB_WORKSPACE} @@ -47,9 +55,10 @@ jobs: - name: Build and test mlir-tutorial run: | mkdir build && cd build - cmake -DLLVM_DIR=${GITHUB_WORKSPACE}/externals/llvm-project/build/lib/cmake/llvm -DMLIR_DIR=${GITHUB_WORKSPACE}/externals/llvm-project/build/lib/cmake/mlir .. + cmake -DLLVM_DIR=${GITHUB_WORKSPACE}/externals/llvm-project/build/lib/cmake/llvm -DMLIR_DIR=${GITHUB_WORKSPACE}/externals/llvm-project/build/lib/cmake/mlir -DBUILD_DEPS="ON" .. cmake --build . --target MLIRAffineFullUnrollPasses cmake --build . --target MLIRMulToAddPasses + cmake --build . --target MLIRNoisyPasses cmake --build . --target mlir-headers cmake --build . --target tutorial-opt cmake --build . --target check-mlir-tutorial diff --git a/.gitignore b/.gitignore index 36802d6..258a6a8 100644 --- a/.gitignore +++ b/.gitignore @@ -4,9 +4,6 @@ bazel-mlir-tutorial bazel-out bazel-testlogs -# git submodule -externals - # cmake related files # ignore the user specified CMake presets in subproject directories. /*/CMakeUserPresets.json diff --git a/CMakeLists.txt b/CMakeLists.txt index 56aab9f..7d3ca4c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -3,6 +3,7 @@ cmake_minimum_required(VERSION 3.20.0) project(mlir-tutorial LANGUAGES CXX C) set(CMAKE_CXX_STANDARD 17 CACHE STRING "C++ standard to conform to") +set(CMAKE_POSITION_INDEPENDENT_CODE ON) find_package(MLIR REQUIRED CONFIG) @@ -22,6 +23,16 @@ include_directories(${PROJECT_SOURCE_DIR}) include_directories(${PROJECT_SOURCE_DIR}/externals/llvm-project) include_directories(${PROJECT_BINARY_DIR}) +message(STATUS "Fetching or-tools...") +include(FetchContent) +FetchContent_Declare( + or-tools + GIT_REPOSITORY https://github.com/google/or-tools.git + GIT_TAG main +) +FetchContent_MakeAvailable(or-tools) +message(STATUS "Done fetching or-tools") + add_subdirectory(tests) add_subdirectory(tools) -add_subdirectory(lib) \ No newline at end of file +add_subdirectory(lib) diff --git a/README.md b/README.md index a894731..99a0e5d 100644 --- a/README.md +++ b/README.md @@ -1,33 +1,32 @@ # MLIR For Beginners -This is the code repository for a series of articles -on the [MLIR framework](https://mlir.llvm.org/) for building compilers. +This is the code repository for a series of articles on the +[MLIR framework](https://mlir.llvm.org/) for building compilers. ## Articles -1. [Build System (Getting Started)](https://jeremykun.com/2023/08/10/mlir-getting-started/) -2. [Running and Testing a Lowering](https://jeremykun.com/2023/08/10/mlir-running-and-testing-a-lowering/) -3. [Writing Our First Pass](https://jeremykun.com/2023/08/10/mlir-writing-our-first-pass/) -4. [Using Tablegen for Passes](https://jeremykun.com/2023/08/10/mlir-using-tablegen-for-passes/) -5. [Defining a New Dialect](https://jeremykun.com/2023/08/21/mlir-defining-a-new-dialect/) -6. [Using Traits](https://jeremykun.com/2023/09/07/mlir-using-traits/) -7. [Folders and Constant Propagation](https://jeremykun.com/2023/09/11/mlir-folders/) -8. [Verifiers](https://jeremykun.com/2023/09/13/mlir-verifiers/) -9. [Canonicalizers and Declarative Rewrite Patterns](https://jeremykun.com/2023/09/20/mlir-canonicalizers-and-declarative-rewrite-patterns/) +1. [Build System (Getting Started)](https://jeremykun.com/2023/08/10/mlir-getting-started/) +2. [Running and Testing a Lowering](https://jeremykun.com/2023/08/10/mlir-running-and-testing-a-lowering/) +3. [Writing Our First Pass](https://jeremykun.com/2023/08/10/mlir-writing-our-first-pass/) +4. [Using Tablegen for Passes](https://jeremykun.com/2023/08/10/mlir-using-tablegen-for-passes/) +5. [Defining a New Dialect](https://jeremykun.com/2023/08/21/mlir-defining-a-new-dialect/) +6. [Using Traits](https://jeremykun.com/2023/09/07/mlir-using-traits/) +7. [Folders and Constant Propagation](https://jeremykun.com/2023/09/11/mlir-folders/) +8. [Verifiers](https://jeremykun.com/2023/09/13/mlir-verifiers/) +9. [Canonicalizers and Declarative Rewrite Patterns](https://jeremykun.com/2023/09/20/mlir-canonicalizers-and-declarative-rewrite-patterns/) 10. [Dialect Conversion](https://jeremykun.com/2023/10/23/mlir-dialect-conversion/) 11. [Lowering through LLVM](https://jeremykun.com/2023/11/01/mlir-lowering-through-llvm/) - ## Bazel build Bazel is one of two supported build systems for this tutorial. The other is CMake. If you're unfamiliar with Bazel, you can read the tutorials at [https://bazel.build/start](https://bazel.build/start). Familiarity with Bazel is not required to build or test, but it is required to follow the articles in -the tutorial series and explained in the first article, [Build System (Getting -Started)](https://jeremykun.com/2023/08/10/mlir-getting-started/). The CMake -build is maintained, but was added at article 10 (Dialect Conversion) and will -not be explained in the articles. +the tutorial series and explained in the first article, +[Build System (Getting Started)](https://jeremykun.com/2023/08/10/mlir-getting-started/). +The CMake build is maintained, but was added at article 10 (Dialect Conversion) +and will not be explained in the articles. ### Prerequisites @@ -51,16 +50,16 @@ bazel test ...:all CMake is one of two supported build systems for this tutorial. The other is Bazel. If you're unfamiliar with CMake, you can read the tutorials at -[https://cmake.org/getting-started/](https://cmake.org/getting-started/). -The CMake build is maintained, but was added at article 10 (Dialect Conversion) -and will not be explained in the articles. +[https://cmake.org/getting-started/](https://cmake.org/getting-started/). The +CMake build is maintained, but was added at article 10 (Dialect Conversion) and +will not be explained in the articles. ### Prerequisites -* Make sure you have installed everything needed to build LLVM -https://llvm.org/docs/GettingStarted.html#software -* For this recipe Ninja is used so be sure to have it as well installed -https://github.com/ninja-build/ninja/wiki/Pre-built-Ninja-packages +* Make sure you have installed everything needed to build LLVM + https://llvm.org/docs/GettingStarted.html#software +* For this recipe Ninja is used so be sure to have it as well installed + https://github.com/ninja-build/ninja/wiki/Pre-built-Ninja-packages ### Checking out the code @@ -131,12 +130,14 @@ LLVM_BUILD_DIR=externals/llvm-project/build cmake -G $BUILD_SYSTEM .. \ -DLLVM_DIR="$LLVM_BUILD_DIR/lib/cmake/llvm" \ -DMLIR_DIR="$LLVM_BUILD_DIR/lib/cmake/mlir" \ + -DBUILD_DEPS="ON" \ -DCMAKE_BUILD_TYPE=Debug popd cmake --build $BUILD_DIR --target MLIRAffineFullUnrollPasses cmake --build $BUILD_DIR --target MLIRMulToAddPasses +cmake --build $BUILD_DIR --target MLIRNoisyPasses cmake --build $BUILD_DIR --target mlir-headers cmake --build $BUILD_DIR --target mlir-doc cmake --build $BUILD_DIR --target tutorial-opt diff --git a/WORKSPACE b/WORKSPACE index 0d5cf04..ed24794 100644 --- a/WORKSPACE +++ b/WORKSPACE @@ -1,6 +1,6 @@ workspace(name = "mlir_tutorial") -load("@bazel_tools//tools/build_defs/repo:git.bzl", "new_git_repository") +load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository", "new_git_repository") load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive") load("@bazel_tools//tools/build_defs/repo:utils.bzl", "maybe") @@ -18,7 +18,9 @@ http_archive( load("@bazel_skylib//:workspace.bzl", "bazel_skylib_workspace") bazel_skylib_workspace() + load("@bazel_skylib//lib:versions.bzl", "versions") + versions.check(minimum_bazel_version = "6.3.2") # A two-step process for buliding LLVM/MLIR with bazel. First the raw source @@ -108,3 +110,122 @@ pip_parse( load("@mlir_tutorial_pip_deps//:requirements.bzl", "install_deps") install_deps() + +##### Deps for or-tools ##### + +## Bazel rules. +git_repository( + name = "platforms", + commit = "380c85cc2c7b126c6e354f517dc16d89fe760c9f", + remote = "https://github.com/bazelbuild/platforms.git", +) + +git_repository( + name = "rules_proto", + commit = "3f1ab99b718e3e7dd86ebdc49c580aa6a126b1cd", + remote = "https://github.com/bazelbuild/rules_proto.git", +) + +## ZLIB +new_git_repository( + name = "zlib", + build_file = "@com_google_protobuf//:third_party/zlib.BUILD", + commit = "04f42ceca40f73e2978b50e93806c2a18c1281fc", + remote = "https://github.com/madler/zlib.git", +) + +## Re2 +git_repository( + name = "com_google_re2", + remote = "https://github.com/google/re2.git", + tag = "2023-07-01", +) + +## Abseil-cpp +git_repository( + name = "com_google_absl", + commit = "c2435f8342c2d0ed8101cb43adfd605fdc52dca2", + patch_args = ["-p1"], + patches = ["@com_google_ortools//patches:abseil-cpp-20230125.3.patch"], + remote = "https://github.com/abseil/abseil-cpp.git", +) + +## Protobuf +git_repository( + name = "com_google_protobuf", + # there's a patch for the CMake build in protobuf, ignoring + # patches = ["@com_google_ortools//patches:protobuf-v23.3.patch"], + commit = "4dd15db6eb3955745f379d28fb4a2fcfb6753de3", + patch_args = ["-p1"], + remote = "https://github.com/protocolbuffers/protobuf.git", +) + +# Load common dependencies. +load("@com_google_protobuf//:protobuf_deps.bzl", "protobuf_deps") + +protobuf_deps() + +## Solvers +http_archive( + name = "glpk", + build_file = "@com_google_ortools//bazel:glpk.BUILD", + sha256 = "4a1013eebb50f728fc601bdd833b0b2870333c3b3e5a816eeba921d95bec6f15", + url = "http://ftp.gnu.org/gnu/glpk/glpk-5.0.tar.gz", +) + +http_archive( + name = "bliss", + build_file = "@com_google_ortools//bazel:bliss.BUILD", + patches = ["@com_google_ortools//bazel:bliss-0.73.patch"], + sha256 = "f57bf32804140cad58b1240b804e0dbd68f7e6bf67eba8e0c0fa3a62fd7f0f84", + url = "https://github.com/google/or-tools/releases/download/v9.0/bliss-0.73.zip", + #url = "http://www.tcs.hut.fi/Software/bliss/bliss-0.73.zip", +) + +new_git_repository( + name = "scip", + build_file = "@com_google_ortools//bazel:scip.BUILD", + commit = "62fab8a2e3708f3452fad473a6f48715c367316b", + patch_args = ["-p1"], + patches = ["@com_google_ortools//bazel:scip.patch"], + remote = "https://github.com/scipopt/scip.git", +) + +# Eigen has no Bazel build. +new_git_repository( + name = "eigen", + build_file_content = + """ +cc_library( + name = 'eigen3', + srcs = [], + includes = ['.'], + hdrs = glob(['Eigen/**']), + visibility = ['//visibility:public'], +) +""", + commit = "3147391d946bb4b6c68edd901f2add6ac1f31f8c", + remote = "https://gitlab.com/libeigen/eigen.git", +) + +git_repository( + name = "highs", + branch = "bazel", + remote = "https://github.com/ERGO-Code/HiGHS.git", +) + +## Swig support +# pcre source code repository +new_git_repository( + name = "pcre2", + build_file = "@com_google_ortools//bazel:pcre2.BUILD", + remote = "https://github.com/PCRE2Project/pcre2.git", + tag = "pcre2-10.42", +) + +git_repository( + name = "com_google_ortools", + commit = "1d696f9108a0ebfd99feb73b9211e2f5a6b0812b", + remote = "https://github.com/google/or-tools.git", + shallow_since = "1647023481 +0100", +) diff --git a/bazel/import_llvm.bzl b/bazel/import_llvm.bzl index 132a6bc..af9a379 100644 --- a/bazel/import_llvm.bzl +++ b/bazel/import_llvm.bzl @@ -8,8 +8,8 @@ load( def import_llvm(name): """Imports LLVM.""" - # 2023-10-30 - LLVM_COMMIT = "896749aa0d420ae573255a64a349bc2a76cfed37" + # 2023-11-13 + LLVM_COMMIT = "f778eafdd878e8b11ad76f9e0a312ce7791a7481" new_git_repository( name = name, diff --git a/externals/llvm-project b/externals/llvm-project index 896749a..ebb8ffd 160000 --- a/externals/llvm-project +++ b/externals/llvm-project @@ -1 +1 @@ -Subproject commit 896749aa0d420ae573255a64a349bc2a76cfed37 +Subproject commit ebb8ffde94476333cb2e95aacfd2e023d7112efb diff --git a/lib/Analysis/CMakeLists.txt b/lib/Analysis/CMakeLists.txt new file mode 100644 index 0000000..a52d48b --- /dev/null +++ b/lib/Analysis/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(ReduceNoiseAnalysis) diff --git a/lib/Analysis/ReduceNoiseAnalysis/BUILD b/lib/Analysis/ReduceNoiseAnalysis/BUILD new file mode 100644 index 0000000..90981d0 --- /dev/null +++ b/lib/Analysis/ReduceNoiseAnalysis/BUILD @@ -0,0 +1,16 @@ +package( + default_visibility = ["//visibility:public"], +) + +cc_library( + name = "ReduceNoiseAnalysis", + srcs = ["ReduceNoiseAnalysis.cpp"], + hdrs = ["ReduceNoiseAnalysis.h"], + deps = [ + "//lib/Dialect/Noisy", + "@com_google_ortools//ortools/base", + "@com_google_ortools//ortools/linear_solver", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:IR", + ], +) diff --git a/lib/Analysis/ReduceNoiseAnalysis/CMakeLists.txt b/lib/Analysis/ReduceNoiseAnalysis/CMakeLists.txt new file mode 100644 index 0000000..1889df5 --- /dev/null +++ b/lib/Analysis/ReduceNoiseAnalysis/CMakeLists.txt @@ -0,0 +1,8 @@ +add_mlir_library(ReduceNoiseAnalysis + ReduceNoiseAnalysis.cpp + + ${PROJECT_SOURCE_DIR}/lib/Analysis/ReduceNoiseAnalysis/ + ADDITIONAL_HEADER_DIRS + LINK_LIBS PUBLIC + ortools::ortools +) diff --git a/lib/Analysis/ReduceNoiseAnalysis/ReduceNoiseAnalysis.cpp b/lib/Analysis/ReduceNoiseAnalysis/ReduceNoiseAnalysis.cpp new file mode 100644 index 0000000..fc75971 --- /dev/null +++ b/lib/Analysis/ReduceNoiseAnalysis/ReduceNoiseAnalysis.cpp @@ -0,0 +1,266 @@ +#include "lib/Analysis/ReduceNoiseAnalysis/ReduceNoiseAnalysis.h" + +#include + +#include "lib/Dialect/Noisy/NoisyOps.h" +#include "mlir/include/mlir/IR/Operation.h" +#include "mlir/include/mlir/IR/Value.h" +#include "ortools/linear_solver/linear_solver.h" +#include "llvm/Support/Debug.h" +#include "llvm/include/llvm/ADT/DenseMap.h" +#include "llvm/include/llvm/ADT/TypeSwitch.h" + +using namespace operations_research; + +namespace mlir { +namespace tutorial { + +#define DEBUG_TYPE "ReduceNoiseAnalysis" + +// This needs only be larger than 32, since we're hard coding i32s in this +// tutorial. +constexpr int IF_THEN_AUX = 100; + +std::string nameAndLoc(Operation *op) { + std::string varName; + llvm::raw_string_ostream ss(varName); + ss << op->getName() << "_" << op->getLoc(); + return ss.str(); +} + +ReduceNoiseAnalysis::ReduceNoiseAnalysis(Operation *op) { + std::unique_ptr solver(MPSolver::CreateSolver("SCIP")); + MPObjective *const objective = solver->MutableObjective(); + objective->SetMinimization(); + + llvm::DenseMap decisionVariables; + llvm::DenseMap ssaNoiseVariables; + std::vector allVariables; + + // First walk the IR to define variables for all values and ops, + // and constraint initial conditions. + op->walk([&](Operation *op) { + // FIXME: assumes all reduce_noise ops have already been removed and their + // values forwarded. + if (!llvm::isa(op)) { + return; + } + + std::string varName = "InsertReduceNoise_" + nameAndLoc(op); + auto decisionVar = solver->MakeIntVar(0, 1, varName); + decisionVariables.insert(std::make_pair(op, decisionVar)); + allVariables.push_back(decisionVar); + objective->SetCoefficient(decisionVar, 1); + + int index = 0; + for (auto operand : op->getOperands()) { + if (ssaNoiseVariables.contains(operand)) { + continue; + } + std::string varName = + "NoiseAt_" + nameAndLoc(op) + "_arg_" + std::to_string(index++); + auto ssaNoiseVar = solver->MakeNumVar(0, MAX_NOISE, varName); + allVariables.push_back(ssaNoiseVar); + ssaNoiseVariables.insert(std::make_pair(operand, ssaNoiseVar)); + } + + if (!ssaNoiseVariables.contains(op->getResult(0))) { + std::string varName = "NoiseAt_" + nameAndLoc(op) + "_result"; + auto ssaNoiseVar = solver->MakeNumVar(0, MAX_NOISE, varName); + allVariables.push_back(ssaNoiseVar); + ssaNoiseVariables.insert(std::make_pair(op->getResult(0), ssaNoiseVar)); + } + }); + + // Define constraints on the noise at each SSA value + for (auto item : ssaNoiseVariables) { + auto value = item.first; + auto var = item.second; + // An input node has noise equal to the initial noise, though we're being a + // bit sloppy by saying that EVERY block argument counts as an input node. + // In the tutorial, there is no control flow, so these are the function + // arguments of the main function being analyzed. A real compiler would + // need to handle this more generically. + if (value.isa() || + llvm::isa(value.getDefiningOp())) { + MPConstraint *const ct = + solver->MakeRowConstraint(INITIAL_NOISE, INITIAL_NOISE, ""); + ct->SetCoefficient(var, 1); + } + } + + std::string cstName; + // Define the decision variable constraints + op->walk([&](Operation *op) { + llvm::TypeSwitch(*op) + .Case([&](auto op) { + // result_noise = input_noise (1 - reduce_decision) + 12 * + // reduce_decision but linearized due to the quadratic term + // input_noise * reduce_decision + + auto inf = solver->infinity(); + auto lhsNoiseVar = ssaNoiseVariables.lookup(op.getLhs()); + auto rhsNoiseVar = ssaNoiseVariables.lookup(op.getRhs()); + auto resultNoiseVar = ssaNoiseVariables.lookup(op.getResult()); + auto reduceNoiseDecision = decisionVariables.lookup(op); + + // result_noise >= 12 * reduce_decision + cstName = "DecisionDynamics_" + nameAndLoc(op) + "_1"; + MPConstraint *const ct1 = + solver->MakeRowConstraint(0.0, inf, cstName); + ct1->SetCoefficient(resultNoiseVar, 1); + ct1->SetCoefficient(reduceNoiseDecision, -INITIAL_NOISE); + + // result_noise <= 12 + (1 - reduce_decision) * BIG_CONST + cstName = "DecisionDynamics_" + nameAndLoc(op) + "_2"; + MPConstraint *const ct2 = solver->MakeRowConstraint( + 0.0, INITIAL_NOISE * IF_THEN_AUX, cstName); + ct2->SetCoefficient(resultNoiseVar, 1); + ct2->SetCoefficient(reduceNoiseDecision, IF_THEN_AUX); + + // result_noise >= input_noise - reduce_decision * BIG_CONST + cstName = "DecisionDynamics_" + nameAndLoc(op) + "_3"; + MPConstraint *const ct3 = + solver->MakeRowConstraint(0.0, inf, cstName); + ct3->SetCoefficient(resultNoiseVar, 1); + ct3->SetCoefficient(reduceNoiseDecision, IF_THEN_AUX); + // The input noise is the sum of the two argument noises + if (op.getLhs() == op.getRhs()) { + ct3->SetCoefficient(lhsNoiseVar, -2); + } else { + ct3->SetCoefficient(lhsNoiseVar, -1); + ct3->SetCoefficient(rhsNoiseVar, -1); + } + + // result_noise <= input_noise + reduce_decision * BIG_CONST + cstName = "DecisionDynamics_" + nameAndLoc(op) + "_4"; + MPConstraint *const ct4 = + solver->MakeRowConstraint(-inf, 0.0, cstName); + ct4->SetCoefficient(resultNoiseVar, 1); + ct4->SetCoefficient(reduceNoiseDecision, -IF_THEN_AUX); + if (op.getLhs() == op.getRhs()) { + ct4->SetCoefficient(lhsNoiseVar, -2); + } else { + ct4->SetCoefficient(lhsNoiseVar, -1); + ct4->SetCoefficient(rhsNoiseVar, -1); + } + + // ensure the noise before the reduce_noise op (input_noise) + // also is not too large + cstName = "DecisionDynamics_" + nameAndLoc(op) + "_5"; + MPConstraint *const ct5 = + solver->MakeRowConstraint(0.0, MAX_NOISE, cstName); + if (op.getLhs() == op.getRhs()) { + ct5->SetCoefficient(lhsNoiseVar, 2); + } else { + ct5->SetCoefficient(lhsNoiseVar, 1); + ct5->SetCoefficient(rhsNoiseVar, 1); + } + }) + .Case([&](auto op) { + // Same as for MulOp, but the noise combination function is more + // complicated because it involves a maximum. + auto inf = solver->infinity(); + auto lhsNoiseVar = ssaNoiseVariables.lookup(op.getLhs()); + auto rhsNoiseVar = ssaNoiseVariables.lookup(op.getRhs()); + auto resultNoiseVar = ssaNoiseVariables.lookup(op.getResult()); + auto reduceNoiseDecision = decisionVariables.lookup(op); + + // result_noise >= 12 * reduce_decision + cstName = "DecisionDynamics_" + nameAndLoc(op) + "_1"; + MPConstraint *const ct1 = + solver->MakeRowConstraint(0.0, inf, cstName); + ct1->SetCoefficient(resultNoiseVar, 1); + ct1->SetCoefficient(reduceNoiseDecision, -INITIAL_NOISE); + + // result_noise <= 12 + (1 - reduce_decision) * BIG_CONST + cstName = "DecisionDynamics_" + nameAndLoc(op) + "_2"; + MPConstraint *const ct2 = solver->MakeRowConstraint( + 0.0, INITIAL_NOISE * IF_THEN_AUX, cstName); + ct2->SetCoefficient(resultNoiseVar, 1); + ct2->SetCoefficient(reduceNoiseDecision, IF_THEN_AUX); + + // for AddOp, the input noise is the max of the two argument noises + // plus one. Model this with an extra variable Z and two constraints: + // + // lhs_noise + 1 <= Z <= MAX_NOISE + // rhs_noise + 1 <= Z <= MAX_NOISE + // input_noise := Z + // + // Then add theze Z variables to the minimization objective, and + // they will be clamped to the larger of the two lower bounds. + cstName = "Z_" + nameAndLoc(op); + auto zVar = solver->MakeNumVar(0, MAX_NOISE, cstName); + allVariables.push_back(zVar); + // The objective coefficient is not all that important: the solver + // cannot cheat by making Z larger than necessary, since making Z + // larger than it needs to be would further increase the need to + // insert reduce_noise ops, which would be more expensive. + objective->SetCoefficient(zVar, 0.1); + + cstName = "DecisionDynamics_" + nameAndLoc(op) + "_z1"; + MPConstraint *const zCt1 = + solver->MakeRowConstraint(1.0, inf, cstName); + zCt1->SetCoefficient(zVar, 1); + zCt1->SetCoefficient(lhsNoiseVar, -1); + + if (op.getLhs() != op.getRhs()) { + cstName = "DecisionDynamics_" + nameAndLoc(op) + "_z2"; + MPConstraint *const zCt2 = + solver->MakeRowConstraint(1.0, inf, cstName); + zCt2->SetCoefficient(zVar, 1); + zCt2->SetCoefficient(rhsNoiseVar, -1); + } + + // result_noise >= input_noise - reduce_decision * BIG_CONST + cstName = "DecisionDynamics_" + nameAndLoc(op) + "_3"; + MPConstraint *const ct3 = + solver->MakeRowConstraint(0.0, inf, cstName); + ct3->SetCoefficient(resultNoiseVar, 1); + ct3->SetCoefficient(reduceNoiseDecision, IF_THEN_AUX); + ct3->SetCoefficient(zVar, -1); + + // result_noise <= input_noise + reduce_decision * BIG_CONST + cstName = "DecisionDynamics_" + nameAndLoc(op) + "_4"; + MPConstraint *const ct4 = + solver->MakeRowConstraint(-inf, 0.0, cstName); + ct4->SetCoefficient(resultNoiseVar, 1); + ct4->SetCoefficient(reduceNoiseDecision, -IF_THEN_AUX); + ct4->SetCoefficient(zVar, -1); + + // ensure the noise before the reduce_noise op (input_noise) + // also is not too large + cstName = "DecisionDynamics_" + nameAndLoc(op) + "_5"; + MPConstraint *const ct5 = + solver->MakeRowConstraint(0.0, MAX_NOISE, cstName); + ct5->SetCoefficient(zVar, 1); + }); + }); + + // Uncomment if you want to read the model's textual description, + // generally not for those unfamiliar with linear programming. + // std::string modelAsString; + // solver->ExportModelAsLpFormat(false, &modelAsString); + // LLVM_DEBUG(llvm::dbgs() << "Model string = " << modelAsString << "\n"); + + solver->Solve(); + LLVM_DEBUG(llvm::dbgs() << "Problem solved in " << solver->wall_time() + << " milliseconds" + << "\n"); + + LLVM_DEBUG(llvm::dbgs() << "Solution:\n"); + LLVM_DEBUG(llvm::dbgs() << "Objective value = " << objective->Value() + << "\n"); + // LLVM_DEBUG(llvm::dbgs() << "Variables:\n"); + // for (auto var : allVariables) { + // LLVM_DEBUG(llvm::dbgs() << " " << var->name() << " = " + // << var->solution_value() << "\n"); + // } + + for (auto item : decisionVariables) { + solution.insert(std::make_pair(item.first, item.second->solution_value())); + } +} + +} // namespace tutorial +} // namespace mlir diff --git a/lib/Analysis/ReduceNoiseAnalysis/ReduceNoiseAnalysis.h b/lib/Analysis/ReduceNoiseAnalysis/ReduceNoiseAnalysis.h new file mode 100644 index 0000000..efffca1 --- /dev/null +++ b/lib/Analysis/ReduceNoiseAnalysis/ReduceNoiseAnalysis.h @@ -0,0 +1,29 @@ +#ifndef LIB_ANALYSIS_REDUCENOISEANALYSIS_REDUCENOISEANALYSIS_H_ +#define LIB_ANALYSIS_REDUCENOISEANALYSIS_REDUCENOISEANALYSIS_H_ + +#include "llvm/include/llvm/ADT/DenseMap.h" // from @llvm-project +#include "mlir/include/mlir/IR/Operation.h" // from @llvm-project +#include "mlir/include/mlir/IR/Value.h" // from @llvm-project + +namespace mlir { +namespace tutorial { + +class ReduceNoiseAnalysis { + public: + ReduceNoiseAnalysis(Operation *op); + ~ReduceNoiseAnalysis() = default; + + /// Return true if a reduce_noise op should be inserted after the given + /// operation, according to the solution to the optimization problem. + bool shouldInsertReduceNoise(Operation *op) const { + return solution.lookup(op); + } + + private: + llvm::DenseMap solution; +}; + +} // namespace tutorial +} // namespace mlir + +#endif // LIB_ANALYSIS_REDUCENOISEANALYSIS_REDUCENOISEANALYSIS_H_ diff --git a/lib/CMakeLists.txt b/lib/CMakeLists.txt index 714031d..c1572fe 100644 --- a/lib/CMakeLists.txt +++ b/lib/CMakeLists.txt @@ -1,3 +1,4 @@ -add_subdirectory(Dialect) +add_subdirectory(Analysis) add_subdirectory(Conversion) +add_subdirectory(Dialect) add_subdirectory(Transform) diff --git a/lib/Dialect/CMakeLists.txt b/lib/Dialect/CMakeLists.txt index aff4d19..1f27a0f 100644 --- a/lib/Dialect/CMakeLists.txt +++ b/lib/Dialect/CMakeLists.txt @@ -1 +1,2 @@ -add_subdirectory(Poly) \ No newline at end of file +add_subdirectory(Noisy) +add_subdirectory(Poly) diff --git a/lib/Dialect/Noisy/BUILD b/lib/Dialect/Noisy/BUILD new file mode 100644 index 0000000..2a93043 --- /dev/null +++ b/lib/Dialect/Noisy/BUILD @@ -0,0 +1,105 @@ +load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library") + +package( + default_visibility = ["//visibility:public"], +) + +td_library( + name = "td_files", + srcs = [ + "NoisyDialect.td", + "NoisyOps.td", + "NoisyTypes.td", + ], + deps = [ + "@llvm-project//mlir:BuiltinDialectTdFiles", + "@llvm-project//mlir:InferIntRangeInterfaceTdFiles", + "@llvm-project//mlir:InferTypeOpInterfaceTdFiles", + "@llvm-project//mlir:OpBaseTdFiles", + "@llvm-project//mlir:SideEffectInterfacesTdFiles", + ], +) + +gentbl_cc_library( + name = "dialect_inc_gen", + tbl_outs = [ + ( + ["-gen-dialect-decls"], + "NoisyDialect.h.inc", + ), + ( + ["-gen-dialect-defs"], + "NoisyDialect.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "NoisyDialect.td", + deps = [ + ":td_files", + ], +) + +gentbl_cc_library( + name = "types_inc_gen", + tbl_outs = [ + ( + ["-gen-typedef-decls"], + "NoisyTypes.h.inc", + ), + ( + ["-gen-typedef-defs"], + "NoisyTypes.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "NoisyTypes.td", + deps = [ + ":dialect_inc_gen", + ":td_files", + ], +) + +gentbl_cc_library( + name = "ops_inc_gen", + tbl_outs = [ + ( + ["-gen-op-decls"], + "NoisyOps.h.inc", + ), + ( + ["-gen-op-defs"], + "NoisyOps.cpp.inc", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "NoisyOps.td", + deps = [ + ":dialect_inc_gen", + ":td_files", + ":types_inc_gen", + ], +) + +cc_library( + name = "Noisy", + srcs = [ + "NoisyDialect.cpp", + "NoisyOps.cpp", + ], + hdrs = [ + "NoisyDialect.h", + "NoisyOps.h", + "NoisyTypes.h", + ], + deps = [ + ":dialect_inc_gen", + ":ops_inc_gen", + ":types_inc_gen", + "@llvm-project//mlir:ComplexDialect", + "@llvm-project//mlir:Dialect", + "@llvm-project//mlir:IR", + "@llvm-project//mlir:InferIntRangeInterface", + "@llvm-project//mlir:InferTypeOpInterface", + "@llvm-project//mlir:Support", + ], +) diff --git a/lib/Dialect/Noisy/CMakeLists.txt b/lib/Dialect/Noisy/CMakeLists.txt new file mode 100644 index 0000000..adb5864 --- /dev/null +++ b/lib/Dialect/Noisy/CMakeLists.txt @@ -0,0 +1,21 @@ +# Inlining `add_mlir_dialect(Noisy noisy)` commands so that +# we can custom name `*.inc` generated files. +set(LLVM_TARGET_DEFINITIONS NoisyOps.td) +mlir_tablegen(NoisyOps.h.inc -gen-op-decls) +mlir_tablegen(NoisyOps.cpp.inc -gen-op-defs) +mlir_tablegen(NoisyTypes.h.inc -gen-typedef-decls -typedefs-dialect=noisy) +mlir_tablegen(NoisyTypes.cpp.inc -gen-typedef-defs -typedefs-dialect=noisy) +mlir_tablegen(NoisyDialect.h.inc -gen-dialect-decls -dialect=noisy) +mlir_tablegen(NoisyDialect.cpp.inc -gen-dialect-defs -dialect=noisy) +add_public_tablegen_target(MLIRNoisyOpsIncGen) +add_dependencies(mlir-headers MLIRNoisyOpsIncGen) + +add_mlir_doc(NoisyDialect NoisyDialect Noisy/ -gen-dialect-doc) + +add_mlir_dialect_library(MLIRNoisy + NoisyDialect.cpp + NoisyOps.cpp + + ADDITIONAL_HEADER_DIRS + ${PROJECT_SOURCE_DIR}/lib/Dialect/Noisy + ) diff --git a/lib/Dialect/Noisy/NoisyDialect.cpp b/lib/Dialect/Noisy/NoisyDialect.cpp new file mode 100644 index 0000000..0a4f5a5 --- /dev/null +++ b/lib/Dialect/Noisy/NoisyDialect.cpp @@ -0,0 +1,31 @@ +#include "lib/Dialect/Noisy/NoisyDialect.h" + +#include "lib/Dialect/Noisy/NoisyOps.h" +#include "lib/Dialect/Noisy/NoisyTypes.h" +#include "mlir/include/mlir/IR/Builders.h" +#include "llvm/include/llvm/ADT/TypeSwitch.h" + +#include "lib/Dialect/Noisy/NoisyDialect.cpp.inc" +#define GET_TYPEDEF_CLASSES +#include "lib/Dialect/Noisy/NoisyTypes.cpp.inc" +#define GET_OP_CLASSES +#include "lib/Dialect/Noisy/NoisyOps.cpp.inc" + +namespace mlir { +namespace tutorial { +namespace noisy { + +void NoisyDialect::initialize() { + addTypes< +#define GET_TYPEDEF_LIST +#include "lib/Dialect/Noisy/NoisyTypes.cpp.inc" + >(); + addOperations< +#define GET_OP_LIST +#include "lib/Dialect/Noisy/NoisyOps.cpp.inc" + >(); +} + +} // namespace noisy +} // namespace tutorial +} // namespace mlir diff --git a/lib/Dialect/Noisy/NoisyDialect.h b/lib/Dialect/Noisy/NoisyDialect.h new file mode 100644 index 0000000..5098e1e --- /dev/null +++ b/lib/Dialect/Noisy/NoisyDialect.h @@ -0,0 +1,14 @@ +#ifndef LIB_DIALECT_NOISY_NOISYDIALECT_H_ +#define LIB_DIALECT_NOISY_NOISYDIALECT_H_ + +// Required because the .h.inc file refers to MLIR classes and does not itself +// have any includes. +#include "mlir/include/mlir/IR/DialectImplementation.h" + +#include "lib/Dialect/Noisy/NoisyDialect.h.inc" + + +constexpr int INITIAL_NOISE = 12; +constexpr int MAX_NOISE = 26; + +#endif // LIB_DIALECT_NOISY_NOISYDIALECT_H_ diff --git a/lib/Dialect/Noisy/NoisyDialect.td b/lib/Dialect/Noisy/NoisyDialect.td new file mode 100644 index 0000000..954e703 --- /dev/null +++ b/lib/Dialect/Noisy/NoisyDialect.td @@ -0,0 +1,15 @@ +#ifndef LIB_DIALECT_NOISY_NOISYDIALECT_TD_ +#define LIB_DIALECT_NOISY_NOISYDIALECT_TD_ + +include "mlir/IR/OpBase.td" + +def Noisy_Dialect : Dialect { + let name = "noisy"; + let summary = "A dialect for arithmetic on noisy i32s"; + + let cppNamespace = "::mlir::tutorial::noisy"; + + let useDefaultTypePrinterParser = 1; +} + +#endif // LIB_DIALECT_NOISY_NOISYDIALECT_TD_ diff --git a/lib/Dialect/Noisy/NoisyOps.cpp b/lib/Dialect/Noisy/NoisyOps.cpp new file mode 100644 index 0000000..97e16a2 --- /dev/null +++ b/lib/Dialect/Noisy/NoisyOps.cpp @@ -0,0 +1,50 @@ +#include "lib/Dialect/Noisy/NoisyOps.h" + +namespace mlir { +namespace tutorial { +namespace noisy { + +ConstantIntRanges initialNoiseRange() { + return ConstantIntRanges::fromUnsigned(APInt(32, 0), + APInt(32, INITIAL_NOISE)); +} + +ConstantIntRanges unionPlusOne(ArrayRef inputRanges) { + auto lhsRange = inputRanges[0]; + auto rhsRange = inputRanges[1]; + auto joined = lhsRange.rangeUnion(rhsRange); + return ConstantIntRanges::fromUnsigned(joined.umin(), joined.umax() + 1); +} + +void EncodeOp::inferResultRanges(ArrayRef inputRanges, + SetIntRangeFn setResultRange) { + setResultRange(getResult(), initialNoiseRange()); +} + +void AddOp::inferResultRanges(ArrayRef inputRanges, + SetIntRangeFn setResultRange) { + setResultRange(getResult(), unionPlusOne(inputRanges)); +} + +void SubOp::inferResultRanges(ArrayRef inputRanges, + SetIntRangeFn setResultRange) { + setResultRange(getResult(), unionPlusOne(inputRanges)); +} + +void MulOp::inferResultRanges(ArrayRef inputRanges, + SetIntRangeFn setResultRange) { + auto lhsRange = inputRanges[0]; + auto rhsRange = inputRanges[1]; + setResultRange(getResult(), ConstantIntRanges::fromUnsigned( + lhsRange.umin() + rhsRange.umin(), + lhsRange.umax() + rhsRange.umax())); +} + +void ReduceNoiseOp::inferResultRanges(ArrayRef inputRanges, + SetIntRangeFn setResultRange) { + setResultRange(getResult(), initialNoiseRange()); +} + +} // namespace noisy +} // namespace tutorial +} // namespace mlir diff --git a/lib/Dialect/Noisy/NoisyOps.h b/lib/Dialect/Noisy/NoisyOps.h new file mode 100644 index 0000000..12f83e4 --- /dev/null +++ b/lib/Dialect/Noisy/NoisyOps.h @@ -0,0 +1,15 @@ +#ifndef LIB_DIALECT_NOISY_NOISYOPS_H_ +#define LIB_DIALECT_NOISY_NOISYOPS_H_ + +#include "lib/Dialect/Noisy/NoisyDialect.h" +#include "lib/Dialect/Noisy/NoisyTypes.h" +#include "mlir/Interfaces/InferTypeOpInterface.h" +#include "mlir/Interfaces/InferIntRangeInterface.h" +#include "mlir/include/mlir/IR/BuiltinOps.h" +#include "mlir/include/mlir/IR/BuiltinTypes.h" +#include "mlir/include/mlir/IR/Dialect.h" + +#define GET_OP_CLASSES +#include "lib/Dialect/Noisy/NoisyOps.h.inc" + +#endif // LIB_DIALECT_NOISY_NOISYOPS_H_ diff --git a/lib/Dialect/Noisy/NoisyOps.td b/lib/Dialect/Noisy/NoisyOps.td new file mode 100644 index 0000000..735e2da --- /dev/null +++ b/lib/Dialect/Noisy/NoisyOps.td @@ -0,0 +1,58 @@ +#ifndef LIB_DIALECT_NOISY_NOISYOPS_TD_ +#define LIB_DIALECT_NOISY_NOISYOPS_TD_ + +include "NoisyDialect.td" +include "NoisyTypes.td" +include "mlir/IR/BuiltinAttributes.td" +include "mlir/IR/CommonTypeConstraints.td" +include "mlir/IR/OpBase.td" +include "mlir/Interfaces/InferIntRangeInterface.td" +include "mlir/Interfaces/InferTypeOpInterface.td" +include "mlir/Interfaces/SideEffectInterfaces.td" + +class Noisy_BinOp : Op +]> { + let arguments = (ins Noisy_I32:$lhs, Noisy_I32:$rhs); + let results = (outs Noisy_I32:$output); + let assemblyFormat = "$lhs `,` $rhs attr-dict `:` qualified(type($output))"; +} + +def Noisy_AddOp : Noisy_BinOp<"add"> { + let summary = "Addition operation between noisy ints. Adds noise."; +} + +def Noisy_SubOp : Noisy_BinOp<"sub"> { + let summary = "Subtraction operation between noisy ints. Adds noise."; +} + +def Noisy_MulOp : Noisy_BinOp<"mul"> { + let summary = "Multiplication operation between noisy ints. Multiplies noise."; +} + +def Noisy_EncodeOp : Op]> { + let summary = "Encodes a noisy i32 from a small-width integer, injecting 12 bits of noise."; + let arguments = (ins AnyIntOfWidths<[1, 2, 3, 4, 5]>:$input); + let results = (outs Noisy_I32:$output); + let assemblyFormat = "$input attr-dict `:` type($input) `->` qualified(type($output))"; +} + +def Noisy_DecodeOp : Op { + let summary = "Decodes a noisy integer to a regular integer, failing if the noise is too high."; + let arguments = (ins Noisy_I32:$input); + let results = (outs AnyIntOfWidths<[1, 2, 3, 4, 5]>:$output); + let assemblyFormat = "$input attr-dict `:` qualified(type($input)) `->` type($output)"; +} + +def Noisy_ReduceNoiseOp : Op]> { + let summary = "Reduces the noise in a noisy integer to a fixed noise level. Expensive!"; + let arguments = (ins Noisy_I32:$input); + let results = (outs Noisy_I32:$output); + let assemblyFormat = "$input attr-dict `:` qualified(type($output))"; +} + +#endif // LIB_DIALECT_NOISY_NOISYOPS_TD_ diff --git a/lib/Dialect/Noisy/NoisyTypes.h b/lib/Dialect/Noisy/NoisyTypes.h new file mode 100644 index 0000000..7b432c9 --- /dev/null +++ b/lib/Dialect/Noisy/NoisyTypes.h @@ -0,0 +1,9 @@ +#ifndef LIB_TYPES_NOISY_NOISYTYPES_H_ +#define LIB_TYPES_NOISY_NOISYTYPES_H_ + +#include "mlir/include/mlir/IR/DialectImplementation.h" + +#define GET_TYPEDEF_CLASSES +#include "lib/Dialect/Noisy/NoisyTypes.h.inc" + +#endif // LIB_TYPES_NOISY_NOISYTYPES_H_ diff --git a/lib/Dialect/Noisy/NoisyTypes.td b/lib/Dialect/Noisy/NoisyTypes.td new file mode 100644 index 0000000..2e65307 --- /dev/null +++ b/lib/Dialect/Noisy/NoisyTypes.td @@ -0,0 +1,15 @@ +#ifndef LIB_DIALECT_NOISY_NOISYTYPES_TD_ +#define LIB_DIALECT_NOISY_NOISYTYPES_TD_ + +include "NoisyDialect.td" +include "mlir/IR/AttrTypeBase.td" + +class Noisy_Type : TypeDef { + let mnemonic = typeMnemonic; +} + +def Noisy_I32 : Noisy_Type<"NoisyI32", "i32"> { + let summary = "A type for approximate 32-bit integers."; +} + +#endif // LIB_DIALECT_NOISY_NOISYTYPES_TD_ diff --git a/lib/Transform/Affine/CMakeLists.txt b/lib/Transform/Affine/CMakeLists.txt index eaa2b68..31a6c14 100644 --- a/lib/Transform/Affine/CMakeLists.txt +++ b/lib/Transform/Affine/CMakeLists.txt @@ -4,6 +4,10 @@ add_mlir_library(AffineFullUnroll ${PROJECT_SOURCE_DIR}/lib/Transform/Affine/ ADDITIONAL_HEADER_DIRS + + DEPENDS + MLIRAffineFullUnrollPasses + LINK_LIBS PUBLIC ) diff --git a/lib/Transform/Arith/CMakeLists.txt b/lib/Transform/Arith/CMakeLists.txt index f73be7c..5ae59e2 100644 --- a/lib/Transform/Arith/CMakeLists.txt +++ b/lib/Transform/Arith/CMakeLists.txt @@ -3,10 +3,14 @@ add_mlir_library(MulToAdd ${PROJECT_SOURCE_DIR}/lib/Transform/Arith/ ADDITIONAL_HEADER_DIRS + + DEPENDS + MLIRMulToAddPasses + LINK_LIBS PUBLIC ) set(LLVM_TARGET_DEFINITIONS Passes.td) mlir_tablegen(Passes.h.inc -gen-pass-decls -name Arith) add_public_tablegen_target(MLIRMulToAddPasses) -add_mlir_doc(Passes ArithPasses ./ -gen-pass-doc) \ No newline at end of file +add_mlir_doc(Passes ArithPasses ./ -gen-pass-doc) diff --git a/lib/Transform/CMakeLists.txt b/lib/Transform/CMakeLists.txt index 9caa085..5534bc3 100644 --- a/lib/Transform/CMakeLists.txt +++ b/lib/Transform/CMakeLists.txt @@ -1,2 +1,3 @@ add_subdirectory(Affine) -add_subdirectory(Arith) \ No newline at end of file +add_subdirectory(Arith) +add_subdirectory(Noisy) diff --git a/lib/Transform/Noisy/BUILD b/lib/Transform/Noisy/BUILD new file mode 100644 index 0000000..369c4f0 --- /dev/null +++ b/lib/Transform/Noisy/BUILD @@ -0,0 +1,57 @@ +# Passes that work with the Noisy dialect + +load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library") + +package( + default_visibility = ["//visibility:public"], +) + +gentbl_cc_library( + name = "pass_inc_gen", + tbl_outs = [ + ( + [ + "-gen-pass-decls", + "-name=Noisy", + ], + "Passes.h.inc", + ), + ( + ["-gen-pass-doc"], + "NoisyPasses.md", + ), + ], + tblgen = "@llvm-project//mlir:mlir-tblgen", + td_file = "Passes.td", + deps = [ + "@llvm-project//mlir:OpBaseTdFiles", + "@llvm-project//mlir:PassBaseTdFiles", + ], +) + +cc_library( + name = "ReduceNoiseOptimizer", + srcs = ["ReduceNoiseOptimizer.cpp"], + hdrs = [ + "Passes.h", + "ReduceNoiseOptimizer.h", + ], + deps = [ + ":pass_inc_gen", + "//lib/Analysis/ReduceNoiseAnalysis", + "//lib/Dialect/Noisy", + "@llvm-project//llvm:Support", + "@llvm-project//mlir:Analysis", + "@llvm-project//mlir:Pass", + "@llvm-project//mlir:Transforms", + ], +) + +cc_library( + name = "Passes", + hdrs = ["Passes.h"], + deps = [ + ":ReduceNoiseOptimizer", + ":pass_inc_gen", + ], +) diff --git a/lib/Transform/Noisy/CMakeLists.txt b/lib/Transform/Noisy/CMakeLists.txt new file mode 100644 index 0000000..94568d8 --- /dev/null +++ b/lib/Transform/Noisy/CMakeLists.txt @@ -0,0 +1,18 @@ +add_mlir_library(NoisyPasses + ReduceNoiseOptimizer.cpp + + ${PROJECT_SOURCE_DIR}/lib/Transform/Noisy/ + ADDITIONAL_HEADER_DIRS + + DEPENDS + MLIRNoisy + MLIRNoisyPasses + + LINK_LIBS PUBLIC + ReduceNoiseAnalysis +) + +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls -name Noisy) +add_public_tablegen_target(MLIRNoisyPasses) +add_mlir_doc(Passes NoisyPasses ./ -gen-pass-doc) diff --git a/lib/Transform/Noisy/Passes.h b/lib/Transform/Noisy/Passes.h new file mode 100644 index 0000000..6711b1f --- /dev/null +++ b/lib/Transform/Noisy/Passes.h @@ -0,0 +1,17 @@ +#ifndef LIB_TRANSFORM_NOISY_PASSES_H_ +#define LIB_TRANSFORM_NOISY_PASSES_H_ + +#include "lib/Transform/Noisy/ReduceNoiseOptimizer.h" + +namespace mlir { +namespace tutorial { +namespace noisy { + +#define GEN_PASS_REGISTRATION +#include "lib/Transform/Noisy/Passes.h.inc" + +} // namespace noisy +} // namespace tutorial +} // namespace mlir + +#endif // LIB_TRANSFORM_NOISY_PASSES_H_ diff --git a/lib/Transform/Noisy/Passes.td b/lib/Transform/Noisy/Passes.td new file mode 100644 index 0000000..a25b53d --- /dev/null +++ b/lib/Transform/Noisy/Passes.td @@ -0,0 +1,15 @@ +#ifndef LIB_TRANSFORM_NOISY_PASSES_TD_ +#define LIB_TRANSFORM_NOISY_PASSES_TD_ + +include "mlir/Pass/PassBase.td" + +def ReduceNoiseOptimizer : Pass<"noisy-reduce-noise-optimizer"> { + let summary = "Insert reduce_noise ops optimally"; + let description = [{ + Solves an integer linear program to select the optimal locations in the IR + to insert `reduce_noise` ops. + }]; + let dependentDialects = ["mlir::tutorial::noisy::NoisyDialect"]; +} + +#endif // LIB_TRANSFORM_NOISY_PASSES_TD_ diff --git a/lib/Transform/Noisy/ReduceNoiseOptimizer.cpp b/lib/Transform/Noisy/ReduceNoiseOptimizer.cpp new file mode 100644 index 0000000..05b4088 --- /dev/null +++ b/lib/Transform/Noisy/ReduceNoiseOptimizer.cpp @@ -0,0 +1,89 @@ +#include "lib/Transform/Noisy/ReduceNoiseOptimizer.h" + +#include "lib/Analysis/ReduceNoiseAnalysis/ReduceNoiseAnalysis.h" +#include "lib/Dialect/Noisy/NoisyOps.h" +#include "lib/Dialect/Noisy/NoisyTypes.h" +#include "mlir/include/mlir/Analysis/DataFlow/DeadCodeAnalysis.h" +#include "mlir/include/mlir/Analysis/DataFlow/IntegerRangeAnalysis.h" +#include "mlir/include/mlir/Analysis/DataFlowFramework.h" +#include "mlir/include/mlir/IR/Visitors.h" +#include "mlir/include/mlir/Pass/Pass.h" + +namespace mlir { +namespace tutorial { +namespace noisy { + +#define GEN_PASS_DEF_REDUCENOISEOPTIMIZER +#include "lib/Transform/Noisy/Passes.h.inc" + +struct ReduceNoiseOptimizer + : impl::ReduceNoiseOptimizerBase { + using ReduceNoiseOptimizerBase::ReduceNoiseOptimizerBase; + + void runOnOperation() { + Operation *module = getOperation(); + + // FIXME: Should have some way to mark failure when solver is infeasible + ReduceNoiseAnalysis analysis(module); + OpBuilder b(&getContext()); + + module->walk([&](Operation *op) { + if (!analysis.shouldInsertReduceNoise(op)) + return; + + b.setInsertionPointAfter(op); + auto reduceOp = b.create(op->getLoc(), op->getResult(0)); + op->getResult(0).replaceAllUsesExcept(reduceOp.getResult(), {reduceOp}); + }); + + // Use the int range analysis to confirm the noise is always below the + // maximum. + DataFlowSolver solver; + // The IntegerRangeAnalysis depends on DeadCodeAnalysis, but this + // dependence is not automatic and fails silently. + solver.load(); + solver.load(); + if (failed(solver.initializeAndRun(module))) { + getOperation()->emitOpError() << "Failed to run the analysis.\n"; + signalPassFailure(); + return; + } + + auto result = module->walk([&](Operation *op) { + if (!llvm::isa(*op)) { + return WalkResult::advance(); + } + const dataflow::IntegerValueRangeLattice *opRange = + solver.lookupState( + op->getResult(0)); + if (!opRange || opRange->getValue().isUninitialized()) { + op->emitOpError() + << "Found op without a set integer range; did the analysis fail?"; + return WalkResult::interrupt(); + } + + ConstantIntRanges range = opRange->getValue().getValue(); + if (range.umax().getZExtValue() > MAX_NOISE) { + op->emitOpError() << "Found op after which the noise exceeds the " + "allowable maximum of " + << MAX_NOISE + << "; it was: " << range.umax().getZExtValue() + << "\n"; + return WalkResult::interrupt(); + } + + return WalkResult::advance(); + }); + + if (result.wasInterrupted()) { + getOperation()->emitOpError() + << "Detected error in the noise analysis.\n"; + signalPassFailure(); + } + } +}; + +} // namespace noisy +} // namespace tutorial +} // namespace mlir diff --git a/lib/Transform/Noisy/ReduceNoiseOptimizer.h b/lib/Transform/Noisy/ReduceNoiseOptimizer.h new file mode 100644 index 0000000..b3afd84 --- /dev/null +++ b/lib/Transform/Noisy/ReduceNoiseOptimizer.h @@ -0,0 +1,17 @@ +#ifndef LIB_TRANSFORM_NOISY_REDUCENOISEOPTIMIZER_H_ +#define LIB_TRANSFORM_NOISY_REDUCENOISEOPTIMIZER_H_ + +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace tutorial { +namespace noisy { + +#define GEN_PASS_DECL_REDUCENOISEOPTIMIZER +#include "lib/Transform/Noisy/Passes.h.inc" + +} // namespace noisy +} // namespace tutorial +} // namespace mlir + +#endif // LIB_TRANSFORM_NOISY_REDUCENOISEOPTIMIZER_H_ diff --git a/tests/noisy_reduce_noise.mlir b/tests/noisy_reduce_noise.mlir new file mode 100644 index 0000000..bd30a0f --- /dev/null +++ b/tests/noisy_reduce_noise.mlir @@ -0,0 +1,160 @@ +// RUN: tutorial-opt %s --noisy-reduce-noise-optimizer | FileCheck %s +// Check for syntax + +// CHECK-LABEL: test_insert_noise_reduction_ops_mul +// CHECK: [[V0:%.*]] = arith.constant 3 +// CHECK-NEXT: [[V1:%.*]] = arith.constant 4 +// CHECK-NEXT: [[V2:%.*]] = noisy.encode [[V0]] +// CHECK-NEXT: [[V3:%.*]] = noisy.encode [[V1]] +// CHECK-NEXT: [[V4:%.*]] = noisy.mul [[V2]], [[V3]] +// CHECK-NEXT: [[V4_R:%.*]] = noisy.reduce_noise [[V4]] +// CHECK-NEXT: [[V5:%.*]] = noisy.mul [[V4_R]], [[V4_R]] +// CHECK-NEXT: [[V5_R:%.*]] = noisy.reduce_noise [[V5]] +// CHECK-NEXT: [[V6:%.*]] = noisy.mul [[V5_R]], [[V5_R]] +// CHECK-NEXT: [[V6_R:%.*]] = noisy.reduce_noise [[V6]] +// This last mul does not need to be reduced +// CHECK-NEXT: [[V7:%.*]] = noisy.mul [[V6_R]], [[V6_R]] +// CHECK-NEXT: [[V8:%.*]] = noisy.decode [[V7]] +// CHECK-NEXT: return +func.func @test_insert_noise_reduction_ops_mul() -> i5 { + %0 = arith.constant 3 : i5 + %1 = arith.constant 4 : i5 + %2 = noisy.encode %0 : i5 -> !noisy.i32 + %3 = noisy.encode %1 : i5 -> !noisy.i32 + %4 = noisy.mul %2, %3 : !noisy.i32 + %5 = noisy.mul %4, %4 : !noisy.i32 + %6 = noisy.mul %5, %5 : !noisy.i32 + %7 = noisy.mul %6, %6 : !noisy.i32 + %8 = noisy.decode %7 : !noisy.i32 -> i5 + return %8 : i5 +} + +// CHECK-LABEL: test_insert_noise_reduction_ops_add_none_needed +// CHECK-NOT: noisy.reduce_noise +func.func @test_insert_noise_reduction_ops_add_none_needed() -> i5 { + %0 = arith.constant 3 : i5 + %1 = arith.constant 4 : i5 + %2 = noisy.encode %0 : i5 -> !noisy.i32 + %3 = noisy.encode %1 : i5 -> !noisy.i32 + %4 = noisy.add %2, %3 : !noisy.i32 + %5 = noisy.add %4, %4 : !noisy.i32 + %6 = noisy.add %5, %5 : !noisy.i32 + %7 = noisy.add %6, %6 : !noisy.i32 + %8 = noisy.decode %7 : !noisy.i32 -> i5 + return %8 : i5 +} + + +// CHECK-LABEL: test_add_after_mul +// CHECK: noisy.mul +// CHECK: noisy.reduce_noise +// CHECK: noisy.add +// CHECK: noisy.add +// CHECK: noisy.add +// CHECK: noisy.add +// CHECK: noisy.decode +// CHECK: return +func.func @test_add_after_mul() -> i5 { + %0 = arith.constant 3 : i5 + %1 = arith.constant 4 : i5 + %2 = noisy.encode %0 : i5 -> !noisy.i32 + %3 = noisy.encode %1 : i5 -> !noisy.i32 + // Noise: 12 + %4 = noisy.mul %2, %3 : !noisy.i32 + // Noise: 24 + %5 = noisy.add %4, %3 : !noisy.i32 + // Noise: 25 + %6 = noisy.add %5, %5 : !noisy.i32 + // Noise: 26 + %7 = noisy.add %6, %6 : !noisy.i32 + // Noise: 27 + %8 = noisy.add %7, %7 : !noisy.i32 + %9 = noisy.decode %8 : !noisy.i32 -> i5 + return %9 : i5 +} + +// This test checks that the solver can find a single insertion point +// for a reduce_noise op that handles two branches, each of which would +// also need a reduce_noise op if handled separately. +// CHECK-LABEL: test_single_insertion_branching +// CHECK: noisy.mul +// CHECK-NOT: noisy.add +// CHECK-COUNT-1: noisy.reduce_noise +// CHECK-NOT: noisy.reduce_noise +func.func @test_single_insertion_branching() -> i5 { + %0 = arith.constant 3 : i5 + %1 = arith.constant 4 : i5 + %2 = noisy.encode %0 : i5 -> !noisy.i32 + %3 = noisy.encode %1 : i5 -> !noisy.i32 + // Noise: 12 + %4 = noisy.mul %2, %3 : !noisy.i32 + // Noise: 24 + + // branch 1 + %b1 = noisy.add %4, %3 : !noisy.i32 + // Noise: 25 + %b2 = noisy.add %b1, %3 : !noisy.i32 + // Noise: 25 + %b3 = noisy.add %b2, %3 : !noisy.i32 + // Noise: 26 + %b4 = noisy.add %b3, %3 : !noisy.i32 + // Noise: 27 + + // branch 2 + %c1 = noisy.sub %4, %2 : !noisy.i32 + // Noise: 25 + %c2 = noisy.sub %c1, %3 : !noisy.i32 + // Noise: 25 + %c3 = noisy.sub %c2, %3 : !noisy.i32 + // Noise: 26 + %c4 = noisy.sub %c3, %3 : !noisy.i32 + // Noise: 27 + + %x1 = noisy.decode %b4 : !noisy.i32 -> i5 + %x2 = noisy.decode %c4 : !noisy.i32 -> i5 + %x3 = arith.addi %x1, %x2 : i5 + return %x3 : i5 +} + +// same as test_single_insertion_branching, but because the last two values +// are multiplied, we need two reduce_noise ops, one on each branch. +// CHECK-LABEL: test_double_insertion_branching +// CHECK: noisy.mul +// CHECK: noisy.add +// CHECK-COUNT-2: noisy.reduce_noise +// CHECK-NOT: noisy.reduce_noise +// CHECK: noisy.mul +func.func @test_double_insertion_branching() -> i5 { + %0 = arith.constant 3 : i5 + %1 = arith.constant 4 : i5 + %2 = noisy.encode %0 : i5 -> !noisy.i32 + %3 = noisy.encode %1 : i5 -> !noisy.i32 + // Noise: 12 + %4 = noisy.mul %2, %3 : !noisy.i32 + // Noise: 24 + + // branch 1 + %b1 = noisy.add %4, %3 : !noisy.i32 + // Noise: 25 + %b2 = noisy.add %b1, %3 : !noisy.i32 + // Noise: 25 + %b3 = noisy.add %b2, %3 : !noisy.i32 + // Noise: 26 + %b4 = noisy.add %b3, %3 : !noisy.i32 + // Noise: 27 + + // branch 2 + %c1 = noisy.add %4, %2 : !noisy.i32 + // Noise: 25 + %c2 = noisy.add %c1, %3 : !noisy.i32 + // Noise: 25 + %c3 = noisy.add %c2, %3 : !noisy.i32 + // Noise: 26 + %c4 = noisy.add %c3, %3 : !noisy.i32 + // Noise: 27 + + %exit = noisy.mul %b4, %c4 : !noisy.i32 + + %x1 = noisy.decode %exit : !noisy.i32 -> i5 + return %x1 : i5 +} diff --git a/tests/noisy_syntax.mlir b/tests/noisy_syntax.mlir new file mode 100644 index 0000000..04f6dc3 --- /dev/null +++ b/tests/noisy_syntax.mlir @@ -0,0 +1,15 @@ +// RUN: tutorial-opt %s | FileCheck %s +// Check for syntax + +// CHECK-LABEL: test_op_syntax +func.func @test_op_syntax() { + %0 = arith.constant 3 : i5 + %1 = arith.constant 4 : i5 + %2 = noisy.encode %0 : i5 -> !noisy.i32 + %3 = noisy.encode %1 : i5 -> !noisy.i32 + %4 = noisy.add %2, %3 : !noisy.i32 + %5 = noisy.mul %4, %4 : !noisy.i32 + %6 = noisy.reduce_noise %5 : !noisy.i32 + %7 = noisy.decode %6 : !noisy.i32 -> i5 + return +} diff --git a/tools/BUILD b/tools/BUILD index 27d82e9..d4397b0 100644 --- a/tools/BUILD +++ b/tools/BUILD @@ -12,9 +12,11 @@ cc_binary( includes = ["include"], deps = [ "//lib/Conversion/PolyToStandard", + "//lib/Dialect/Noisy", "//lib/Dialect/Poly", "//lib/Transform/Affine:Passes", "//lib/Transform/Arith:Passes", + "//lib/Transform/Noisy:Passes", "@llvm-project//mlir:AllPassesAndDialects", "@llvm-project//mlir:ArithToLLVM", "@llvm-project//mlir:BufferizationPipelines", diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index 59f33b1..8cd5559 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -4,12 +4,15 @@ get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS) set (LIBS ${dialect_libs} ${conversion_libs} - MLIRPoly AffineFullUnroll - MulToAdd + MLIRNoisy MLIROptLib MLIRPass + MLIRPoly + MulToAdd + NoisyPasses PolyToStandard + ortools::ortools ) add_llvm_executable(tutorial-opt tutorial-opt.cpp) diff --git a/tools/tutorial-opt.cpp b/tools/tutorial-opt.cpp index 91f4a95..87f213d 100644 --- a/tools/tutorial-opt.cpp +++ b/tools/tutorial-opt.cpp @@ -1,7 +1,9 @@ #include "lib/Conversion/PolyToStandard/PolyToStandard.h" +#include "lib/Dialect/Noisy/NoisyDialect.h" #include "lib/Dialect/Poly/PolyDialect.h" #include "lib/Transform/Affine/Passes.h" #include "lib/Transform/Arith/Passes.h" +#include "lib/Transform/Noisy/Passes.h" #include "mlir/include/mlir/Conversion/ArithToLLVM/ArithToLLVM.h" #include "mlir/include/mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h" #include "mlir/include/mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h" @@ -16,7 +18,6 @@ #include "mlir/include/mlir/Pass/PassRegistry.h" #include "mlir/include/mlir/Tools/mlir-opt/MlirOptMain.h" #include "mlir/include/mlir/Transforms/Passes.h" -// void polyToLLVMPipelineBuilder(mlir::OpPassManager &manager) { // Poly @@ -33,7 +34,8 @@ void polyToLLVMPipelineBuilder(mlir::OpPassManager &manager) { manager.addPass( mlir::bufferization::createOneShotBufferizePass(bufferizationOptions)); mlir::bufferization::BufferDeallocationPipelineOptions deallocationOptions; - mlir::bufferization::buildBufferDeallocationPipeline(manager, deallocationOptions); + mlir::bufferization::buildBufferDeallocationPipeline(manager, + deallocationOptions); manager.addPass(mlir::createConvertLinalgToLoopsPass()); @@ -57,18 +59,20 @@ void polyToLLVMPipelineBuilder(mlir::OpPassManager &manager) { int main(int argc, char **argv) { mlir::DialectRegistry registry; registry.insert(); + registry.insert(); mlir::registerAllDialects(registry); mlir::registerAllPasses(); mlir::tutorial::registerAffinePasses(); mlir::tutorial::registerArithPasses(); + mlir::tutorial::noisy::registerNoisyPasses(); // Dialect conversion passes mlir::tutorial::poly::registerPolyToStandardPasses(); - mlir::PassPipelineRegistration<>("poly-to-llvm", - "Run passes to lower the poly dialect to LLVM", - polyToLLVMPipelineBuilder); + mlir::PassPipelineRegistration<>( + "poly-to-llvm", "Run passes to lower the poly dialect to LLVM", + polyToLLVMPipelineBuilder); return mlir::asMainReturnCode( mlir::MlirOptMain(argc, argv, "Tutorial Pass Driver", registry));