diff --git a/tests/pytorch/lenet/.gitignore b/tests/pytorch/lenet/.gitignore new file mode 100644 index 000000000..4a44e5828 --- /dev/null +++ b/tests/pytorch/lenet/.gitignore @@ -0,0 +1,3 @@ +arg0.data +forward.mlir +subgraph0.mlir diff --git a/tests/pytorch/lenet/build.nix b/tests/pytorch/lenet/build.nix new file mode 100644 index 000000000..e4d110341 --- /dev/null +++ b/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" + ) + ''; +} diff --git a/tests/pytorch/lenet/lenet.cc b/tests/pytorch/lenet/lenet.cc new file mode 100644 index 000000000..3cb5d1a88 --- /dev/null +++ b/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 input(input_0, sizesInput); +MemRef output(output_0, sizesOutput); +MemRef params(param_0, 2.0, sizesParams); + +// Declare the target model C interface. +extern "C" { +void _mlir_ciface_forward(MemRef *output, MemRef *arg0, + MemRef *input); +} + +extern "C" int test() { + _mlir_ciface_forward(&output, ¶ms, &input); + return 0; +} diff --git a/tests/pytorch/lenet/lenet.py b/tests/pytorch/lenet/lenet.py new file mode 100644 index 000000000..9fd364f68 --- /dev/null +++ b/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() diff --git a/tests/pytorch/lenet/lenet_model.pth b/tests/pytorch/lenet/lenet_model.pth new file mode 100644 index 000000000..1b86e3bd5 Binary files /dev/null and b/tests/pytorch/lenet/lenet_model.pth differ diff --git a/tests/pytorch/lenet/model.py b/tests/pytorch/lenet/model.py new file mode 100644 index 000000000..2f32951a6 --- /dev/null +++ b/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