Enable fine-grained control over whether to enable air ping-pong buffering transformation passes (Xilinx#784)

Author: erwei-xilinx

python/air/backend/xrt.py

@@ -48,18 +48,21 @@ def __init__(
         verbose=False,
         experimental_passes=False,
         omit_while_true_loop=False,
+        omit_pingpong=False,
     ):
         """Constructor for XRTBackend
 
         Args:
             verbose: verbose output
             experimental_passes: configure aircc to run additional experimental passes
-            omit_while_true_loop: configure aircc to comit the while true loop it traditionally emits.
+            omit_while_true_loop: configure aircc to omit the while true loop it traditionally emits.
+            omit_pingpong: configure aircc to omit the generation of ping-pong buffering.
         """
         super().__init__()
         self.verbose = verbose
         self.experimental_passes = experimental_passes
         self.omit_while_true_loop = omit_while_true_loop
+        self.omit_pingpong = omit_pingpong
         self.currently_loaded = False
 
     def __del__(self):
@@ -116,6 +119,9 @@ def compile(
             if self.omit_while_true_loop:
                 aircc_options += ["--omit-while-true-loop"]
 
+            if self.omit_pingpong:
+                aircc_options += ["--omit-ping-pong-transform"]
+
             aircc.run(air_module, aircc_options)
 
         return XRTCompileArtifact(xclbin, kernel, insts)

python/air/compiler/aircc/cl_arguments.py

@@ -123,6 +123,13 @@ def parse_args(args=None):
         action="store_true",
         help="By default, aircc may output a while(true) loop around per-core logic. If this option is specified, a while(true) loop will not be added.",
     )
+    parser.add_argument(
+        "--omit-ping-pong-transform",
+        dest="omit_pingpong",
+        default=False,
+        action="store_true",
+        help="Whether to run passes which generate ping-pong buffering patterns or not. This will only change the behavior for this program for npu devices",
+    )
 
     opts = parser.parse_args(args)
     return opts

python/air/compiler/aircc/main.py

@@ -24,24 +24,32 @@
 
 import aie.compiler.aiecc.main as aiecc
 
-EXPERIMENTAL_PASSES = [
-    "air-dependency",
-    "air-dependency-schedule-opt",
-    "air-specialize-dma-broadcast",
-    "air-dma-to-channel",
-    "canonicalize",
-    "cse",
-    "air-dependency-canonicalize",
-    "canonicalize",
-    "cse",
-    "func.func(air-loop-fusion)",
-    "air-label-scf-for-to-ping-pong",
-    "air-ping-pong-transform{keep-memref-dealloc=true}",
-    "air-isolate-async-dma-loop-nests",
-    "air-linalg-to-func",
-    "canonicalize",
-    "cse",
-]
+
+def get_experimental_passes(omit_pingpong=True):
+    EXPERIMENTAL_PASSES = [
+        "air-dependency",
+        "air-dependency-schedule-opt",
+        "air-specialize-dma-broadcast",
+        "air-dma-to-channel",
+        "canonicalize",
+        "cse",
+        "air-dependency-canonicalize",
+        "canonicalize",
+        "cse",
+        "func.func(air-loop-fusion)",
+    ]
+    if not omit_pingpong:
+        EXPERIMENTAL_PASSES += [
+            "air-label-scf-for-to-ping-pong",
+            "air-ping-pong-transform{keep-memref-dealloc=true}",
+        ]
+    EXPERIMENTAL_PASSES += [
+        "air-isolate-async-dma-loop-nests",
+        "air-linalg-to-func",
+        "canonicalize",
+        "cse",
+    ]
+    return EXPERIMENTAL_PASSES
 
 
 def emit_wrapper(herd_name="segment", include_name="aie.inc"):
@@ -378,7 +386,7 @@ def run(mlir_module, args=None):
                 "func.func(air-lower-herd-parallel)",
             ]
             + (
-                EXPERIMENTAL_PASSES
+                get_experimental_passes(opts.omit_pingpong)
                 if "npu" in opts.device and opts.experimental_passes
                 else []
             )