[frontend] Add Initial Graph Infra.

Co-authored-by: xtayex <liyuliang2022_sei@stu.ecnu.edu.cn>
Co-authored-by: zhanghb97 <hongbin2019@iscas.ac.cn>

examples/BuddyBert/import-bert.py

@@ -46,12 +46,16 @@
     "attention_mask": torch.tensor([[1 for _ in range(5)]], dtype=torch.int64),
 }
 with torch.no_grad():
-    module, params = dynamo_compiler.importer(model, **inputs)
+    graphs = dynamo_compiler.importer(model, **inputs)
 
+assert len(graphs) == 1
+graph = graphs[0]
+params = dynamo_compiler.imported_params[graph]
+graph.lower_to_top_level_ir(do_params_pack=True)
 current_path = os.path.dirname(os.path.abspath(__file__))
 
 with open(Path(current_path) / "bert.mlir", "w") as module_file:
-    module_file.write(str(module))
+    module_file.write(str(graph._imported_module))
 
 float32_param = np.concatenate(
     [param.detach().numpy().reshape([-1]) for param in params[:-1]]

examples/BuddyGraph/README.md

@@ -0,0 +1,23 @@
+# Buddy Graph Representation Examples
+
+## Run the Examples
+
+0. Enter your Python Env
+```
+(base)$ conda activate buddy
+(buddy)$ ...
+```
+1. Build Python Packages
+2. Configure Python Path
+```
+(buddy)$ cd buddy-mlir/build
+(buddy)$ export BUDDY_MLIR_BUILD_DIR=$PWD
+(buddy)$ export LLVM_MLIR_BUILD_DIR=$PWD/../llvm/build
+(buddy)$ export PYTHONPATH=${LLVM_MLIR_BUILD_DIR}/tools/mlir/python_packages/mlir_core:${BUDDY_MLIR_BUILD_DIR}/python_packages:${PYTHONPATH}
+
+```
+3. Run the Examples
+```
+(buddy)$ cd examples/BuddyGraph
+(buddy)$ python import-dynamo-break.py
+```

examples/BuddyGraph/import-dynamo-break.py

@@ -0,0 +1,63 @@
+# ===- import-dynamo-break.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.
+#
+# ===---------------------------------------------------------------------------
+#
+# The example for dynamo graph break, import, and execute.
+#
+# ===---------------------------------------------------------------------------
+
+import torch
+import torch._dynamo as dynamo
+from torch._inductor.decomposition import decompositions as inductor_decomp
+from torch._functorch.aot_autograd import aot_autograd_decompositions
+
+from buddy.compiler.frontend import DynamoCompiler
+from buddy.compiler.ops import tosa
+
+
+class TestModule(torch.nn.Module):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+
+    def forward(self, b, c):
+        if torch.nn.functional.silu(b)[0][0]:
+            return torch.add(b, c)
+        else:
+            return torch.matmul(b, c)
+
+# Define a PyTorch model and run it with PyTorch runtime.
+model = TestModule()
+a, b = torch.randn((1024, 1024)), torch.randn((1024, 1024))
+print(model(a, b))
+
+# JIT Mode
+# Initialize Buddy Dynamo Compiler to compile and execute the PyTorch model.
+dynamo_compiler = DynamoCompiler(
+    primary_registry=tosa.ops_registry,
+    aot_autograd_decomposition=aot_autograd_decompositions
+)
+model_opt = torch.compile(model, backend=dynamo_compiler)
+print(model_opt(a, b))
+
+torch._dynamo.reset()
+
+# AOT Mode
+# Import PyTorch model to Buddy Graph and MLIR/LLVM IR.
+graphs = dynamo_compiler.importer(
+    model, a, b
+)
+for g in graphs:
+    g.lower_to_top_level_ir()
+    print(g._imported_module)

examples/BuddyLlama/import-llama2.py

@@ -19,11 +19,13 @@
 # ===---------------------------------------------------------------------------
 
 import os
+import time
 
 import numpy
 import torch
 from transformers import LlamaForCausalLM, LlamaTokenizer
 from torch._functorch.aot_autograd import aot_autograd_decompositions
+from torch._inductor.decomposition import decompositions as inductor_decomp
 
 from buddy.compiler.frontend import DynamoCompiler
 from buddy.compiler.ops import tosa
@@ -44,19 +46,22 @@
 # Initialize Dynamo Compiler with specific configurations as an importer.
 dynamo_compiler = DynamoCompiler(
     primary_registry=tosa.ops_registry,
-    aot_autograd_decomposition=aot_autograd_decompositions,
+    aot_autograd_decomposition=inductor_decomp,
 )
 
 # Import the model into MLIR module and parameters.
 with torch.no_grad():
-    gm, params = dynamo_compiler.importer(
-        model, torch.tensor([[1 for _ in range(40)]], dtype=torch.int64)
-    )
+    data = torch.tensor([[1 for i in range(40)]], dtype=torch.int64)
+    graphs = dynamo_compiler.importer(model, data)
 
+assert len(graphs) == 1
+graph = graphs[0]
+params = dynamo_compiler.imported_params[graph]
+graph.lower_to_top_level_ir(True)
 path_prefix = os.path.dirname(os.path.abspath(__file__))
 # Write the MLIR module to the file.
 with open(os.path.join(path_prefix, "llama.mlir"), "w") as module_file:
-    print(gm, file=module_file)
+    print(graph._imported_module, file=module_file)
 
 # Concatenate all parameters into a single numpy array and write to a file.
 all_param = numpy.concatenate(

examples/BuddyLlama/llama-main.cpp

@@ -18,12 +18,9 @@
 #include <buddy/LLM/TextContainer.h>
 #include <chrono>
 #include <cstddef>
-#include <cstdint>
 #include <filesystem>
 #include <fstream>
 #include <iostream>
-#include <limits>
-#include <type_traits>
 
 using namespace buddy;
 

examples/BuddyPython/bert.py

@@ -15,6 +15,10 @@
 text = "Replace me by any text you'd like."
 encoded_text = tokenizer(text, return_tensors="pt")
 with torch.no_grad():
-    module, params = dynamo_compiler.importer(model, **encoded_text)
-    print(module)
-    print(params)
+    graphs = dynamo_compiler.importer(model, **encoded_text)
+
+graph = graphs[0]
+params = dynamo_compiler.imported_params[graph]    
+graph.lower_to_top_level_ir(do_params_pack=True)
+print(graph._imported_module)
+print(params)

examples/BuddyPython/module_gen.py

@@ -43,23 +43,12 @@ def foo(x, y):
     aot_autograd_decomposition=inductor_decomp,
 )
 
-# The first way to generate an MLIR Module:
-# Pass the function and input data to the dynamo compiler's importer,
-# and accepts the generated module and weight parameters.
-module, params = dynamo_compiler.importer(foo, *(float32_in1, float32_in2))
-
-print(module)
-print(params)
-
-# The second way to generate an MLIR Module:
-# Execute the target function using a define-by-run style,
-# and get the module and weight parameters from the dynamo compiler's attribute.
-foo_mlir = dynamo.optimize(dynamo_compiler)(foo)
-
-foo_mlir(float32_in1, float32_in2)
-print(dynamo_compiler.imported_module)
-print(dynamo_compiler.imported_params)
-
-foo_mlir(int32_in1, int32_in2)
-print(dynamo_compiler.imported_module)
-print(dynamo_compiler.imported_params)
+# Pass the function and input data to the dynamo compiler's importer, the 
+# importer will first build a graph. Then, lower the graph to top-level IR. 
+# (tosa, linalg, etc.). Finally, accepts the generated module and weight parameters.
+graphs = dynamo_compiler.importer(foo, *(float32_in1, float32_in2))
+graph = graphs[0]
+graph.lower_to_top_level_ir(do_params_pack=True)
+
+print(graph._imported_module)
+print(dynamo_compiler.imported_params[graph])

examples/BuddyResNet18/import-resnet18.py

@@ -0,0 +1,45 @@
+# ===- import-resnet18.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 test of resnet18 model.
+#
+# ===---------------------------------------------------------------------------
+
+import torch
+import torchvision
+from torch._inductor.decomposition import decompositions as inductor_decomp
+
+from buddy.compiler.frontend import DynamoCompiler
+from buddy.compiler.ops import tosa
+
+
+model = torchvision.models.resnet18()
+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
+graphs[0].lower_to_top_level_ir(do_params_pack=True)
+print(graphs[0]._imported_module)

examples/lit.cfg.py

@@ -36,6 +36,7 @@
 # subdirectories contain auxiliary inputs for various tests in their parent
 # directories.
 config.excludes = [
+    'BuddyBert',
     'BuddyLlama',
     'BuddyBert',
     'ConvOpt',

frontend/Interfaces/buddy/LLM/TextContainer.h

@@ -325,7 +325,7 @@ template <typename T, size_t N> std::string Text<T, N>::revertLlama() {
   const int CLS_ID = 1;
   const int SEP_ID = 2;
 
-  for (size_t i = 0; i < this->getSize(); i++) {
+  for (size_t i = 0; i < this->tokenCnt; i++) {
     int id = this->aligned[i];
     if (id == PAD_ID || id == CLS_ID)
       continue;