Remove print statements from compile utils (#593)

shark/examples/shark_inference/stable_diffusion/stable_args.py

@@ -1,8 +1,4 @@
 import argparse
-from shark.parser import shark_args
-
-# Don't enable nchw-to-nhwc for SD models.
-shark_args.enable_conv_transform = False
 
 p = argparse.ArgumentParser(
     description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter

shark/iree_utils/compile_utils.py

@@ -239,7 +239,6 @@ def compile_module_to_flatbuffer(
     args += get_iree_common_args()
     args += get_model_specific_args()
     args += extra_args
-    print(args)
 
     if frontend in ["tensorflow", "tf"]:
         input_type = "mhlo"

shark/parser.py

@@ -107,7 +107,7 @@ parser.add_argument(
 
 parser.add_argument(
     "--enable_conv_transform",
-    default=True,
+    default=False,
     action="store",
     help="Enables the --iree-flow-enable-conv-nchw-to-nhwc-transform flag.",
 )

shark/shark_benchmark_runner.py

@@ -80,7 +80,6 @@ class SharkBenchmarkRunner(SharkRunner):
             input_tensors,
             mlir_dialect=self.mlir_dialect,
         )
-        print(self.benchmark_cl)
 
     def benchmark_frontend(self, modelname):
         if self.mlir_dialect in ["linalg", "torch"]:
@@ -150,7 +149,6 @@ class SharkBenchmarkRunner(SharkRunner):
             ]
 
     def benchmark_c(self):
-        print(self.benchmark_cl)
         result = run_benchmark_module(self.benchmark_cl)
         print(f"Shark-IREE-C benchmark:{result} iter/second")
         return [f"{result}", f"{1000/result}"]

tank/all_models.csv

@@ -1,35 +1,35 @@
-resnet50,mhlo,tf,1e-2,1e-3,default
-albert-base-v2,mhlo,tf,1e-2,1e-2,default
-roberta-base,mhlo,tf,1e-02,1e-3,default
-bert-base-uncased,mhlo,tf,1e-2,1e-3,default
-camembert-base,mhlo,tf,1e-2,1e-3,default
-dbmdz/convbert-base-turkish-cased,mhlo,tf,1e-2,1e-3,default
-distilbert-base-uncased,mhlo,tf,1e-2,1e-3,default
-facebook/convnext-tiny-224,mhlo,tf,1e-2,1e-3,tf_vit
-funnel-transformer/small,mhlo,tf,1e-2,1e-3,default
-google/electra-small-discriminator,mhlo,tf,1e-2,1e-3,default
-google/mobilebert-uncased,mhlo,tf,1e-2,1e-3,default
-google/vit-base-patch16-224,mhlo,tf,1e-2,1e-3,tf_vit
-hf-internal-testing/tiny-random-flaubert,mhlo,tf,1e-2,1e-3,default
-microsoft/MiniLM-L12-H384-uncased,mhlo,tf,1e-2,1e-3,tf_hf
-microsoft/layoutlm-base-uncased,mhlo,tf,1e-2,1e-3,default
-microsoft/mpnet-base,mhlo,tf,1e-2,1e-2,default
-albert-base-v2,linalg,torch,1e-2,1e-3,default
-alexnet,linalg,torch,1e-2,1e-3,default
-bert-base-cased,linalg,torch,1e-2,1e-3,default
-bert-base-uncased,linalg,torch,1e-2,1e-3,default
-facebook/deit-small-distilled-patch16-224,linalg,torch,1e-2,1e-3,default
-google/vit-base-patch16-224,linalg,torch,1e-2,1e-3,default
-microsoft/beit-base-patch16-224-pt22k-ft22k,linalg,torch,1e-2,1e-3,default
-microsoft/MiniLM-L12-H384-uncased,linalg,torch,1e-2,1e-3,default
-microsoft/resnet-50,linalg,torch,1e-2,1e-3,default
-google/mobilebert-uncased,linalg,torch,1e-2,1e-3,default
-mobilenet_v3_small,linalg,torch,1e-1,1e-2,default
-nvidia/mit-b0,linalg,torch,1e-2,1e-3,default
-resnet101,linalg,torch,1e-2,1e-3,default
-resnet18,linalg,torch,1e-2,1e-3,default
-resnet50,linalg,torch,1e-2,1e-3,default
-squeezenet1_0,linalg,torch,1e-2,1e-3,default
-wide_resnet50_2,linalg,torch,1e-2,1e-3,default
-efficientnet-v2-s,mhlo,tf,1e-02,1e-3,default
-mnasnet1_0,linalg,torch,1e-2,1e-3,default
+resnet50,mhlo,tf,1e-2,1e-3,default,nhcw-nhwc
+albert-base-v2,mhlo,tf,1e-2,1e-2,default,None
+roberta-base,mhlo,tf,1e-02,1e-3,default,nhcw-nhwc
+bert-base-uncased,mhlo,tf,1e-2,1e-3,default,None
+camembert-base,mhlo,tf,1e-2,1e-3,default,None
+dbmdz/convbert-base-turkish-cased,mhlo,tf,1e-2,1e-3,default,nhcw-nhwc
+distilbert-base-uncased,mhlo,tf,1e-2,1e-3,default,None
+facebook/convnext-tiny-224,mhlo,tf,1e-2,1e-3,tf_vit,nhcw-nhwc,
+funnel-transformer/small,mhlo,tf,1e-2,1e-3,default,None
+google/electra-small-discriminator,mhlo,tf,1e-2,1e-3,default,None
+google/mobilebert-uncased,mhlo,tf,1e-2,1e-3,default,None
+google/vit-base-patch16-224,mhlo,tf,1e-2,1e-3,tf_vit,nhcw-nhwc
+hf-internal-testing/tiny-random-flaubert,mhlo,tf,1e-2,1e-3,default,None
+microsoft/MiniLM-L12-H384-uncased,mhlo,tf,1e-2,1e-3,tf_hf,None
+microsoft/layoutlm-base-uncased,mhlo,tf,1e-2,1e-3,default,None
+microsoft/mpnet-base,mhlo,tf,1e-2,1e-2,default,None
+albert-base-v2,linalg,torch,1e-2,1e-3,default,None
+alexnet,linalg,torch,1e-2,1e-3,default,None
+bert-base-cased,linalg,torch,1e-2,1e-3,default,None
+bert-base-uncased,linalg,torch,1e-2,1e-3,default,None
+facebook/deit-small-distilled-patch16-224,linalg,torch,1e-2,1e-3,default,nhcw-nhwc
+google/vit-base-patch16-224,linalg,torch,1e-2,1e-3,default,nhcw-nhwc
+microsoft/beit-base-patch16-224-pt22k-ft22k,linalg,torch,1e-2,1e-3,default,nhcw-nhwc
+microsoft/MiniLM-L12-H384-uncased,linalg,torch,1e-2,1e-3,default,None
+microsoft/resnet-50,linalg,torch,1e-2,1e-3,default,nhcw-nhwc
+google/mobilebert-uncased,linalg,torch,1e-2,1e-3,default,None
+mobilenet_v3_small,linalg,torch,1e-1,1e-2,default,nhcw-nhwc
+nvidia/mit-b0,linalg,torch,1e-2,1e-3,default,None
+resnet101,linalg,torch,1e-2,1e-3,default,nhcw-nhwc
+resnet18,linalg,torch,1e-2,1e-3,default,None
+resnet50,linalg,torch,1e-2,1e-3,default,nhcw-nhwc
+squeezenet1_0,linalg,torch,1e-2,1e-3,default,nhcw-nhwc
+wide_resnet50_2,linalg,torch,1e-2,1e-3,default,nhcw-nhwc
+efficientnet-v2-s,mhlo,tf,1e-02,1e-3,default,nhcw-nhwc
+mnasnet1_0,linalg,torch,1e-2,1e-3,default,nhcw-nhwc

tank/test_models.py

@@ -38,6 +38,7 @@ def load_csv_and_convert(filename, gen=False):
                     "rtol": float(row[3]),
                     "atol": float(row[4]),
                     "out_type": row[5],
+                    "flags": row[6],
                 }
             )
     # This is a pytest workaround
@@ -130,11 +131,10 @@ class SharkModuleTester:
 
         shark_args.local_tank_cache = self.local_tank_cache
         shark_args.update_tank = self.update_tank
-        if self.config["model_name"] in [
-            "alexnet",
-            "resnet18",
-        ] or os.path.isfile("./.use-iree"):
-            shark_args.enable_conv_transform = False
+        if "nhcw-nhwc" in self.config["flags"] and not os.path.isfile(
+            ".use-iree"
+        ):
+            shark_args.enable_conv_transform = True
 
         model, func_name, inputs, golden_out = download_model(
             self.config["model_name"],

web/models/stable_diffusion/stable_args.py

@@ -1,8 +1,4 @@
 import argparse
-from shark.parser import shark_args
-
-# Don't enable nchw-to-nhwc-for SD models.
-shark_args.enable_conv_transpose = False
 
 p = argparse.ArgumentParser(
     description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter