Add raw_ir output for pytorch (#4)

With the output type a user can specify their own way to print IR in the
code and this output will be respected by the tool. Example use case:
import torch
import torch.nn as nn
from torch_mlir import fx
from torch_mlir.compiler_utils import run_pipeline_with_repro_report
from torch_mlir.fx import OutputType

```
class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(4, 4)

    def forward(self, x):
        return torch.relu(self.linear(x))

model = MyModel()
example_input = torch.randn(4, 4)
module = fx.export_and_import(model, example_input, output_type=OutputType.LINALG_ON_TENSORS) print(module)
```

This patch also brings few modifications to README.

Author: MrSidims

README.md

@@ -26,8 +26,6 @@ tracing models through various IR stages and transformations.
 
 - (PyTorch) The model and input tensor must be initialized in the provided code. If multiple models are defined, it is recommended to explicitly pair each model and its input tensor using the internal `__explore__(model, input)` function.
 
-- (PyTorch) The current version does not recognize or capture user attempts to dump IR inside the input PyTorch module. It is planned that, in the future, if the user manually calls `fx.export_and_import()` (or similar IR-producing APIs), the app will use that IR as the base and apply the user-defined custom toolchain.
-
 - (Triton) The current implementation runs Triton kernels and retrieves IR dumps from the Triton cache directory. Timeout is set to 20s.
 
 ## Getting Started
@@ -41,11 +39,13 @@ tracing models through various IR stages and transformations.
 - Triton
 - LLVM with mlir-opt
 
-Current version is tested on Ubuntu 22.04 windows subsystem using LLVM 21 dev.
+To setup PyTorch and Torch-MLIR it's a good idea to visit https://github.com/llvm/torch-mlir repository and follow instructions from there.
+
+Current version of the application is tested on Ubuntu 22.04 windows subsystem using LLVM 21 dev.
 
 ### Install dependencies
 
-In case of missing prerequisites here are some scripts to help set them up.
+In case of missing prerequisites here are some scripts to help set them up (runs on Debian and its derivatives).
 
 ```bash
 git clone https://github.com/MrSidims/PytorchExplorer.git
@@ -74,6 +74,8 @@ If you want to use your builds of the tools like `torch-mlir-opt`, `mlir-opt` et
 npm run start:all
 ```
 
+Then open http://localhost:3000/ in your browser and enjoy!
+
 ### Run the tests
 
 With the application (or just backend) started, run:
@@ -82,9 +84,14 @@ With the application (or just backend) started, run:
 pytest tests -v
 ```
 
+## User manual
+
+TBD
+
 ## Implementation details
 
-The app uses fx.export_and_import to inpect IR output for PyTorch.
+The app uses `fx.export_and_import` under the hood to inpect IR output for PyTorch, therefore for pre-defined lowering paths it's required for a module to have `forward` method.
+
 Lowering to LLVM IR goes through:
 
 ```bash

backend/server.py

@@ -25,6 +25,10 @@
 import atexit
 import shutil
 
+import io
+import sys
+from contextlib import redirect_stdout, redirect_stderr
+
 app = FastAPI()
 
 app.add_middleware(
@@ -44,6 +48,7 @@
 class CodeRequest(BaseModel):
     code: str
     ir_type: str
+    selected_language: Optional[str] = "pytorch"
     custom_pipeline: Optional[List[str]] = []
     torch_mlir_opt: Optional[str] = ""
     mlir_opt: Optional[str] = ""
@@ -60,15 +65,20 @@ class FreeIRCacheRequest(BaseModel):
 
 # Get model-tensor pairs to process.
 def extract_model_input_pairs(code: str):
-    import traceback
-
     explore_pairs = []
 
     def __explore__(model, input_tensor):
         explore_pairs.append((model, input_tensor))
 
     exec_globals = {"__explore__": __explore__}
-    exec(code, exec_globals)
+
+    # Suppress stdout and stderr.
+    with open(os.devnull, "w") as devnull:
+        try:
+            with redirect_stdout(devnull), redirect_stderr(devnull):
+                exec(code, exec_globals)
+        except Exception:
+            pass
 
     # If no __explore__ calls found, try matching models and tensors heuristically.
     if not explore_pairs:
@@ -429,6 +439,26 @@ def process_model(request: CodeRequest) -> str:
         if request.ir_type.startswith("triton"):
             return compile_triton_ir(request.code, request.ir_type)
 
+        if request.ir_type == "raw_ir" and request.selected_language == "pytorch":
+            # Execute user Python, capture stdout.
+            try:
+                with tempfile.TemporaryDirectory() as tmpdir:
+                    stdout_path = os.path.join(tmpdir, "captured_output.txt")
+                    with open(stdout_path, "w") as f, redirect_stdout(
+                        f
+                    ), redirect_stderr(f):
+                        exec_globals = {}
+                        exec(request.code, exec_globals)
+
+                    with open(stdout_path, "r") as f:
+                        captured = f.read()
+
+                return apply_optional_passes(
+                    captured, build_pipeline(request), request.dump_after_each_opt
+                )
+            except Exception as e:
+                return f"Error executing user code: {str(e)}"
+
         if request.ir_type == "raw_ir":
             return apply_optional_passes(
                 request.code,

src/app/page.js

@@ -78,6 +78,7 @@ const pytorchIROptions = [
   { value: "stablehlo_mlir", label: "StableHLO MLIR" },
   { value: "llvm_mlir", label: "LLVM MLIR" },
   { value: "llvm_ir", label: "LLVM IR" },
+  { value: "raw_ir", label: "Raw IR Output" },
 ];
 
 const tritonIROptions = [
@@ -308,6 +309,7 @@ export default function PyTorchTritonExplorer() {
     const body = {
       code,
       ir_type: window.selectedIR,
+      selected_language: selectedLanguage,
       custom_pipeline: [],
       torch_mlir_opt: window.pipeline
         .filter((p) => p.tool === "torch-mlir-opt")

tests/test_torch_user_controlled_ir_print.py

@@ -0,0 +1,50 @@
+import pytest
+import httpx
+
+API_URL = "http://localhost:8000/generate_ir"
+
+
+def test_user_controlled_ir_print():
+    code = """
+import torch
+import torch.nn as nn
+from torch_mlir import fx
+from torch_mlir.compiler_utils import run_pipeline_with_repro_report
+from torch_mlir.fx import OutputType
+
+class MyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear = nn.Linear(4, 4)
+
+    def forward(self, x):
+        return torch.relu(self.linear(x))
+
+model = MyModel()
+example_input = torch.randn(4, 4)
+module = fx.export_and_import(model, example_input, output_type=OutputType.LINALG_ON_TENSORS)
+print(module)
+"""
+
+    payload = {
+        "code": code,
+        "ir_type": "raw_ir",
+        "custom_pipeline": [],
+        "torch_mlir_opt": "",
+        "mlir_opt": '--one-shot-bufferize="bufferize-function-boundaries"',
+        "mlir_translate": "",
+        "llvm_opt": "",
+        "llc": "",
+        "user_tool": "",
+        "dump_after_each_opt": False,
+    }
+
+    response = httpx.post(API_URL, json=payload)
+    assert response.status_code == 200
+
+    ir = response.json()["output"]
+
+    assert "affine_map" in ir
+    assert "memref.global" in ir
+    assert "arith.constant" in ir
+    assert "linalg.matmul" in ir