Move location and give more helpful output

Author: rolfmorel

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "ingress/KernelBench/KernelBench"]
+	path = ingress/KernelBench/KernelBench
+	url = https://github.com/ScalingIntelligence/KernelBench.git

ingress/KernelBench/KernelBench

@@ -1,12 +1,25 @@
 #!/usr/bin/env python3
 
+import sys
 from pathlib import Path
 
 from mlir import ir, passmanager
 from lighthouse.ingress import torch as torch_ingress
 
 
 kernels_as_pytorch_folder = Path(__file__).parent / "KernelBench" / "KernelBench"
+
+if not (kernels_as_pytorch_folder.exists() and kernels_as_pytorch_folder.is_dir()):
+    print(
+        "ERROR: KernelBench repo not found.\n"
+        "NOTE: Pull in dependency with: git submodule update "
+        + str(kernels_as_pytorch_folder.parent.relative_to(Path.cwd()))
+        + "",
+        file=sys.stderr,
+    )
+    sys.exit(1)
+
+
 kernels_as_pytorch_level1 = kernels_as_pytorch_folder / "level1"
 kernels_as_pytorch_level2 = kernels_as_pytorch_folder / "level2"
 
@@ -105,6 +118,7 @@
 pm = passmanager.PassManager(context=ctx)
 pm.add("linalg-specialize-generic-ops")
 
+print("Output directory:", kernels_as_mlir_folder)
 for pytorch_level, mlir_level in (
     (kernels_as_pytorch_level1, kernels_as_mlir_level1),
     (kernels_as_pytorch_level2, kernels_as_mlir_level2),
@@ -133,16 +147,14 @@
         mlir_kernel = torch_ingress.import_from_file(
             kernel_pytorch_file, ir_context=ctx
         )
+        assert isinstance(mlir_kernel, ir.Module)
 
-        before_clean_up = "//" + str(mlir_kernel)[:-1].replace("\n", "\n//") + "\n"
         try:
             pm.run(mlir_kernel.operation)  # cleanup
         except Exception as e:
-            print(f"Error: got the following error cleaning up {kernel_name}")
+            print(f"Error: got the following error cleaning up '{kernel_name}'")
             raise e
 
         with kernel_as_mlir_path.open("w") as f:
-            print("// Torch-MLIR output:", file=f)
-            print(before_clean_up, file=f)
-            print("// MLIR output after clean-up:", file=f)
+            print("// MLIR output after conversion and clean-up:", file=f)
             print(mlir_kernel, file=f)