Merge pull request #56 from opencompl/sasha/alexnet

generate alexnet linalg
opencompl · Nov 3, 2023 · 1e309b7 · 1e309b7
2 parents 153d8bd + 3266da9
commit 1e309b7
Show file tree

Hide file tree

Showing 6 changed files with 1,132 additions and 0 deletions.
diff --git a/alexnet/flow.mlir b/alexnet/flow.mlir
diff --git a/alexnet/input.mlir b/alexnet/input.mlir
diff --git a/alexnet/iree.mlir b/alexnet/iree.mlir
diff --git a/alexnet/requirements.txt b/alexnet/requirements.txt
@@ -0,0 +1,2 @@
+shark-turbine
+torchvision
diff --git a/alexnet/run.sh b/alexnet/run.sh
@@ -0,0 +1,5 @@
+# Script that generated the mlir files in this folder
+
+python test.py > iree.mlir
+iree-compile --iree-input-type=torch --compile-to=input iree.mlir > input.mlir
+iree-compile --iree-input-type=torch --compile-to=flow --iree-util-zero-fill-elided-attrs iree.mlir > flow.mlir
diff --git a/alexnet/test.py b/alexnet/test.py
@@ -0,0 +1,12 @@
+import torch
+import torch.nn as nn
+
+import shark_turbine.aot as aot
+from iree.compiler import compile_str
+
+from torchvision.models import AlexNet
+
+model = AlexNet()
+example_x = torch.empty(64, 3, 224, 224)
+exported = aot.export(model, example_x)
+exported.print_readable()