Add opt_causallm_samples.py. (#1916)

* Updates to opt_causallm example

* Fixup opt_perf_comparison.py

* Use same filenames across opt examples.

apps/stable_diffusion/src/utils/resources/opt_flags.json

@@ -11,12 +11,12 @@
     "untuned": {
       "fp16": {
         "default_compilation_flags": [
-          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32}))"
+          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-global-opt-detach-elementwise-from-named-ops,iree-global-opt-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32}))"
         ]
       },
       "fp32": {
         "default_compilation_flags": [
-          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=16}))"
+          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-global-opt-detach-elementwise-from-named-ops,iree-global-opt-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=16}))"
         ]
       }
     }
@@ -28,7 +28,7 @@
         "specified_compilation_flags": {
           "cuda": [],
           "default_device": [
-            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32},iree-linalg-ext-convert-conv2d-to-winograd))"
+            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-global-opt-detach-elementwise-from-named-ops,iree-global-opt-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32},iree-linalg-ext-convert-conv2d-to-winograd))"
           ]
         }
       },
@@ -37,7 +37,7 @@
         "specified_compilation_flags": {
           "cuda": [],
           "default_device": [
-            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=16},iree-linalg-ext-convert-conv2d-to-winograd))"
+            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-global-opt-detach-elementwise-from-named-ops,iree-global-opt-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=16},iree-linalg-ext-convert-conv2d-to-winograd))"
           ]
         }
       }
@@ -45,12 +45,12 @@
     "untuned": {
       "fp16": {
         "default_compilation_flags": [
-          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32},iree-linalg-ext-convert-conv2d-to-winograd))"
+          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-global-opt-detach-elementwise-from-named-ops,iree-global-opt-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32},iree-linalg-ext-convert-conv2d-to-winograd))"
         ]
       },
       "fp32": {
         "default_compilation_flags": [
-          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-flow-detach-elementwise-from-named-ops,iree-flow-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=16},iree-linalg-ext-convert-conv2d-to-winograd))"
+          "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-global-opt-detach-elementwise-from-named-ops,iree-global-opt-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=16},iree-linalg-ext-convert-conv2d-to-winograd))"
         ]
       }
     }

apps/stable_diffusion/src/utils/sd_annotation.py

@@ -203,8 +203,8 @@ def dump_after_mlir(input_mlir, use_winograd):
     if use_winograd:
         preprocess_flag = (
             "--iree-preprocessing-pass-pipeline=builtin.module"
-            "(func.func(iree-flow-detach-elementwise-from-named-ops,"
-            "iree-flow-convert-1x1-filter-conv2d-to-matmul,"
+            "(func.func(iree-global-opt-detach-elementwise-from-named-ops,"
+            "iree-global-opt-convert-1x1-filter-conv2d-to-matmul,"
             "iree-preprocessing-convert-conv2d-to-img2col,"
             "iree-preprocessing-pad-linalg-ops{pad-size=32},"
             "iree-linalg-ext-convert-conv2d-to-winograd))"
@@ -212,8 +212,8 @@ def dump_after_mlir(input_mlir, use_winograd):
     else:
         preprocess_flag = (
             "--iree-preprocessing-pass-pipeline=builtin.module"
-            "(func.func(iree-flow-detach-elementwise-from-named-ops,"
-            "iree-flow-convert-1x1-filter-conv2d-to-matmul,"
+            "(func.func(iree-global-opt-detach-elementwise-from-named-ops,"
+            "iree-global-opt-convert-1x1-filter-conv2d-to-matmul,"
             "iree-preprocessing-convert-conv2d-to-img2col,"
             "iree-preprocessing-pad-linalg-ops{pad-size=32}))"
         )

apps/stable_diffusion/src/utils/utils.py

@@ -895,6 +895,13 @@ def save_output_img(output_img, img_seed, extra_info=None):
         pngInfo = PngImagePlugin.PngInfo()
 
         if args.write_metadata_to_png:
+            # Using a conditional expression caused problems, so setting a new
+            # variable for now.
+            if args.use_hiresfix:
+                png_size_text = f"{args.hiresfix_width}x{args.hiresfix_height}"
+            else:
+                png_size_text = f"{args.width}x{args.height}"
+
             pngInfo.add_text(
                 "parameters",
                 f"{args.prompts[0]}"
@@ -903,7 +910,7 @@ def save_output_img(output_img, img_seed, extra_info=None):
                 f"Sampler: {args.scheduler}, "
                 f"CFG scale: {args.guidance_scale}, "
                 f"Seed: {img_seed},"
-                f"Size: {args.width}x{args.height}, "
+                f"Size: {png_size_text}, "
                 f"Model: {img_model}, "
                 f"VAE: {img_vae}, "
                 f"LoRA: {img_lora}",
@@ -930,8 +937,10 @@ def save_output_img(output_img, img_seed, extra_info=None):
         "CFG_SCALE": args.guidance_scale,
         "PRECISION": args.precision,
         "STEPS": args.steps,
-        "HEIGHT": args.height,
-        "WIDTH": args.width,
+        "HEIGHT": args.height
+        if not args.use_hiresfix
+        else args.hiresfix_height,
+        "WIDTH": args.width if not args.use_hiresfix else args.hiresfix_width,
         "MAX_LENGTH": args.max_length,
         "OUTPUT": out_img_path,
         "VAE": img_vae,
@@ -969,6 +978,10 @@ def get_generation_text_info(seeds, device):
     )
     text_output += (
         f"\nsize={args.height}x{args.width}, "
+        if not args.use_hiresfix
+        else f"\nsize={args.hiresfix_height}x{args.hiresfix_width}, "
+    )
+    text_output += (
         f"batch_count={args.batch_count}, "
         f"batch_size={args.batch_size}, "
         f"max_length={args.max_length}"

apps/stable_diffusion/web/ui/img2img_ui.py

@@ -638,16 +638,6 @@ with gr.Blocks(title="Image-to-Image") as img2img_web:
                         choices=available_devices,
                         allow_custom_value=True,
                     )
-                with gr.Row():
-                    random_seed = gr.Button("Randomize Seed")
-                    random_seed.click(
-                        lambda: -1,
-                        inputs=[],
-                        outputs=[seed],
-                        queue=False,
-                    )
-                    stop_batch = gr.Button("Stop Batch")
-                    stable_diffusion = gr.Button("Generate Image(s)")
 
             with gr.Column(scale=1, min_width=600):
                 with gr.Group():
@@ -667,6 +657,18 @@ with gr.Blocks(title="Image-to-Image") as img2img_web:
                         show_label=False,
                     )
                     img2img_status = gr.Textbox(visible=False)
+                with gr.Row():
+                    stable_diffusion = gr.Button("Generate Image(s)")
+                    random_seed = gr.Button("Randomize Seed")
+                    random_seed.click(
+                        lambda: -1,
+                        inputs=[],
+                        outputs=[seed],
+                        queue=False,
+                    )
+                    stop_batch = gr.Button("Stop Batch")
+                with gr.Row():
+                    blank_thing_for_row = None
                 with gr.Row():
                     img2img_sendto_inpaint = gr.Button(value="SendTo Inpaint")
                     img2img_sendto_outpaint = gr.Button(

apps/stable_diffusion/web/ui/inpaint_ui.py

@@ -514,16 +514,6 @@ with gr.Blocks(title="Inpainting") as inpaint_web:
                         choices=available_devices,
                         allow_custom_value=True,
                     )
-                with gr.Row():
-                    random_seed = gr.Button("Randomize Seed")
-                    random_seed.click(
-                        lambda: -1,
-                        inputs=[],
-                        outputs=[seed],
-                        queue=False,
-                    )
-                    stop_batch = gr.Button("Stop Batch")
-                    stable_diffusion = gr.Button("Generate Image(s)")
 
             with gr.Column(scale=1, min_width=600):
                 with gr.Group():
@@ -543,7 +533,18 @@ with gr.Blocks(title="Inpainting") as inpaint_web:
                         show_label=False,
                     )
                     inpaint_status = gr.Textbox(visible=False)
-
+                with gr.Row():
+                    stable_diffusion = gr.Button("Generate Image(s)")
+                    random_seed = gr.Button("Randomize Seed")
+                    random_seed.click(
+                        lambda: -1,
+                        inputs=[],
+                        outputs=[seed],
+                        queue=False,
+                    )
+                    stop_batch = gr.Button("Stop Batch")
+                with gr.Row():
+                    blank_thing_for_row = None
                 with gr.Row():
                     inpaint_sendto_img2img = gr.Button(value="SendTo Img2Img")
                     inpaint_sendto_outpaint = gr.Button(

apps/stable_diffusion/web/ui/outpaint_ui.py

@@ -540,16 +540,6 @@ with gr.Blocks(title="Outpainting") as outpaint_web:
                         choices=available_devices,
                         allow_custom_value=True,
                     )
-                with gr.Row():
-                    random_seed = gr.Button("Randomize Seed")
-                    random_seed.click(
-                        lambda: -1,
-                        inputs=[],
-                        outputs=[seed],
-                        queue=False,
-                    )
-                    stop_batch = gr.Button("Stop Batch")
-                    stable_diffusion = gr.Button("Generate Image(s)")
 
             with gr.Column(scale=1, min_width=600):
                 with gr.Group():
@@ -569,6 +559,18 @@ with gr.Blocks(title="Outpainting") as outpaint_web:
                         show_label=False,
                     )
                     outpaint_status = gr.Textbox(visible=False)
+                with gr.Row():
+                    stable_diffusion = gr.Button("Generate Image(s)")
+                    random_seed = gr.Button("Randomize Seed")
+                    random_seed.click(
+                        lambda: -1,
+                        inputs=[],
+                        outputs=[seed],
+                        queue=False,
+                    )
+                    stop_batch = gr.Button("Stop Batch")
+                with gr.Row():
+                    blank_thing_for_row = None
                 with gr.Row():
                     outpaint_sendto_img2img = gr.Button(value="SendTo Img2Img")
                     outpaint_sendto_inpaint = gr.Button(value="SendTo Inpaint")

apps/stable_diffusion/web/ui/stablelm_ui.py

@@ -32,36 +32,39 @@ model_map = {
 # NOTE: Each `model_name` should have its own start message
 start_message = {
     "llama2_7b": (
-        "System: You are a helpful, respectful and honest assistant. Always answer "
-        "as helpfully as possible, while being safe.  Your answers should not "
-        "include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal "
-        "content. Please ensure that your responses are socially unbiased and positive "
-        "in nature. If a question does not make any sense, or is not factually coherent, "
-        "explain why instead of answering something not correct. If you don't know the "
-        "answer to a question, please don't share false information."
+        "You are a helpful, respectful and honest assistant. Always answer "
+        "as helpfully as possible, while being safe. Your answers should not "
+        "include any harmful, unethical, racist, sexist, toxic, dangerous, or "
+        "illegal content. Please ensure that your responses are socially "
+        "unbiased and positive in nature. If a question does not make any "
+        "sense, or is not factually coherent, explain why instead of "
+        "answering something not correct. If you don't know the answer "
+        "to a question, please don't share false information."
     ),
     "llama2_13b": (
-        "System: You are a helpful, respectful and honest assistant. Always answer "
-        "as helpfully as possible, while being safe.  Your answers should not "
-        "include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal "
-        "content. Please ensure that your responses are socially unbiased and positive "
-        "in nature. If a question does not make any sense, or is not factually coherent, "
-        "explain why instead of answering something not correct. If you don't know the "
-        "answer to a question, please don't share false information."
+        "You are a helpful, respectful and honest assistant. Always answer "
+        "as helpfully as possible, while being safe. Your answers should not "
+        "include any harmful, unethical, racist, sexist, toxic, dangerous, or "
+        "illegal content. Please ensure that your responses are socially "
+        "unbiased and positive in nature. If a question does not make any "
+        "sense, or is not factually coherent, explain why instead of "
+        "answering something not correct. If you don't know the answer "
+        "to a question, please don't share false information."
     ),
     "llama2_70b": (
-        "System: You are a helpful, respectful and honest assistant. Always answer "
-        "as helpfully as possible, while being safe.  Your answers should not "
-        "include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal "
-        "content. Please ensure that your responses are socially unbiased and positive "
-        "in nature. If a question does not make any sense, or is not factually coherent, "
-        "explain why instead of answering something not correct. If you don't know the "
-        "answer to a question, please don't share false information."
+        "You are a helpful, respectful and honest assistant. Always answer "
+        "as helpfully as possible, while being safe. Your answers should not "
+        "include any harmful, unethical, racist, sexist, toxic, dangerous, or "
+        "illegal content. Please ensure that your responses are socially "
+        "unbiased and positive in nature. If a question does not make any "
+        "sense, or is not factually coherent, explain why instead of "
+        "answering something not correct. If you don't know the answer "
+        "to a question, please don't share false information."
     ),
     "vicuna": (
-        "A chat between a curious user and an artificial intelligence assistant. "
-        "The assistant gives helpful, detailed, and polite answers to the user's "
-        "questions.\n"
+        "A chat between a curious user and an artificial intelligence "
+        "assistant. The assistant gives helpful, detailed, and "
+        "polite answers to the user's questions.\n"
     ),
 }
 
@@ -77,7 +80,10 @@ def create_prompt(model_name, history, prompt_prefix):
         conversation = "".join(
             [f"{B_INST} {item[0]} {E_INST} {item[1]} " for item in history[1:]]
         )
-        msg = f"{B_INST} {B_SYS} {system_message} {E_SYS} {history[0][0]} {E_INST} {history[0][1]} {conversation}"
+        if prompt_prefix:
+            msg = f"{B_INST} {B_SYS}{system_message}{E_SYS}{history[0][0]} {E_INST} {history[0][1]} {conversation}"
+        else:
+            msg = f"{B_INST} {history[0][0]} {E_INST} {history[0][1]} {conversation}"
     elif model_name in ["vicuna"]:
         conversation = "".join(
             [

apps/stable_diffusion/web/ui/txt2img_ui.py

@@ -140,6 +140,11 @@ def txt2img_inf(
         args.max_length = max_length
         args.height = height
         args.width = width
+        args.use_hiresfix = use_hiresfix
+        args.hiresfix_height = hiresfix_height
+        args.hiresfix_width = hiresfix_width
+        args.hiresfix_strength = hiresfix_strength
+        args.resample_type = resample_type
         args.device = device.split("=>", 1)[1].strip()
         args.iree_vulkan_target_triple = init_iree_vulkan_target_triple
         args.iree_metal_target_platform = init_iree_metal_target_platform

apps/stable_diffusion/web/ui/upscaler_ui.py

@@ -533,16 +533,6 @@ with gr.Blocks(title="Upscaler") as upscaler_web:
                         choices=available_devices,
                         allow_custom_value=True,
                     )
-                with gr.Row():
-                    random_seed = gr.Button("Randomize Seed")
-                    random_seed.click(
-                        lambda: -1,
-                        inputs=[],
-                        outputs=[seed],
-                        queue=False,
-                    )
-                    stop_batch = gr.Button("Stop Batch")
-                    stable_diffusion = gr.Button("Generate Image(s)")
 
             with gr.Column(scale=1, min_width=600):
                 with gr.Group():
@@ -562,7 +552,18 @@ with gr.Blocks(title="Upscaler") as upscaler_web:
                         show_label=False,
                     )
                     upscaler_status = gr.Textbox(visible=False)
-
+                with gr.Row():
+                    stable_diffusion = gr.Button("Generate Image(s)")
+                    random_seed = gr.Button("Randomize Seed")
+                    random_seed.click(
+                        lambda: -1,
+                        inputs=[],
+                        outputs=[seed],
+                        queue=False,
+                    )
+                    stop_batch = gr.Button("Stop Batch")
+                with gr.Row():
+                    blank_thing_for_row = None
                 with gr.Row():
                     upscaler_sendto_img2img = gr.Button(value="SendTo Img2Img")
                     upscaler_sendto_inpaint = gr.Button(value="SendTo Inpaint")

tank/examples/opt/opt_causallm.py

@@ -1,3 +1,4 @@
+import argparse
 import os
 import torch
 import numpy as np
@@ -9,22 +10,18 @@ from shark.iree_utils._common import (
 from shark.shark_inference import SharkInference
 from shark.shark_importer import import_with_fx
 from transformers import AutoTokenizer, OPTForCausalLM
-
-OPT_MODEL = "opt-1.3b"
-OPT_FS_NAME = "opt-1_3b"
-MAX_SEQUENCE_LENGTH = 128
-MAX_NEW_TOKENS = 60
+from typing import Iterable
 
 
-def create_module(model_name, tokenizer, device):
-    opt_base_model = OPTForCausalLM.from_pretrained("facebook/" + model_name)
+def create_module(model_name, tokenizer, device, args):
+    opt_base_model = OPTForCausalLM.from_pretrained(model_name)
     opt_base_model.eval()
     opt_model = OPTForCausalLMModel(opt_base_model)
     encoded_inputs = tokenizer(
         "What is the meaning of life?",
         padding="max_length",
         truncation=True,
-        max_length=MAX_SEQUENCE_LENGTH,
+        max_length=args.max_seq_len,
         return_tensors="pt",
     )
     inputs = (
@@ -33,8 +30,11 @@ def create_module(model_name, tokenizer, device):
     )
     # np.save("model_inputs_0.npy", inputs[0])
     # np.save("model_inputs_1.npy", inputs[1])
+    opt_fs_name = "-".join(
+        "_".join(args.model_name.split("/")[1].split("-")).split(".")
+    )
 
-    mlir_path = f"./{OPT_FS_NAME}_causallm_{MAX_SEQUENCE_LENGTH}_torch.mlir"
+    mlir_path = f"./{opt_fs_name}_causallm_{args.max_seq_len}_torch.mlir"
     if os.path.isfile(mlir_path):
         print(f"Found .mlir from {mlir_path}")
     else:
@@ -42,7 +42,7 @@ def create_module(model_name, tokenizer, device):
             model=opt_model,
             inputs=inputs,
             is_f16=False,
-            model_name=OPT_FS_NAME,
+            model_name=opt_fs_name,
             return_str=True,
         )
         with open(mlir_path, "w") as f:
@@ -57,7 +57,7 @@ def create_module(model_name, tokenizer, device):
         is_benchmark=False,
     )
 
-    vmfb_name = f"{OPT_FS_NAME}_causallm_{MAX_SEQUENCE_LENGTH}_torch_{device}"
+    vmfb_name = f"{opt_fs_name}_causallm_{args.max_seq_len}_torch_cpu"
     shark_module.save_module(module_name=vmfb_name, debug=False)
     vmfb_path = vmfb_name + ".vmfb"
     return vmfb_path
@@ -71,11 +71,11 @@ def shouldStop(tokens):
     return False
 
 
-def generate_new_token(shark_model, tokenizer, new_text):
+def generate_new_token(shark_module, tokenizer, new_text, max_seq_len: int):
     model_inputs = tokenizer(
         new_text,
         padding="max_length",
-        max_length=MAX_SEQUENCE_LENGTH,
+        max_length=max_seq_len,
         truncation=True,
         return_tensors="pt",
     )
@@ -84,7 +84,7 @@ def generate_new_token(shark_model, tokenizer, new_text):
         model_inputs["attention_mask"],
     )
     sum_attentionmask = torch.sum(model_inputs.attention_mask)
-    output = shark_model("forward", inputs)
+    output = shark_module("forward", inputs)
     output = torch.FloatTensor(output[0])
     next_toks = torch.topk(output, 1)
     stop_generation = False
@@ -104,39 +104,87 @@ def generate_new_token(shark_model, tokenizer, new_text):
     return ret_dict
 
 
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--max-seq-len", type=int, default=32)
+    parser.add_argument(
+        "--model-name",
+        help="Model name",
+        type=str,
+        choices=[
+            "facebook/opt-125m",
+            "facebook/opt-350m",
+            "facebook/opt-1.3b",
+            "facebook/opt-6.7b",
+        ],
+        default="facebook/opt-1.3b",
+    )
+    parser.add_argument(
+        "--recompile",
+        help="If set, recompiles MLIR -> .vmfb",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+    )
+    parser.add_argument(
+        "--plugin-path",
+        help="path to executable plugin",
+        type=str,
+        default=None,
+    )
+    args = parser.parse_args()
+    print("args={}".format(args))
+    return args
+
+
+def generate_tokens(
+    opt_shark_module: "SharkInference",
+    tokenizer,
+    input_text: str,
+    max_output_len: int,
+    print_intermediate_results: True,
+) -> Iterable[str]:
+    words_list = []
+    new_text = input_text
+    try:
+        for _ in range(max_output_len):
+            generated_token_op = generate_new_token(
+                opt_shark_module, tokenizer, new_text, max_output_len
+            )
+            detok = generated_token_op["detok"]
+            if generated_token_op["stop_generation"]:
+                break
+            if print_intermediate_results:
+                print(detok, end="", flush=True)
+            words_list.append(detok)
+            if detok == "":
+                break
+            new_text += detok
+    except KeyboardInterrupt as e:
+        print("Exiting token generation.")
+    return words_list
+
+
 if __name__ == "__main__":
-    tokenizer = AutoTokenizer.from_pretrained(
-        "facebook/" + OPT_MODEL, use_fast=False
+    args = parse_args()
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name, use_fast=False)
+    opt_fs_name = "-".join(
+        "_".join(args.model_name.split("/")[1].split("-")).split(".")
     )
-    vmfb_path = (
-        f"./{OPT_FS_NAME}_causallm_{MAX_SEQUENCE_LENGTH}_torch_cpu-task.vmfb"
+    vmfb_path = f"./{opt_fs_name}_causallm_{args.max_seq_len}_torch_cpu.vmfb"
+    if args.plugin_path is not None:
+        rt_flags = [f"--executable_plugin={args.plugin_path}"]
+    else:
+        rt_flags = []
+    opt_shark_module = SharkInference(
+        mlir_module=None, device="cpu-task", rt_flags=rt_flags
     )
-    opt_shark_module = SharkInference(mlir_module=None, device="cpu-task")
     if os.path.isfile(vmfb_path):
         opt_shark_module.load_module(vmfb_path)
     else:
-        vmfb_path = create_module(OPT_MODEL, tokenizer, "cpu-task")
+        vmfb_path = create_module(args.model_name, tokenizer, "cpu-task", args)
         opt_shark_module.load_module(vmfb_path)
     while True:
-        try:
-            new_text = input("Give me a sentence to complete:")
-            new_text_init = new_text
-            words_list = []
-
-            for i in range(MAX_NEW_TOKENS):
-                generated_token_op = generate_new_token(
-                    opt_shark_module, tokenizer, new_text
-                )
-                detok = generated_token_op["detok"]
-                stop_generation = generated_token_op["stop_generation"]
-                if stop_generation:
-                    break
-                print(detok, end="", flush=True)
-                words_list.append(detok)
-                if detok == "":
-                    break
-                new_text = new_text + detok
-
-        except KeyboardInterrupt:
-            print("Exiting program.")
-            break
+        input_text = input("Give me a sentence to complete:")
+        generate_tokens(
+            opt_shark_module, tokenizer, input_text, args.max_seq_len
+        )

tank/examples/opt/opt_causallm_samples.py

@@ -0,0 +1,74 @@
+import argparse
+import os
+
+import opt_causallm
+import opt_util
+
+from shark.shark_inference import SharkInference
+from transformers import AutoTokenizer, OPTForCausalLM
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--max-seq-len", type=int, default=32)
+    parser.add_argument(
+        "--model-name",
+        help="Model name",
+        type=str,
+        choices=[
+            "facebook/opt-125m",
+            "facebook/opt-350m",
+            "facebook/opt-1.3b",
+            "facebook/opt-6.7b",
+        ],
+        default="facebook/opt-1.3b",
+    )
+    parser.add_argument(
+        "--recompile",
+        help="If set, recompiles MLIR -> .vmfb",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+    )
+    parser.add_argument(
+        "--plugin-path",
+        help="path to executable plugin",
+        type=str,
+        default=None,
+    )
+    args = parser.parse_args()
+    print("args={}".format(args))
+    return args
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name, use_fast=False)
+    opt_fs_name = "-".join(
+        "_".join(args.model_name.split("/")[1].split("-")).split(".")
+    )
+    vmfb_path = f"./{opt_fs_name}_causallm_{args.max_seq_len}_torch_cpu.vmfb"
+    if args.plugin_path is not None:
+        rt_flags = [f"--executable_plugin={args.plugin_path}"]
+    else:
+        rt_flags = []
+    opt_shark_module = SharkInference(
+        mlir_module=None, device="cpu-task", rt_flags=rt_flags
+    )
+    if os.path.isfile(vmfb_path):
+        opt_shark_module.load_module(vmfb_path)
+    else:
+        vmfb_path = opt_causallm.create_module(
+            args.model_name, tokenizer, "cpu-task", args
+        )
+        opt_shark_module.load_module(vmfb_path)
+
+    for prompt in opt_util.PROMPTS:
+        print("\n\nprompt: {}".format(prompt))
+        response = opt_causallm.generate_tokens(
+            opt_shark_module,
+            tokenizer,
+            prompt,
+            args.max_seq_len,
+            print_intermediate_results=False,
+        )
+        print("reponse: {}".format("".join(response)))

tank/examples/opt/opt_perf_comparison.py

@@ -22,6 +22,7 @@ import time
 import numpy as np
 from typing import Tuple
 
+from opt_util import PROMPTS
 from shark.shark_inference import SharkInference
 from shark.shark_importer import import_with_fx
 from transformers import AutoTokenizer, OPTForCausalLM
@@ -44,19 +45,6 @@ REPORT_LOAD_VIRTUAL_MEMORY_MB = "load_virtual_MB"
 REPORT_RUN_PHYSICAL_MEMORY_MB = "run_physical_MB"
 REPORT_RUN_VIRTUAL_MEMORY_MB = "run_virtual_MB"
 
-PROMPTS = [
-    "What is the meaning of life?",
-    "Tell me something you don't know.",
-    "What does Xilinx do?",
-    "What is the mass of earth?",
-    "What is a poem?",
-    "What is recursion?",
-    "Tell me a one line joke.",
-    "Who is Gilgamesh?",
-    "Tell me something about cryptocurrency.",
-    "How did it all begin?",
-]
-
 ModelWrapper = collections.namedtuple("ModelWrapper", ["model", "tokenizer"])
 
 
@@ -147,7 +135,7 @@ def load_shark_model(
     plugin_path: str = [],
 ) -> ModelWrapper:
     opt_fs_name = get_opt_fs_name(model_name)
-    vmfb_name = f"{opt_fs_name}_causallm_{max_seq_len}_torch_{DEVICE}_tiled_ukernels.vmfb"
+    vmfb_name = f"{opt_fs_name}_causallm_{max_seq_len}_torch_{DEVICE}.vmfb"
     tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
     if recompile_shark or not os.path.isfile(vmfb_name):
         print(f"vmfb not found. compiling and saving to {vmfb_name}")
@@ -344,7 +332,7 @@ def parse_args():
         default=PLATFORM_SHARK,
     )
     parser.add_argument(
-        "--plugin_path",
+        "--plugin-path",
         help="path to executable plugin",
         type=str,
         default=None,

tank/examples/opt/opt_util.py

@@ -0,0 +1,12 @@
+PROMPTS = [
+    "What is the meaning of life?",
+    "Tell me something you don't know.",
+    "What does Xilinx do?",
+    "What is the mass of earth?",
+    "What is a poem?",
+    "What is recursion?",
+    "Tell me a one line joke.",
+    "Who is Gilgamesh?",
+    "Tell me something about cryptocurrency.",
+    "How did it all begin?",
+]