[tests] add pytorch.lenet

Signed-off-by: Avimitin <[email protected]>

tests/pytorch/lenet/.gitignore

@@ -0,0 +1,3 @@
+arg0.data
+forward.mlir
+subgraph0.mlir

tests/pytorch/lenet/build.nix

@@ -0,0 +1,60 @@
+{ buildBuddyE2ETest, fetchurl }:
+let
+  lenetModel = fetchurl {
+    url = "https://github.com/buddy-compiler/buddy-benchmark/blob/1e166d53faae6d96a209645688cd9ab1d6eb604d/benchmarks/DeepLearning/Models/LeNet/lenet_model.pth";
+    hash = "sha256-imM6Hbl//AXQd1aCgJGB1S3eweavxOVc+bup9B/MFpA=";
+  };
+in
+buildBuddyE2ETest {
+  caseName = "lenet";
+
+  optPhase = ''
+    export LENET_MODEL_PATH=${lenetModel}
+    python ./lenet.py
+
+    echo "Lowering forward.mlir"
+    buddy-opt 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, 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), \
+                                  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)" \
+      > forward-lowered.mlir
+
+    echo "Lowering subgraphs[0]"
+    buddy-opt subgraphs0.mlir -pass-pipeline \
+        "builtin.module(func.func(tosa-to-linalg-named, tosa-to-arith, tosa-to-linalg, tosa-to-tensor))" \
+      | buddy-opt \
+          -convert-elementwise-to-linalg \
+          -func-bufferize-dynamic-offset \
+          -arith-bufferize \
+          -func-bufferize \
+          -tensor-bufferize \
+          -linalg-bufferize \
+          -finalizing-bufferize \
+          -batchmatmul-optimize \
+          -convert-linalg-to-affine-loops \
+          -lower-affine \
+          -convert-vector-to-scf \
+          -convert-scf-to-cf \
+          -llvm-request-c-wrappers \
+          -convert-vector-to-llvm \
+          -convert-math-to-llvm \
+          -convert-math-to-libm \
+          -convert-arith-to-llvm \
+          -convert-func-to-llvm \
+          -expand-strided-metadata \
+          -finalize-memref-to-llvm \
+          -reconcile-unrealized-casts \
+      > subgraphs0-lowered.mlir
+
+    optArtifacts+=(
+      "forward-lowered.mlir"
+      "subgraphs0-lowered.mlir"
+    )
+  '';
+}

tests/pytorch/lenet/lenet.cc

@@ -0,0 +1,34 @@
+#include "memref.hpp"
+
+#define INPUT_N 1
+#define INPUT_C 1
+#define INPUT_H 28
+#define INPUT_W 28
+#define INPUT_TOTAL (INPUT_N * INPUT_C * INPUT_H * INPUT_W)
+#define OUTPUT_N 10
+#define PARAM_N 44426
+
+__attribute((section(".vdata"))) float input_0[INPUT_TOTAL];
+__attribute((section(".vdata"))) float output_0[OUTPUT_N];
+__attribute((section(".vdata"))) float param_0[PARAM_N];
+
+// Define the sizes of the input and output tensors.
+static const int32_t sizesInput[4] = {INPUT_N, INPUT_C, INPUT_H, INPUT_W};
+static const int32_t sizesOutput[2] = {1, OUTPUT_N};
+static const int32_t sizesParams[1] = {PARAM_N};
+
+// Create input and output containers for the image and model output.
+MemRef<float, 4> input(input_0, sizesInput);
+MemRef<float, 2> output(output_0, sizesOutput);
+MemRef<float, 1> params(param_0, 2.0, sizesParams);
+
+// Declare the target model C interface.
+extern "C" {
+void _mlir_ciface_forward(MemRef<float, 2> *output, MemRef<float, 1> *arg0,
+                          MemRef<float, 4> *input);
+}
+
+extern "C" int test() {
+  _mlir_ciface_forward(&output, &params, &input);
+  return 0;
+}

tests/pytorch/lenet/lenet.py

@@ -0,0 +1,47 @@
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch._inductor.decomposition import decompositions as inductor_decomp
+
+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 model import LeNet
+
+def main():
+    model_path = os.environ.get("LENET_MODEL_PATH")
+    if model_path is None:
+        sys.exit("Error: No model path was provided. Please set $LENET_MODEL_PATH")
+    model = torch.load(model_path)
+    model = model.eval()
+
+    # Initialize Dynamo Compiler with specific configurations as an importer.
+    dynamo_compiler = DynamoCompiler(
+        primary_registry=tosa.ops_registry,
+        aot_autograd_decomposition=inductor_decomp,
+    )
+
+    data = torch.randn([1, 1, 28, 28])
+    # 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("subgraphs0.mlir", "w") as module_file:
+        print(driver.subgraphs[0]._imported_module, file=module_file)
+    with open("forward.mlir", "w") as module_file:
+        print(driver.construct_main_graph(True), file=module_file)
+
+if __name__ == "__main__":
+    main()

tests/pytorch/lenet/model.py

@@ -0,0 +1,42 @@
+# ===- model.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.
+#
+# ===---------------------------------------------------------------------------
+#
+# LeNet model definition.
+#
+# ===---------------------------------------------------------------------------
+
+import torch
+import torch.nn as nn
+
+
+class LeNet(nn.Module):
+    def __init__(self):
+        super(LeNet, self).__init__()
+        self.conv1 = nn.Conv2d(1, 6, 5)
+        self.pool = nn.MaxPool2d(2, 2)
+        self.conv2 = nn.Conv2d(6, 16, 5)
+        self.fc1 = nn.Linear(16 * 4 * 4, 120)
+        self.fc2 = nn.Linear(120, 84)
+        self.fc3 = nn.Linear(84, 10)
+
+    def forward(self, x):
+        x = self.pool(torch.relu(self.conv1(x)))
+        x = self.pool(torch.relu(self.conv2(x)))
+        x = x.view(-1, 16 * 4 * 4)
+        x = torch.relu(self.fc1(x))
+        x = torch.relu(self.fc2(x))
+        x = self.fc3(x)
+        return x