diff --git a/.github/cases/blastoise/default.json b/.github/cases/blastoise/default.json index 61d8f3bc1..2b03351dd 100644 --- a/.github/cases/blastoise/default.json +++ b/.github/cases/blastoise/default.json @@ -2,6 +2,7 @@ "pytorch.demo": -1, "pytorch.lenet": -1, "pytorch.matmul": -1, + "pytorch.mobilenet": -1, "mlir.rvv_vp_intrinsic_add": 436, "mlir.rvv_vp_intrinsic_add_scalable": 584, "mlir.hello": 146, diff --git a/nix/pkgs/buddy-mlir.nix b/nix/pkgs/buddy-mlir.nix index 2f520e491..b15cb9397 100644 --- a/nix/pkgs/buddy-mlir.nix +++ b/nix/pkgs/buddy-mlir.nix @@ -57,6 +57,7 @@ let pyenv = python3.withPackages (ps: [ self ps.torch + ps.torchvision ]); }; }; diff --git a/tests/default.nix b/tests/default.nix index 1f0428112..ab934a2bc 100644 --- a/tests/default.nix +++ b/tests/default.nix @@ -80,7 +80,7 @@ let findAndBuild = dir: build: lib.recurseIntoAttrs (lib.pipe (builtins.readDir dir) [ # filter out all non-directory entrires and underscore-prefixed directories - (lib.filterAttrs (name: type: type == "directory" && ! lib.hasPrefix "_" name)) + (lib.filterAttrs (name: type: type == "directory" && ! (lib.hasPrefix "_" name) && ! (lib.elem name [ "lib" "include" ]))) # prepend path with base directory (lib.mapAttrs (subDirName: _: (lib.path.append dir subDirName))) # build. If {sourcePath}/default.nix exists, call it. Otherwise call the generic builder diff --git a/tests/pytorch/default.nix b/tests/pytorch/default.nix index 5e8f57715..9b50d9cef 100644 --- a/tests/pytorch/default.nix +++ b/tests/pytorch/default.nix @@ -67,10 +67,7 @@ let fi echo "Building final binary" - mkdir -p _include - cp ${./memref.hpp} _include/memref.hpp - - $CXX -nostdlib -I _include -c ${caseName}.cc -o host.o + $CXX -nostdlib -I${./include} -c ${caseName}.cc -o host.o $CC -T${linkerScript} \ host.o ''${llcArtifacts[@]} ${t1main} \ -o $pname.elf diff --git a/tests/pytorch/include/img.hpp b/tests/pytorch/include/img.hpp new file mode 100644 index 000000000..f839df342 --- /dev/null +++ b/tests/pytorch/include/img.hpp @@ -0,0 +1,234 @@ +//===- ImgContainer.h -----------------------------------------------------===// +// +// 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. +// +//===----------------------------------------------------------------------===// +// +// original source from https://github.com/buddy-compiler/buddy-mlir +// Modified for embedded environment. +// +//===----------------------------------------------------------------------===// + +#ifndef IMG_H +#define IMG_H + +#include "memref.hpp" + +enum ImageModes { + DIP_GRAYSCALE = 0, + DIP_RGB = 1, +}; + +template class Image : public MemRef { +public: + // Constructor initializes the image by loading from a file. + // Params: + // file: Raw data to image file memory + // filesize: size of the file memory + // buffer: pre-allocated space for image modification + // mode: Specifies the image mode (e.g., DIP_GRAYSCALE, DIP_RGB). + // norm: Indicates whether to normalize pixel values (default is false). + constexpr Image(const char *file, int32_t filesize, T *buffer, + ImageModes mode, bool norm = false); + + // Overload + constexpr Image(T *data, const int32_t sizes[N]); + constexpr Image(T *data, T init, const int32_t sizes[N]); + + // Retrieves the name of the current image format as a string. + inline const char *getFormatName() const { + switch (this->imageFormat) { + case ImageFormat::BMP: + return "BMP"; + default: + return "Unsupported format"; + } + } + // Returns the width of the image in pixels. + inline size_t getWidth() const { return this->width; } + // Returns the height of the image in pixels. + inline size_t getHeight() const { return this->height; } + // Returns the bit depth of the image. + inline int getBitDepth() const { return this->bitDepth; } + +private: + // Enum to represent supported image formats. + enum class ImageFormat { + BMPDecodeError, // Represents an error or unsupported format. + BMP, // BMP file format. + Unsupported, + } imageFormat; + // Mode of the image (e.g., DIP_GRAYSCALE, DIP_RGB). + ImageModes imageMode; + // Width of the image in pixels. + int32_t width; + // Height of the image in pixels. + int32_t height; + // Bit depth of the image. + int32_t bitDepth; + // Normalization flag. + bool isNorm; + // Determines the image format from raw file data. + inline void determineFormat(const uint8_t *fileData, uint32_t filesize); + // Decodes a BMP image from raw file data. + inline bool decodeBMP(const uint8_t *fileData, uint32_t filesize); +}; + +template +constexpr Image::Image(T *data, const int32_t sizes[N]) + : MemRef(data, sizes){}; + +template +constexpr Image::Image(T *data, T init, const int32_t sizes[N]) + : MemRef(data, init, sizes){}; + +// Image Container Constructor +// Constructs an image container object from the image file path. +template +constexpr Image::Image(const char *file, int32_t filesize, T *buffer, + ImageModes mode, bool norm) + : imageMode(mode), isNorm(norm) { + + this->allocated = buffer; + this->aligned = this->allocated; + + determineFormat(file, filesize); + if (this->imageFormat == ImageFormat::BMP) { + bool success = decodeBMP(file, filesize); + if (!success) { + this->imageFormat = ImageFormat::BMPDecodeError; + }; + } else { + this->imageFormat = ImageFormat::Unsupported; + } +} + +// Determines the image format by inspecting the header of the file data. +template +inline void Image::determineFormat(const uint8_t *fileData, + uint32_t filesize) { + if (filesize > 2 && fileData[0] == 'B' && fileData[1] == 'M') { + this->imageFormat = ImageFormat::BMP; + } else { + this->imageFormat = ImageFormat::BMPDecodeError; + } +} + +// BMP Image File Decoder +template +inline bool Image::decodeBMP(const uint8_t *fileData, uint32_t filesize) { + // Check if the provided data is large enough to contain a minimal BMP header + // (54 bytes). + if (filesize < 54) { + return false; + } + + // Extract image information from BMP header + this->width = *reinterpret_cast(&fileData[18]); + this->height = *reinterpret_cast(&fileData[22]); + this->bitDepth = *reinterpret_cast(&fileData[28]); + uint32_t compression = *reinterpret_cast(&fileData[30]); + size_t pixelDataOffset = *reinterpret_cast(&fileData[10]); + + // Currently, only the BI_RGB (value 0) compression method is supported. + if (compression != 0) { + return false; + } + + // Currently, only the NCHW format with 4 dimensions is supported. + if (N == 4) { + if (this->imageMode == ImageModes::DIP_GRAYSCALE) { + // TODO: Add batch setting. + this->sizes[0] = 1; + this->sizes[1] = 1; + this->sizes[2] = this->height; + this->sizes[3] = this->width; + this->setStrides(); + size_t size = this->product(this->sizes); + // Fullfill data to memref container. + size_t memrefIndex = 0; + if (this->bitDepth == 32) { + // BMP file is upside-down storage. + for (size_t i = this->height; i > 0; i--) { + for (size_t j = 0; j < this->width; j++) { + // Locate the current pixel. + size_t pixelIndex = + pixelDataOffset + (((i - 1) * this->width) + j) * 4; + // Extract the blue, green, and red value from the current pixel. + int bluePixel = + *reinterpret_cast(&fileData[pixelIndex]); + int greenPixel = + *reinterpret_cast(&fileData[pixelIndex + 1]); + int redPixel = + *reinterpret_cast(&fileData[pixelIndex + 2]); + // Calculate the gray scale value. + int grayScaleValue = static_cast( + 0.299 * redPixel + 0.587 * greenPixel + 0.114 * bluePixel); + // Store the gray scale value into memref container. + this->aligned[memrefIndex] = + this->isNorm ? static_cast(grayScaleValue) / 255 + : static_cast(grayScaleValue); + memrefIndex++; + } + } + } else { + return false; + } + } else if (this->imageMode == ImageModes::DIP_RGB) { + // TODO: Add batch setting. + this->sizes[0] = 1; + this->sizes[1] = 3; + this->sizes[2] = this->height; + this->sizes[3] = this->width; + this->setStrides(); + size_t size = this->product(this->sizes); + // Fullfill data to memref container. + size_t memrefIndex = 0; + size_t colorStride = this->height * this->width; + if (this->bitDepth == 32) { + // BMP file is upside-down storage. + for (size_t i = height; i > 0; i--) { + for (size_t j = 0; j < width; j++) { + // Locate the current pixel. + size_t pixelIndex = pixelDataOffset + (((i - 1) * width) + j) * 4; + // Extract the blue, green, and red value from the current pixel. + int bluePixel = + *reinterpret_cast(&fileData[pixelIndex]); + int greenPixel = + *reinterpret_cast(&fileData[pixelIndex + 1]); + int redPixel = + *reinterpret_cast(&fileData[pixelIndex + 2]); + // Store the values into memref container as RGB order. (BGR -> RGB) + this->aligned[memrefIndex] = this->isNorm + ? static_cast(redPixel) / 255 + : static_cast(redPixel); + this->aligned[memrefIndex + colorStride] = + this->isNorm ? static_cast(greenPixel) / 255 + : static_cast(greenPixel); + this->aligned[memrefIndex + 2 * colorStride] = + this->isNorm ? static_cast(bluePixel) / 255 + : static_cast(bluePixel); + memrefIndex++; + } + } + } else { + return false; + } + } + } else { + return false; + } + return true; +} + +#endif // IMG_H diff --git a/tests/pytorch/memref.hpp b/tests/pytorch/include/memref.hpp similarity index 58% rename from tests/pytorch/memref.hpp rename to tests/pytorch/include/memref.hpp index b37b96181..95e6725f2 100644 --- a/tests/pytorch/memref.hpp +++ b/tests/pytorch/include/memref.hpp @@ -1,3 +1,24 @@ +//===- ImgContainer.h -----------------------------------------------------===// +// +// 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. +// +//===----------------------------------------------------------------------===// +// +// Original source from https://github.com/buddy-compiler/buddy-mlir. +// Modified for embedded environment. +// +//===----------------------------------------------------------------------===// + #ifndef MEMREF_H #define MEMREF_H @@ -56,4 +77,4 @@ template inline void MemRef::setStrides() { } } -#endif +#endif // MEMREF_H diff --git a/tests/pytorch/lenet/build.nix b/tests/pytorch/lenet/build.nix index 1bab5a80b..baabc0454 100644 --- a/tests/pytorch/lenet/build.nix +++ b/tests/pytorch/lenet/build.nix @@ -29,27 +29,29 @@ buildBuddyE2ETest { 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 \ + --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 \ + --lower-vector-exp \ + --lower-rvv=rv32 \ + --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+=( diff --git a/tests/pytorch/lib/MemrefCopy.cc b/tests/pytorch/lib/MemrefCopy.cc new file mode 100644 index 000000000..c7f15ff79 --- /dev/null +++ b/tests/pytorch/lib/MemrefCopy.cc @@ -0,0 +1,132 @@ +//===- CRunnerUtils.cpp - Utils for MLIR execution ------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// Modified to basic data type for linkage usage only. +// +//===----------------------------------------------------------------------===// + +#include +#include +#include + +// Unranked MemRef +template struct UnrankedMemRefType { + int32_t rank; + void *descriptor; +}; + +/// StridedMemRef descriptor type with static rank. +template struct StridedMemRefType { + T *basePtr; + T *data; + int64_t offset; + int64_t sizes[N]; + int64_t strides[N]; +}; + +/// StridedMemRef descriptor type specialized for rank 1. +template struct StridedMemRefType { + T *basePtr; + T *data; + int64_t offset; + int64_t sizes[1]; + int64_t strides[1]; + + T &operator[](int64_t idx) { return *(data + offset + idx * strides[0]); } +}; + +/// StridedMemRef descriptor type specialized for rank 0. +template struct StridedMemRefType { + T *basePtr; + T *data; + int64_t offset; +}; + +// A reference to one of the StridedMemRef types. +template class DynamicMemRefType { +public: + int64_t rank; + T *basePtr; + T *data; + int64_t offset; + const int64_t *sizes; + const int64_t *strides; + + explicit DynamicMemRefType(const StridedMemRefType &memRef) + : rank(0), basePtr(memRef.basePtr), data(memRef.data), + offset(memRef.offset), sizes(nullptr), strides(nullptr) {} + template + explicit DynamicMemRefType(const StridedMemRefType &memRef) + : rank(N), basePtr(memRef.basePtr), data(memRef.data), + offset(memRef.offset), sizes(memRef.sizes), strides(memRef.strides) {} + explicit DynamicMemRefType(const ::UnrankedMemRefType &memRef) + : rank(memRef.rank) { + auto *desc = static_cast *>(memRef.descriptor); + basePtr = desc->basePtr; + data = desc->data; + offset = desc->offset; + sizes = rank == 0 ? nullptr : desc->sizes; + strides = sizes + rank; + } +}; + +extern "C" void memrefCopy(int32_t elemSize, UnrankedMemRefType *srcArg, + UnrankedMemRefType *dstArg) { + DynamicMemRefType src(*srcArg); + DynamicMemRefType dst(*dstArg); + + int32_t rank = src.rank; + + // Handle empty shapes -> nothing to copy. + for (int rankp = 0; rankp < rank; ++rankp) + if (src.sizes[rankp] == 0) + return; + + char *srcPtr = src.data + src.offset * elemSize; + char *dstPtr = dst.data + dst.offset * elemSize; + + if (rank == 0) { + memcpy(dstPtr, srcPtr, elemSize); + return; + } + + int32_t *indices = static_cast(alloca(sizeof(int32_t) * rank)); + int32_t *srcStrides = static_cast(alloca(sizeof(int32_t) * rank)); + int32_t *dstStrides = static_cast(alloca(sizeof(int32_t) * rank)); + + // Initialize index and scale strides. + for (int rankp = 0; rankp < rank; ++rankp) { + indices[rankp] = 0; + srcStrides[rankp] = src.strides[rankp] * elemSize; + dstStrides[rankp] = dst.strides[rankp] * elemSize; + } + + int32_t readIndex = 0, writeIndex = 0; + for (;;) { + // Copy over the element, byte by byte. + memcpy(dstPtr + writeIndex, srcPtr + readIndex, elemSize); + // Advance index and read position. + for (int32_t axis = rank - 1; axis >= 0; --axis) { + // Advance at current axis. + auto newIndex = ++indices[axis]; + readIndex += srcStrides[axis]; + writeIndex += dstStrides[axis]; + // If this is a valid index, we have our next index, so continue copying. + if (src.sizes[axis] != newIndex) + break; + // We reached the end of this axis. If this is axis 0, we are done. + if (axis == 0) + return; + // Else, reset to 0 and undo the advancement of the linear index that + // this axis had. Then continue with the axis one outer. + indices[axis] = 0; + readIndex -= src.sizes[axis] * srcStrides[axis]; + writeIndex -= dst.sizes[axis] * dstStrides[axis]; + } + } +} diff --git a/tests/pytorch/mobilenet/build.nix b/tests/pytorch/mobilenet/build.nix new file mode 100644 index 000000000..7aaa44d03 --- /dev/null +++ b/tests/pytorch/mobilenet/build.nix @@ -0,0 +1,73 @@ +{ fetchurl +, buildBuddyE2ETest +}: +let + checkpointFile = "mobilenet_v3_small-047dcff4.pth"; + modelCache = fetchurl { + url = "https://download.pytorch.org/models/${checkpointFile}"; + hash = "sha256-BH3P9K3e+G6lvC7/E8lhTcEfR6sRYNCnGiXn25lPTh8="; + }; +in +buildBuddyE2ETest { + caseName = "mobilenet"; + + optPhase = '' + mkdir -p pytorchCache/hub/checkpoints/ + cp -v ${modelCache} pytorchCache/hub/checkpoints/${checkpointFile} + export TORCH_HOME=pytorchCache + python ./mobilenet.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 \ + --lower-vector-exp \ + --lower-rvv=rv32 \ + --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 + + echo "Compiling memrefCopy library" + $CXX -nostdlib -c ${../lib/MemrefCopy.cc} -o memrefCopy.o + llcArtifacts+=( + memrefCopy.o + ) + + optArtifacts+=( + "forward-lowered.mlir" + "subgraphs0-lowered.mlir" + ) + ''; +} diff --git a/tests/pytorch/mobilenet/mobilenet.cc b/tests/pytorch/mobilenet/mobilenet.cc new file mode 100644 index 000000000..d64064fbb --- /dev/null +++ b/tests/pytorch/mobilenet/mobilenet.cc @@ -0,0 +1,45 @@ +#include "img.hpp" + +#define INPUT_N 1 +#define INPUT_C 3 +#define INPUT_H 224 +#define INPUT_W 224 +#define INPUT_TOTAL (INPUT_N * INPUT_C * INPUT_H * INPUT_W) +#define OUTPUT_N 1000 +// #define PARAM_N0 2554968 +#define PARAM_N0 25549 +#define PARAM_N1 34 + +__attribute((section(".vdata"))) float input_0[INPUT_TOTAL]; +__attribute((section(".vdata"))) float output_0[OUTPUT_N]; +__attribute((section(".vdata"))) float param_0[PARAM_N0]; +__attribute((section(".vdata"))) int64_t param_1[PARAM_N1]; + +// 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 sizesParam0[1] = {PARAM_N0}; +static const int32_t sizesParam1[1] = {PARAM_N1}; + +extern "C" { +void _mlir_ciface_forward(MemRef *output, MemRef *arg0, + MemRef *arg1, Image *input); +} + +extern "C" int test() { + + // Generate input memref container with random numbers. + const int inputSize = INPUT_N * INPUT_C * INPUT_H * INPUT_W; + + // Create input and output containers for the image and model output. + Image input(input_0, sizesInput); + MemRef output(output_0, sizesOutput); + + // Set random model parameters. + MemRef paramsF32(param_0, 2.0, sizesParam0); + MemRef paramsI64(param_1, 1, sizesParam1); + + _mlir_ciface_forward(&output, ¶msF32, ¶msI64, &input); + + return 0; +} diff --git a/tests/pytorch/mobilenet/mobilenet.py b/tests/pytorch/mobilenet/mobilenet.py new file mode 100644 index 000000000..57299ef69 --- /dev/null +++ b/tests/pytorch/mobilenet/mobilenet.py @@ -0,0 +1,60 @@ +# ===- buddy_mobilenetv3_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 MobileNet V3 model AOT importer. +# +# Source: https://github.com/buddy-compiler/buddy-benchmark +# +# ===--------------------------------------------------------------------------- + +import os + +from pathlib import Path +import numpy as np +import torch +import torchvision.models as models +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 + +model = models.mobilenet_v3_small( + weights=models.MobileNet_V3_Small_Weights.IMAGENET1K_V1, pretrained=True +) +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, 3, 224, 224]) +# 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)