[doc] update document for the new framework

Signed-off-by: Avimitin <[email protected]>
chipsalliance · Aug 24, 2024 · bb1e81c · bb1e81c
1 parent 99db93f
commit bb1e81c
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 48 deletions.
diff --git a/tests/pytorch/README.md b/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:
 

diff --git a/tests/pytorch/default.nix b/tests/pytorch/default.nix
@@ -15,16 +15,17 @@ let
       buildBuddyE2ETest = { optPhase, ... }@overrides: builder
         ({
           inherit caseName;
-          configurePhase = ''
-            declare -A optArtifacts translateArtifacts llcArtifacts
-          '';
 
           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