fix mlperftiny flow

KULeuven-MICAS · Jan 20, 2025 · 2cf1ab6 · 2cf1ab6
1 parent 31a0a37
commit 2cf1ab6
Show file tree

Hide file tree

Showing 5 changed files with 7 additions and 28 deletions.
diff --git a/compiler/transforms/frontend/preprocess_mlir.py b/compiler/transforms/frontend/preprocess_mlir.py
@@ -42,4 +42,4 @@ class PreprocessPass(ModulePass):
 
     def apply(self, ctx: MLContext, op: builtin.ModuleOp) -> None:
         for flags in MLIR_PREPOC_FLAGS:
-            MLIROptPass(generic=False, arguments=flags).apply(ctx, op)
+            MLIROptPass(generic=True, arguments=flags).apply(ctx, op)
diff --git a/compiler/transforms/snax_bufferize.py b/compiler/transforms/snax_bufferize.py
@@ -76,6 +76,7 @@ class SnaxBufferize(ModulePass):
 
     mlir_bufferization_pass = MLIROptPass(
         arguments=(
+            "--empty-tensor-to-alloc-tensor",
             "--one-shot-bufferize=bufferize-function-boundaries allow-return-allocs-from-loops allow-unknown-ops"
             + " function-boundary-type-conversion=identity-layout-map",
             "--mlir-print-op-generic",

diff --git a/kernels/mlperf_tiny/Snakefile b/kernels/mlperf_tiny/Snakefile
@@ -25,7 +25,7 @@ use rule * from tf_rules as tf_*
 
 
 nets = ["pretrainedResnet_quant", "kws_ref_model", "vww_96_int8", "ad01_int8"]
-extensions = [".o", ".no-snax-opt.o"]
+extensions = [".o"]
 
 
 rule all:
@@ -49,11 +49,11 @@ rule snax_opt_mlir:
     Similar to snax_opt_mlir, but with print-op-generic
     """
     input:
-        "{file}.preprocfinal.mlir",
+        "{file}.mlir",
     output:
         temp("{file}.snax-opt.mlir"),
     shell:
-        "{config[snax-opt]} -p {config[snaxoptflags]} --print-op-generic -o {output} {input}"
+        "{config[snax-opt]} -p {config[snaxoptflags]} --allow-unregistered-dialect --print-op-generic -o {output} {input}"
 
 
 # Override default rule, for output not to be temporary
@@ -64,13 +64,3 @@ rule compile_llvm_module:
         "{file}.o",
     shell:
         "{config[cc]} {config[clangflags]} -x ir -c {input} -o {output}"
-
-
-# Special preprocessing for anomaly detection
-rule preprocess_anomaly_detection:
-    input:
-        "ad01_int8.mlir",
-    output:
-        "ad01_int8.preprocfinal.mlir",
-    shell:
-        "{config[snax-opt]} -p preprocess-mlperftiny --print-op-generic --allow-unregistered-dialect -o {output} {input}"
diff --git a/util/snake/configs.py b/util/snake/configs.py
@@ -59,6 +59,8 @@ def get_mlperf_tiny_config():
         {
             "snaxoptflags": ",".join(
                 [
+                    "preprocess",
+                    "snax-bufferize",
                     "dispatch-kernels",
                     "set-memory-space",
                     "set-memory-layout",

diff --git a/util/snake/default_rules.smk b/util/snake/default_rules.smk
@@ -25,20 +25,6 @@ rule postprocess_mlir:
         "{config[mlir-opt]} {config[mlirpostprocflags]} -o {output} {input}"
 
 
-rule postprocess_no_snax_mlir:
-    """
-    Bypass snax-opt, and apply various postprocessing transformations to mlir files with upstream mlir.
-    Goal is to lower everything to LLVM dialect after this step.
-    Options controlled with `mlirpostprocflags` defined in config.
-    """
-    input:
-        "{file}.preprocfinal.mlir",
-    output:
-        temp("{file}.no-snax-opt.ll.mlir"),
-    shell:
-        "{config[mlir-opt]} {config[mlirpostprocflags]} -o {output} {input}"
-
-
 rule translate_mlir:
     """
     Translate MLIR LLVM dialect to actual LLVM.