[test] Add PyTorch Lenet model

nix/pkgs/buddy-mlir.nix

@@ -18,8 +18,8 @@ let
     src = fetchFromGitHub {
       owner = "buddy-compiler";
       repo = "buddy-mlir";
-      rev = "d7d90a488ac0d6fc1e700e932f842c7b2bcad816";
-      hash = "sha256-MhykCa6Z7Z8PpAlNh+vMuWYEOZZDyWhtMzMnFlNbGIk=";
+      rev = "802cefe91199c0935122546d463e400bee8635a6";
+      hash = "sha256-d8e/VM5LrsEwsC7NyNy/kdBp0fpY/CWeItrk4adOK0A=";
     };
 
     nativeBuildInputs = [ cmake ninja bintools ];

tests/pytorch/README.md

@@ -8,11 +8,10 @@ Assuming that the new PyTorch test have project name call `demo`, let's create t
 cd tests/pytorch
 mkdir -p demo
 cd demo
-touch demo.cc demo.py config.nix
+touch demo.cc demo.py build.nix
 ```
 
 Developers should put their PyTorch implementation into "<project-name>.py" file.
-For each PyTorch tests, developers must write the MLIR model to "forward.mlir" file.
 
 ```python
 # demo.py
@@ -49,57 +48,55 @@ extern "C" int test() {
 }
 ```
 
-After PyTorch model and the C entry is correctly created, developers should declare a "config.nix"
+After PyTorch model and the C entry is correctly created, developers should declare a "build.nix"
 file to indicate our build system to find and build the test case:
 
 ```nix
-{
-  # Tell our build system to include the memref.h header.
-  # Developer could add extra headers here.
-  includes = [
-    ../memref.hpp
-  ];
-
-  # Tell the build system to run buddy-opt with three phrase, with arguments to run in each phrase
-  buddyOptArgs = [
-    [
-      "--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, arith-bufferize, func.func(linalg-bufferize, tensor-bufferize), func-bufferize)"
-    ]
-    [
-      "--pass-pipeline"
-      "builtin.module(func.func(buffer-deallocation-simplification, convert-linalg-to-loops), eliminate-empty-tensors, func.func(llvm-request-c-wrappers))"
-    ]
-    [
-      "--lower-affine"
-      "--convert-math-to-llvm"
-      "--convert-math-to-libm"
-      "--convert-scf-to-cf"
-      "--convert-arith-to-llvm"
-      "--expand-strided-metadata"
-      "--finalize-memref-to-llvm"
-      "--lower-vector-exp"
-      "--lower-rvv=rv32"
-      "--convert-vector-to-llvm"
-      "--convert-func-to-llvm"
-      "--reconcile-unrealized-casts"
-    ]
-  ];
+{ buildBuddyE2ETest }:
+buildBuddyE2ETest {
+  caseName = "demo";
+
+  optPhase = ''
+    echo "Lowering MLIR"
+    python ./demo.py \
+    | buddy-opt --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, arith-bufferize, \
+        func.func(linalg-bufferize, tensor-bufferize), func-bufferize)" \
+    | buddy-opt --pass-pipeline "builtin.module(func.func(buffer-deallocation-simplification, convert-linalg-to-loops), \
+        eliminate-empty-tensors, func.func(llvm-request-c-wrappers))" \
+    | buddy-opt --lower-affine \
+        --convert-math-to-llvm \
+        --convert-math-to-libm \
+        --convert-scf-to-cf \
+        --convert-arith-to-llvm \
+        --expand-strided-metadata \
+        --finalize-memref-to-llvm \
+        --lower-vector-exp \
+        --lower-rvv=rv32 \
+        --convert-vector-to-llvm \
+        --convert-func-to-llvm \
+        --reconcile-unrealized-casts \
+        -o forward-lowered.mlir
+
+    optArtifacts+=(
+      "forward-lowered.mlir"
+    )
+  '';
 }
 ```
 
-Our build system accept the below data layout for the "config.nix" file:
+Here you can think `optPhase` as a bash function. Developers can write their own pass in this function.
+Each `optPhase` should modify the `optArtifacts` array, to indicate our build system about the final output.
 
-```text
-Set {
-    buddyOptArgs: Array<Array<String>>,
+The `caseName` and `optPhase` attribute is always required.
+We also offer the below attribute for you to override:
 
-    includes: Optional<Array<String>>,
-    pythonArgs: Optional<Array<String>>,
-    buddyTranslateArgs: Optional<Array<String>>,
-    buddyLLCArgs: Optional<Array<String>>,
-}
-```
+* `translatePhase`: By default, run `buddy-translate --buddy-to-llvmir` for each file in `optArtifacts` array,
+  add output into `translateArtifacts` array.
+* `llcPhase`: By default run `buddy-llc` for each file in `translateArtifacts` array, add output to `llcArtifacts` array.
+* `linkPhase`: By default link all the .o object file present in `llcArtifacts` array and the `caseName.cc` C++ file.
+
+> Developer can add other `stdenv.mkDerivation` attribute to override if they know the risks.
 
 After the project have been implemented, developers can run the below commands to build and test the ELF:
 

tests/pytorch/default.nix

@@ -5,96 +5,88 @@
 , findAndBuild
 , getTestRequiredFeatures
 , t1main
+, callPackage
 }:
 
 let
-
-  builder = makeBuilder { casePrefix = "mlir"; };
+  builder = makeBuilder { casePrefix = "pytorch"; };
   build = { caseName, sourcePath }:
-    let
-      buddyBuildConfig = import (sourcePath + "/config.nix");
-      defaultBuddyTranslateArgs = [ "--buddy-to-llvmir" ];
-      defaultBuddyLLCArgs = [
-        "-mtriple=riscv32"
-        "-target-abi=ilp32f"
-        "-mattr=+m,+f,+zve32f"
-        "-riscv-v-vector-bits-min=128"
-      ];
-    in
-    builder rec {
-      inherit caseName;
-
-      src = sourcePath;
-
-      featuresRequired = getTestRequiredFeatures sourcePath;
-
-      nativeBuildInputs = [ buddy-mlir.pyenv buddy-mlir ];
-
-      pythonArgs = buddyBuildConfig.pythonArgs or [ ];
-      buddyTranslateArgs = buddyBuildConfig.buddyTranslateArgs or defaultBuddyTranslateArgs;
-      buddyLLCArgs = buddyBuildConfig.buddyLLCArgs or defaultBuddyLLCArgs;
-      buddyIncludes = buddyBuildConfig.includes or [ ];
-
-      postUnpack = ''
-        buddyIncludeDir="."
-        if [ "x$buddyIncludes" != "x" ]; then
-          mkdir -p buddyInclude
-          _buddyHeaderArray=( $buddyIncludes )
-          for h in "''${_buddyHeaderArray}"; do
-            cp -v "$h" buddyInclude/"$(stripHash $h)"
-          done
-
-          buddyIncludeDir=$PWD/buddyInclude
-        fi
-      '';
-
-      buildPhase = ''
-        runHook preBuild
-
-        echo "Running python with args $pythonArgs"
-        python $pythonArgs ${caseName}.py
-
-        # Generate multiple buddy-opt call, each will read input from former pipeline
-        # For example, for buddyOptArgs = [ [ "--arg-a" ], [ "--arg-b" ], [ "--arg-c" ] ]
-        # This will generate
-        #
-        #   echo "..."
-        #   buddy-opt forward.mlir --arg-a -o forward-1.mlir
-        #   echo "..."
-        #   buddy-opt forward-1.mlir --arg-b -o forward-2.mlir
-        #   echo "..."
-        #   buddy-opt forward-2.mlir --arg-c -o forward-3.mlir
-        #
-        ${lib.concatStringsSep "\n" (
-          lib.imap0
-          (idx: args: ''
-            echo "Running buddy-opt with args ${lib.escapeShellArgs args}"
-            buddy-opt \
-              forward${if idx == 0 then "" else "-${toString idx}"}.mlir \
-              ${lib.escapeShellArgs args} \
-              -o forward-${toString (idx+1)}.mlir
-          '')
-          buddyBuildConfig.buddyOptArgs
-        )}
-
-        # Pick up the last optimized MLIR file
-        echo "Running buddy-translate with args $buddyTranslateArgs"
-        buddy-translate forward-${with builtins; toString (length buddyBuildConfig.buddyOptArgs)}.mlir \
-          $buddyTranslateArgs -o forward.ll
-
-        echo "Running buddy-llc with args $buddyLLCArgs"
-        buddy-llc forward.ll $buddyLLCArgs --filetype=obj -o forward.o
-
-        echo "Using include dir $buddyIncludeDir"
-        $CXX -nostdlib -I$buddyIncludeDir -c ${caseName}.cc -o host.o
-        $CC -T${linkerScript} \
-          host.o forward.o ${t1main} \
-          -o $pname.elf
-
-        runHook postBuild
-      '';
-
-      meta.description = "testcase '${caseName}', written in MLIR";
+    callPackage (sourcePath + "/build.nix") {
+      buildBuddyE2ETest = { optPhase, ... }@overrides: builder
+        ({
+          inherit caseName;
+
+          featuresRequired = getTestRequiredFeatures sourcePath;
+
+          nativeBuildInputs = [ buddy-mlir.pyenv buddy-mlir ];
+
+          src = sourcePath;
+
+          configurePhase = ''
+            declare -A optArtifacts translateArtifacts llcArtifacts
+          '';
+
+          translatePhase = ''
+            if [[ -z "$optArtifacts" ]]; then
+              echo "optPhase doesn't produce optArtifacts, abort" >&2
+              exit 1
+            fi
+
+            for mlir in ''${optArtifacts[@]}; do
+              echo "Translating $mlir"
+              buddy-translate --buddy-to-llvmir "$mlir" -o "$mlir.ll"
+
+              translateArtifacts+=("$mlir.ll")
+            done
+          '';
+
+          llcPhase = ''
+            if [[ -z "$translateArtifacts" ]]; then
+              echo "translatePhase doesn't produce translateArtifacts, abort" >&2
+              exit 1
+            fi
+
+            for llvmir in ''${translateArtifacts[@]}; do
+              echo "Compiling $llvmir"
+              buddy-llc "$llvmir" \
+                -mtriple=riscv32 \
+                -target-abi=ilp32f \
+                -mattr=+m,+f,+zve32f \
+                -riscv-v-vector-bits-min=128 \
+                --filetype=obj \
+                -o "$llvmir.o"
+
+              llcArtifacts+=("$llvmir.o")
+            done
+          '';
+
+          linkPhase = ''
+            if [[ -z "$llcArtifacts" ]]; then
+              echo "llcPhase doesn't produce any llcArtifacts" >&2
+              exit 1
+            fi
+
+            echo "Building final binary"
+            mkdir -p _include
+            cp ${./memref.hpp} _include/memref.hpp
+
+            $CXX -nostdlib -I _include -c ${caseName}.cc -o host.o
+            $CC -T${linkerScript} \
+              host.o ''${llcArtifacts[@]} ${t1main} \
+              -o $pname.elf
+          '';
+
+          buildPhase = ''
+            runHook preBuild
+
+            runPhase optPhase
+            runPhase translatePhase
+            runPhase llcPhase
+            runPhase linkPhase
+
+            runHook postBuild
+          '';
+        } // overrides);
     };
 in
 findAndBuild ./. build

tests/pytorch/demo/build.nix

@@ -0,0 +1,31 @@
+{ buildBuddyE2ETest }:
+buildBuddyE2ETest {
+  caseName = "demo";
+
+  optPhase = ''
+    echo "Lowering MLIR"
+    python ./demo.py \
+    | buddy-opt --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, arith-bufferize, \
+        func.func(linalg-bufferize, tensor-bufferize), func-bufferize)" \
+    | buddy-opt --pass-pipeline "builtin.module(func.func(buffer-deallocation-simplification, convert-linalg-to-loops), \
+        eliminate-empty-tensors, func.func(llvm-request-c-wrappers))" \
+    | buddy-opt --lower-affine \
+        --convert-math-to-llvm \
+        --convert-math-to-libm \
+        --convert-scf-to-cf \
+        --convert-arith-to-llvm \
+        --expand-strided-metadata \
+        --finalize-memref-to-llvm \
+        --lower-vector-exp \
+        --lower-rvv=rv32 \
+        --convert-vector-to-llvm \
+        --convert-func-to-llvm \
+        --reconcile-unrealized-casts \
+        -o forward-lowered.mlir
+
+    optArtifacts+=(
+      "forward-lowered.mlir"
+    )
+  '';
+}