iree

Dockerfile

@@ -29,14 +29,20 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     --mount=type=cache,target=/var/lib/apt,sharing=locked \
     apt-get -y update && apt-get install -y mlir-18-tools llvm-18
 
+# IREE
+RUN pip install \
+  --index-url https://download.pytorch.org/whl/test/cpu torch==2.4.1
+RUN pip install iree-turbine
+
+
 # Install pytorch-mlir
 
-RUN pip install --pre torch-mlir torch \
-  --extra-index-url https://download.pytorch.org/whl/nightly/cpu \
-  -f https://github.com/llvm/torch-mlir-release/releases/expanded_assets/dev-wheels
+#RUN pip install --pre torch-mlir torch \
+#  --extra-index-url https://download.pytorch.org/whl/nightly/cpu \
+#  -f https://github.com/llvm/torch-mlir-release/releases/expanded_assets/dev-wheels
 
 # Add PyTorch-MLIR tools to PATH
-ENV PATH="/root/.local/bin:${PATH}"
+#ENV PATH="/root/.local/bin:${PATH}"
 
 # Verify installations
 RUN python -c "import torch; print(torch.__version__)" && \

main.py

@@ -13,13 +13,15 @@ def process_code(pytorch_code):
             code_file.flush()
 
             # Step 2: Run the PyTorch code to generate TorchScript
-            script_output_path = tempfile.NamedTemporaryFile(delete=False, suffix=".pt").name
+            #script_output_path = tempfile.NamedTemporaryFile(delete=False, suffix=".pt").name
             subprocess.run(
-                ["python3", code_file.name, script_output_path],
+                ["python3", code_file.name],
                 check=True,
                 capture_output=True,
             )
 
+            print("Converted to onnx.")
+
             # Step 3: Convert TorchScript to MLIR (torch dialect)
             torch_mlir_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mlir").name
             subprocess.run(
@@ -70,7 +72,7 @@ iface = gr.Interface(
     fn=process_code,
     inputs="text",
     outputs=[
-        gr.Textbox(label="Torch Dialect"),
+        gr.Textbox(label="Torch Dialect", lines=10),
         gr.Textbox(label="Linalg Dialect"),
         gr.Textbox(label="GPU Dialect"),
         gr.Textbox(label="LLVM Dialect"),