Add post-transform memory space override, rms_norm blocked on alloc memory space

Add air-override-memref-memory-space{scope=herd memory-space=2} and
{scope=func memory-space=1} as post-transform passes in driver.py to fix
memory spaces for allocs created by one_shot_bufferize.

rms_norm: simplified transform (no pad, no L1 alloc) that tiles the fused
output with forall[32] and creates a 2x1 herd. Blocked on mlir-air #1384
(herd block arg type propagation) and #1387 (air-par-to-launch hang when
types are fixed). Filed 7 mlir-air issues total (#1367-#1387).

Existing elementwise examples (relu, sigmoid, silu, gelu, leaky_relu)
verified working with the updated driver.py.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: erwei-xilinx

amd_triton_npu/backend/driver.py

@@ -290,72 +290,7 @@ def _ttshared_to_air(mod, gridX, gridY, gridZ):
             context=air_context,
         )
         pm_ms.run(air_module.operation)
-        # Step 2.6: DISABLED -- fixing herd block arg types causes
-        # air-par-to-launch to hang (applyPatternsGreedily convergence issue).
-        # See mlir-air #1384. Waiting for upstream fix.
-        if False:
-            import re as _re
-            import subprocess as _sp
-        _ir = str(air_module)
-        _hm = _re.search(
-            r"air\.herd\s+@\w+\s+tile\s*\([^)]+\)\s+in\s*\([^)]+\)\s+"
-            r"args\(([^)]+)\)\s*:\s*([^{]+)\{",
-            _ir,
-        )
-        if _hm:
-            _types = []
-            _d, _c = 0, ""
-            for ch in _hm.group(2):
-                if ch == "<": _d += 1
-                elif ch == ">": _d -= 1
-                if ch == "," and _d == 0:
-                    _types.append(_c.strip()); _c = ""
-                else: _c += ch
-            if _c.strip(): _types.append(_c.strip())
-
-            _hs = _hm.end()
-            _d, _he = 1, _hs
-            for i in range(_hs, len(_ir)):
-                if _ir[i] == "{": _d += 1
-                elif _ir[i] == "}":
-                    _d -= 1
-                    if _d == 0: _he = i + 1; break
-
-            _hb = _ir[_hs:_he]
-            _changed = False
-            for _t in _types:
-                if "memref<" not in _t: continue
-                _ms = _re.search(r",\s*(\d+\s*:\s*i\d+)\s*>$", _t)
-                if not _ms: continue
-                _ms_s = _ms.group(1)
-                _base = _t[:_ms.start()] + ">"
-                if _base == _t: continue
-                _hb = _hb.replace(_base, _t)
-                _el = _re.search(r"memref<\d+x(\w+)>", _base)
-                if _el:
-                    for _sm in set(_re.findall(
-                        r"memref<\d+x" + _re.escape(_el.group(1))
-                        + r",\s*strided<[^\]]*\],\s*offset:\s*\??>>", _hb
-                    )):
-                        if _ms_s not in _sm:
-                            _hb = _hb.replace(_sm, _sm[:-1] + ", " + _ms_s + ">")
-                _changed = True
-
-            if _changed:
-                _fixed = _ir[:_hs] + _hb + _ir[_he:]
-                _tmpf = os.path.join(tmpdir, "step26_fixed.mlir")
-                with open(_tmpf, "w") as f:
-                    f.write(_fixed)
-                # Re-parse via air-opt (which has all dialects registered)
-                _air_opt = _get_air_opt_path()
-                _result = _sp.run(
-                    [_air_opt, _tmpf, "--allow-unregistered-dialect"],
-                    capture_output=True, text=True, timeout=30,
-                )
-                if _result.returncode == 0:
-                    air_module = Module.parse(
-                        _result.stdout, context=air_context
-                    )
+        # Step 2.6: REMOVED -- type fix moved to step 3.5 (after air-par-to-launch).
         # MLIR-AIR compilation step 3: converting to AIR
         pipeline = (
             "builtin.module("

examples/rms_norm/transform_aie2.mlir

@@ -45,15 +45,7 @@ module attributes {transform.with_named_sequence} {
     } : !transform.any_op
     transform.apply_cse to %func_2 : !transform.any_op
 
-    // Phase 4: Promote generic inside forall to L1 (no pad needed)
-    %gen_in_forall = transform.structured.match ops{["linalg.generic"]} in %forall : (!transform.any_op) -> !transform.any_op
-    %gen_l1_buf, %gen_l1_new = transform.structured.bufferize_to_allocation %gen_in_forall
-        {memory_space = 2, bufferize_destination_only, emit_dealloc} : !transform.any_op
-
-    %func_3 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-    transform.apply_patterns to %func_3 { transform.apply_patterns.canonicalization } : !transform.any_op
-    transform.apply_cse to %func_3 : !transform.any_op
-
+    // Phase 4: No L1 alloc -- let post-passes handle memory hierarchy.
     // Phase 5: Bufferize
     // Memory spaces for remaining allocs handled by post-transform overrides
     // in driver.py (scope=herd→2, scope=func→1 with exclusive scopes).

examples/rms_norm/transform_aie2p.mlir

@@ -45,15 +45,7 @@ module attributes {transform.with_named_sequence} {
     } : !transform.any_op
     transform.apply_cse to %func_2 : !transform.any_op
 
-    // Phase 4: Promote generic inside forall to L1 (no pad needed)
-    %gen_in_forall = transform.structured.match ops{["linalg.generic"]} in %forall : (!transform.any_op) -> !transform.any_op
-    %gen_l1_buf, %gen_l1_new = transform.structured.bufferize_to_allocation %gen_in_forall
-        {memory_space = 2, bufferize_destination_only, emit_dealloc} : !transform.any_op
-
-    %func_3 = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.any_op
-    transform.apply_patterns to %func_3 { transform.apply_patterns.canonicalization } : !transform.any_op
-    transform.apply_cse to %func_3 : !transform.any_op
-
+    // Phase 4: No L1 alloc -- let post-passes handle memory hierarchy.
     // Phase 5: Bufferize
     // Memory spaces for remaining allocs handled by post-transform overrides
     // in driver.py (scope=herd→2, scope=func→1 with exclusive scopes).