feat: Add Linear-SiLU fusion pass and integrate into BuddyMlp example

examples/BuddyMlp/CMakeLists.txt

@@ -0,0 +1,79 @@
+add_custom_command(
+  OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/forward.mlir
+         ${CMAKE_CURRENT_BINARY_DIR}/subgraph0.mlir
+         ${CMAKE_CURRENT_BINARY_DIR}/arg0.data
+  COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/buddy-mlp-import.py
+          --output-dir ${CMAKE_CURRENT_BINARY_DIR}
+  COMMENT "Generating forward.mlir, subgraph0.mlir and parameter files"
+)
+
+
+add_custom_command(
+  OUTPUT forward.o
+  COMMAND ${LLVM_TOOLS_BINARY_DIR}/mlir-opt ${CMAKE_CURRENT_BINARY_DIR}/forward.mlir
+            -pass-pipeline "builtin.module(func.func(tosa-to-linalg-named, tosa-to-linalg, tosa-to-tensor, tosa-to-arith), empty-tensor-to-alloc-tensor, convert-elementwise-to-linalg)" |
+          ${LLVM_TOOLS_BINARY_DIR}/mlir-opt
+            -pass-pipeline "builtin.module(func.func(buffer-deallocation-simplification, convert-linalg-to-loops), eliminate-empty-tensors, func.func(llvm-request-c-wrappers),convert-math-to-llvm, convert-math-to-libm, convert-scf-to-cf,  convert-arith-to-llvm, expand-strided-metadata, finalize-memref-to-llvm, convert-func-to-llvm, reconcile-unrealized-casts)" |
+          ${LLVM_TOOLS_BINARY_DIR}/mlir-translate -mlir-to-llvmir |
+          ${LLVM_TOOLS_BINARY_DIR}/llvm-as |
+          ${LLVM_TOOLS_BINARY_DIR}/llc -filetype=obj  -relocation-model=pic -O0 -o ${CMAKE_CURRENT_BINARY_DIR}/forward.o
+  DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/forward.mlir
+  COMMENT "Building forward.o"
+  VERBATIM)
+
+add_custom_command(
+  OUTPUT subgraph0.o
+  COMMAND ${LLVM_TOOLS_BINARY_DIR}/mlir-opt ${CMAKE_CURRENT_BINARY_DIR}/subgraph0.mlir
+            -pass-pipeline "builtin.module(func.func(tosa-to-linalg-named, tosa-to-linalg, tosa-to-tensor, tosa-to-arith))" |
+          ${BUDDY_BINARY_DIR}/buddy-opt
+            -linear-silu-fusion
+            -linear-silu-lower
+            -arith-expand 
+            -eliminate-empty-tensors 
+            -empty-tensor-to-alloc-tensor 
+            -one-shot-bufferize="bufferize-function-boundaries" 
+            -convert-linalg-to-affine-loops 
+            -affine-loop-fusion 
+            -buffer-deallocation 
+            -convert-vector-to-scf 
+            -expand-strided-metadata 
+            -lower-affine 
+            -convert-vector-to-llvm 
+            -convert-arith-to-llvm 
+            -finalize-memref-to-llvm 
+            -convert-scf-to-cf 
+            -convert-cf-to-llvm
+            -llvm-request-c-wrappers
+            -convert-arith-to-llvm 
+            -convert-math-to-llvm 
+            -convert-math-to-libm 
+            -convert-func-to-llvm 
+            -reconcile-unrealized-casts | 
+          ${LLVM_TOOLS_BINARY_DIR}/mlir-translate -mlir-to-llvmir |
+          ${LLVM_TOOLS_BINARY_DIR}/llvm-as |
+          ${LLVM_TOOLS_BINARY_DIR}/llc -filetype=obj  -relocation-model=pic -O0 -o ${CMAKE_CURRENT_BINARY_DIR}/subgraph0.o
+  DEPENDS ${CMAKE_CURRENT_BINARY_DIR}/subgraph0.mlir
+  COMMENT "Building subgraph0.o"
+  VERBATIM)
+
+set(MLP_EXAMPLE_PATH ${CMAKE_CURRENT_SOURCE_DIR})
+set(MLP_EXAMPLE_BUILD_PATH ${CMAKE_CURRENT_BINARY_DIR})
+
+add_library(MLP STATIC subgraph0.o forward.o)
+
+SET_TARGET_PROPERTIES(MLP PROPERTIES LINKER_LANGUAGE C)
+
+add_executable(buddy-mlp-run buddy-mlp-main.cpp)
+target_link_directories(buddy-mlp-run PRIVATE ${LLVM_LIBRARY_DIR})
+
+if(NOT DEFINED BUDDY_ENABLE_PNG)
+  message(FATAL_ERROR "To run MLP inference, the png library is required. Please define BUDDY_ENABLE_PNG for CMake.")
+endif()
+set(BUDDY_MLP_LIBS MLP mlir_c_runner_utils mlir_async_runtime mlir_runner_utils ${PNG_LIBRARIES})
+
+target_link_libraries(buddy-mlp-run ${BUDDY_MLP_LIBS})
+
+target_compile_definitions(buddy-mlp-run PRIVATE
+    MLP_EXAMPLE_PATH="${MLP_EXAMPLE_PATH}"
+    MLP_EXAMPLE_BUILD_PATH="${MLP_EXAMPLE_BUILD_PATH}"
+)

examples/BuddyMlp/README.md

@@ -0,0 +1,88 @@
+# Buddy Compiler MLP Example
+
+## Train the MLP Model
+
+Activate your python environment.
+
+```bash
+$ cd buddy-mlir
+$ cd examples/BuddyMlp
+$ python pytorch-mlp-train.py
+```
+
+## MLP Model Inference
+
+### Activate your python environment.
+
+```bash
+$ conda activate <your env>
+```
+
+### Build LLVM
+
+```bash
+$ cd buddy-mlir
+$ mkdir llvm/build
+$ cd llvm/build
+
+// CPU
+$ cmake -G Ninja ../llvm \
+    -DLLVM_ENABLE_PROJECTS="mlir;clang;openmp" \
+    -DLLVM_TARGETS_TO_BUILD="host;RISCV" \
+    -DLLVM_ENABLE_ASSERTIONS=ON \
+    -DOPENMP_ENABLE_LIBOMPTARGET=OFF \
+    -DCMAKE_BUILD_TYPE=RELEASE \
+    -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
+    -DPython3_EXECUTABLE=$(which python3)
+
+$ ninja check-clang check-mlir omp
+```
+
+### Build buddy-mlir
+
+```bash
+$ cd buddy-mlir
+$ mkdir build && cd build
+$ cmake -G Ninja .. \
+    -DMLIR_DIR=$PWD/../llvm/build/lib/cmake/mlir \
+    -DLLVM_DIR=$PWD/../llvm/build/lib/cmake/llvm \
+    -DLLVM_ENABLE_ASSERTIONS=ON \
+    -DCMAKE_BUILD_TYPE=RELEASE \
+    -DBUDDY_MLIR_ENABLE_PYTHON_PACKAGES=ON \
+    -DPython3_EXECUTABLE=$(which python3) \
+    -DBUDDY_MLIR_ENABLE_DIP_LIB=ON \
+    -DBUDDY_ENABLE_PNG=ON
+$ ninja
+$ ninja check-buddy
+```
+
+### Set the `PYTHONPATH` environment variable.
+
+Make sure you are in the build directory.
+
+```bash
+$ export BUDDY_MLIR_BUILD_DIR=$PWD
+$ export LLVM_MLIR_BUILD_DIR=$PWD/../llvm/build
+$ export PYTHONPATH=${LLVM_MLIR_BUILD_DIR}/tools/mlir/python_packages/mlir_core:${BUDDY_MLIR_BUILD_DIR}/python_packages:${PYTHONPATH}
+```
+
+### Build and run the MLP example
+
+```bash
+$ cmake -G Ninja .. -DBUDDY_MLP_EXAMPLES=ON
+
+$ ninja buddy-mlp-run
+$ cd bin
+$ ./buddy-mlp-run
+
+```
+
+## Debug the Lowering Pass Pipeline with Fake Parameters.
+
+```bash
+$ cd buddy-mlir
+$ cd examples/BuddyMlp
+$ make buddy-mlp-fusion
+$ make buddy-mlp-fused-run
+$ make buddy-mlp-run
+```

examples/BuddyMlp/buddy-mlp-import.py

@@ -0,0 +1,85 @@
+# ===- buddy-mlp-import.py ---------------------------------------------------
+#
+# 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.
+#
+# ===---------------------------------------------------------------------------
+#
+# This is the MLP model AOT importer.
+#
+# ===---------------------------------------------------------------------------
+
+import os
+from pathlib import Path
+import argparse
+
+import numpy as np
+import torch
+
+from buddy.compiler.frontend import DynamoCompiler
+from buddy.compiler.graph import GraphDriver
+from buddy.compiler.graph.transform import simply_fuse
+from buddy.compiler.ops import tosa
+from torch._inductor.decomposition import decompositions as inductor_decomp
+from model import MLP
+
+# Parse command-line arguments.
+parser = argparse.ArgumentParser(description="MLP model AOT importer")
+parser.add_argument(
+    "--output-dir",
+    type=str,
+    default="./",
+    help="Directory to save output files.",
+)
+args = parser.parse_args()
+
+# Ensure output directory exists.
+output_dir = Path(args.output_dir)
+output_dir.mkdir(parents=True, exist_ok=True)
+
+# Retrieve the MLP model path.
+model_path = os.path.dirname(os.path.abspath(__file__))
+
+model = MLP()
+model = torch.load(model_path + "/mlp-model.pth", weights_only=False)
+model = model.eval()
+print(model)
+# Initialize Dynamo Compiler with specific configurations as an importer.
+dynamo_compiler = DynamoCompiler(
+    primary_registry=tosa.ops_registry, verbose=True,
+    aot_autograd_decomposition=inductor_decomp,
+)
+
+data = torch.randn([1, 1, 784])
+# Import the model into MLIR module and parameters.
+with torch.no_grad():
+    graphs = dynamo_compiler.importer(model, data)
+
+assert len(graphs) == 1
+graph = graphs[0]
+params = dynamo_compiler.imported_params[graph]
+pattern_list = [simply_fuse]
+graphs[0].fuse_ops(pattern_list)
+driver = GraphDriver(graphs[0])
+driver.subgraphs[0].lower_to_top_level_ir()
+with open(output_dir / "subgraph0.mlir", "w") as module_file:
+    print(driver.subgraphs[0]._imported_module, file=module_file)
+with open(output_dir / "forward.mlir", "w") as module_file:
+    print(driver.construct_main_graph(True), file=module_file)
+
+params = dynamo_compiler.imported_params[graph]
+
+float32_param = np.concatenate(
+    [param.detach().numpy().reshape([-1]) for param in params]
+)
+
+float32_param.tofile(output_dir / "arg0.data")

examples/BuddyMlp/buddy-mlp-main.cpp

@@ -0,0 +1,137 @@
+//===- buddy-mlp-main.cpp -------------------------------------------------===//
+//
+// 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 <buddy/Core/Container.h>
+#include <buddy/DIP/ImgContainer.h>
+#include <chrono>
+#include <cmath>
+#include <cstdlib>
+#include <filesystem>
+#include <fstream>
+#include <limits>
+#include <string>
+#include <utility>
+#include <vector>
+
+constexpr size_t ParamsSize = 26506;
+const std::string ImgName = "3.png";
+
+/// Declare MLP forward function.
+extern "C" void _mlir_ciface_forward(MemRef<float, 3> *output,
+                                     MemRef<float, 1> *params,
+                                     MemRef<float, 3> *input);
+
+/// Print [Log] label in bold blue format.
+void printLogLabel() { std::cout << "\033[34;1m[Log] \033[0m"; }
+
+/// Load parameters into data container.
+void loadParameters(const std::string &paramFilePath,
+                    MemRef<float, 1> &params) {
+  const auto loadStart = std::chrono::high_resolution_clock::now();
+  // Open the parameter file in binary mode.
+  std::ifstream paramFile(paramFilePath, std::ios::in | std::ios::binary);
+  if (!paramFile.is_open()) {
+    throw std::runtime_error("[Error] Failed to open params file!");
+  }
+  printLogLabel();
+  std::cout << "Loading params..." << std::endl;
+  printLogLabel();
+  // Print the canonical path of the parameter file.
+  std::cout << "Params file: " << std::filesystem::canonical(paramFilePath)
+            << std::endl;
+  // Read the parameter data into the provided memory reference.
+  paramFile.read(reinterpret_cast<char *>(params.getData()),
+                 sizeof(float) * (params.getSize()));
+  if (paramFile.fail()) {
+    throw std::runtime_error("Error occurred while reading params file!");
+  }
+  paramFile.close();
+  const auto loadEnd = std::chrono::high_resolution_clock::now();
+  const std::chrono::duration<double, std::milli> loadTime =
+      loadEnd - loadStart;
+  printLogLabel();
+  std::cout << "Params load time: " << (double)(loadTime.count()) / 1000
+            << "s\n"
+            << std::endl;
+}
+
+/// Softmax function to convert logits to probabilities.
+void softmax(float *input, size_t size) {
+  size_t i;
+  float max_value = -INFINITY;
+  double sum = 0.0;
+  // Find the maximum value in the input array for numerical stability.
+  for (i = 0; i < size; ++i) {
+    if (max_value < input[i]) {
+      max_value = input[i];
+    }
+  }
+  // Calculate the sum of the exponentials of the input elements, normalized by
+  // the max value.
+  for (i = 0; i < size; ++i) {
+    sum += exp(input[i] - max_value);
+  }
+  // Normalize the input array with the softmax calculation.
+  for (i = 0; i < size; ++i) {
+    input[i] = exp(input[i] - max_value) / sum;
+  }
+}
+
+int main() {
+  // Print the title of this example.
+  const std::string title = "MLP Inference Powered by Buddy Compiler";
+  std::cout << "\033[33;1m" << title << "\033[0m" << std::endl;
+
+  // Define the sizes of the output tensors.
+  intptr_t sizesOutput[3] = {1, 1, 10};
+  intptr_t sizesInput[3] = {1, 1, 784};
+
+  // Create input and output containers for the image and model output.
+  std::string mlpDir = MLP_EXAMPLE_PATH;
+  std::string mlpBuildDir = MLP_EXAMPLE_BUILD_PATH;
+  std::string imgPath = mlpDir + "/images/" + ImgName;
+  dip::Image<float, 4> image(imgPath, dip::DIP_GRAYSCALE, true /* norm */);
+  MemRef<float, 3> input(sizesInput);
+  std::memcpy(input.getData(), image.getData(), 784 * sizeof(float));
+  MemRef<float, 3> output(sizesOutput);
+
+  // Load model parameters from the specified file.
+  std::string paramsDir = mlpBuildDir + "/arg0.data";
+  MemRef<float, 1> paramsContainer({ParamsSize});
+  loadParameters(paramsDir, paramsContainer);
+
+  // Call the forward function of the model.
+  _mlir_ciface_forward(&output, &paramsContainer, &input);
+
+  // Apply softmax to the output logits to get probabilities.
+  auto out = output.getData();
+  softmax(out, 10);
+
+  // Find the classification and print the result.
+  float maxVal = 0;
+  float maxIdx = 0;
+  for (int i = 0; i < 10; ++i) {
+    if (out[i] > maxVal) {
+      maxVal = out[i];
+      maxIdx = i;
+    }
+  }
+
+  std::cout << "Classification: " << maxIdx << std::endl;
+  std::cout << "Probability: " << maxVal << std::endl;
+
+  return 0;
+}