hexagon: support for multi-device model split (aka row-split) (#28589)

* hex-row-split: add support for multi-device row spliting

Co-authored-by: Max Krasnyansky <maxk@qti.qualcomm.com>

* hex-mdev: add work splitting to fused kernels

* hex-mdev: use mdev_ prefix for all multi-device state

* hex-mdev: make device configuration more expressive to support device groups

* hex-mdev: fix mdev session init

* hex-mdev: fused nx (2x,3x) matmuls must update row counts for each w/o

* hex-mdev: fix MUL_MAT work partitioning bugs introduced by mdev

* hex-cont: fix crashes with new tests due to wrong striding

* hex-mdev: move fences after l2flushes

* hex-cont: fix work splitting for mnpu -- align chunks to cachelines

* hex-mdev: fix CPY tests with multi-dev

* hex-mmid: fix work partitioning with mnpu

* hex-mm: fix test failures with mdev

* hex-binary: fix work partitioning for mdev

* hex-argsort: fix mdev partitioning

* hex-mdev: fix work partitioning and general updates for all simple ops

* hex-fa: fix mdev work splitting issues

* hex-mdev: fixing more failing ops test

* hex-mdev: update the rest of the ops

* hex-mdev: refactor all mdev splitting logic to be contained within if (mdev_count > 1) {...}

* hex-mdev: fix macros

* hex-mdev: simplify session flush logic

* hex-sync: fix recursion in session flush

* hex-mdev: factor out fence buffer and allocator

* hex-fence: make fence allocation more robust with reserved slots for mdev

* hex-mdev: keep all mdev state in htp_mdev_group

* hex-mdev: further cleanup mdev group handling at the host

* hex-mdev: update group idx in the opbatch before serializing

* hex-batch: remove separate op_pending and use batch_req/rsp_seq

* hex-async: workaround another missing tensor_init in ggml-meta

* hex-fence: cleanup and robustify fences and error handling in multi-device scenarios

* hex-ar: improve ALLREDUCE error handling

* hex-async: robust error handling for op_cpy_fence

* hex-async: use seq0 from allreduce context to allocate fence_seq

* hex-mdev: fix remaining issues with fence and barrier clearing in CPY_FENCE

* hex-misc: realign macros and fix misplaces trace events

* hex-misc: align macros

* hex-mdev: fix unclone buffer re-entrancy

* hex-glu: fix mdev partitioning logic

* hex-mdev: make buffer uncloning/cleanup work with tensor-split scenarios

* hex-mdev: tighten up the can_split check in act-ops

* hex-mdev: factor out common bits of the partitioning logic

* hex-mm: minor realignment of the macros

* hex-bufs: fix incorrectly placed assert for MAX_BUFS

* hex-pad: tighten up gating checks for PAD

* hex-kparams: make sure all kernels properly use kparams->n_threads

* hex-docs: update user and developer docs with new features and detailed guide for ops development

* hex-scripts: update run script to properly parse dev groups

* hex-misc: formatting

* hex-sess: minor cleanup for session init

* hex-ar: fix vtcm size calc in allreduce kparams

* hex-scripts: fix flake8 warnings

* hex-rope: update ROPE to support mdev work split

* hex-ops: remove redunant checks and minor reformat

* hex-dev-guide: update dev-guide to avoid redundant null checks

* hex-async: improve event_wait, event_sync and fence implementations

* hex-async: remove synchronous flush from event_sync

* hex-async: symplify fence recovery protocol and make sync more robust

* hex-async: futher simplify error recovery for fences

* hex-err: return status instead of just -1

* hex-async: print all seq nums in hex

* hex-async: make sure fences flush dirty ranges

* hex-async: add dirty ranges merging to reduce fence flushes

* hex-async: properly sync before freeing the event

* hex-async: make sure fence owner session is not overriden

* hex-async: more fence write order more robust

* hex-async: make sure not to fuse ALLREDUCE+ADD if their dsts overlap

* hex-fusion: cleanup redundant checks

---------

Co-authored-by: Alexander Lu <alexlu@qti.qualcomm.com>

scripts/snapdragon/ggml-hexagon-align-macros.py

@@ -0,0 +1,296 @@
+#!/usr/bin/env python3
+"""
+align-macros.py - Inspect and align trailing backslashes in multiline C/C++ macros.
+
+Usage:
+    align-macros.py [paths...]                 # Check and report misaligned macros
+    align-macros.py --diff [paths...]          # Show unified diff of fixes
+    align-macros.py --fix [paths...]           # Fix misaligned macros in-place
+    align-macros.py --fix --mode majority ...  # Align to the dominant column
+    align-macros.py --fix --pad 2 ...          # Align to (max_content_len + pad)
+
+Safety rules:
+    - Macros that are ALREADY aligned are NEVER touched (unless --all is given).
+    - Whitespace after trailing backslashes is flagged and cleaned.
+"""
+
+import argparse
+import difflib
+import logging
+import os
+import re
+import sys
+from collections import Counter
+from typing import List, Optional, Tuple, NamedTuple
+
+logger = logging.getLogger("ggml-hexagon-align-macros")
+
+
+class MacroLine(NamedTuple):
+    line_num: int       # 1-indexed
+    raw: str            # Original line including newline
+    content: str        # Line content before trailing backslash (stripped of trailing whitespace)
+    bs_col: Optional[int]  # 1-indexed column of backslash, or None if last line has no backslash
+    trailing_ws: bool   # True if whitespace existed after the backslash
+
+
+class MacroDef(NamedTuple):
+    name: str
+    filepath: str
+    start_line: int
+    end_line: int
+    lines: List[MacroLine]
+
+
+def parse_macros(filepath: str) -> List[MacroDef]:
+    """Extract all multiline macros from a C/C++ source file."""
+    try:
+        with open(filepath, "r", encoding="utf-8", errors="replace") as f:
+            lines = f.readlines()
+    except Exception as e:
+        logger.error(f"Error reading {filepath}: {e}")
+        return []
+
+    macros: List[MacroDef] = []
+    i = 0
+    n = len(lines)
+
+    while i < n:
+        line = lines[i]
+        m = re.match(r"^\s*#\s*define\s+([A-Za-z_][A-Za-z0-9_]*)", line)
+        if m:
+            macro_name = m.group(1)
+            macro_start = i + 1
+            macro_lines: List[MacroLine] = []
+            cur = i
+
+            while cur < n:
+                l_raw = lines[cur]
+                l_rstrip = l_raw.rstrip("\r\n")
+
+                # Check if line has a trailing backslash
+                # Note: handle possible accidental spaces after backslash
+                match_bs = re.search(r"\\([ \t]*)$", l_rstrip)
+                if match_bs:
+                    has_trailing_ws = len(match_bs.group(1)) > 0
+                    bs_index = match_bs.start()
+                    content = l_rstrip[:bs_index].rstrip()
+                    # 1-indexed column of the backslash
+                    bs_col = bs_index + 1
+                    macro_lines.append(MacroLine(
+                        line_num=cur + 1,
+                        raw=l_raw,
+                        content=content,
+                        bs_col=bs_col,
+                        trailing_ws=has_trailing_ws
+                    ))
+                    cur += 1
+                else:
+                    # Line does not end with backslash
+                    if cur == i:
+                        # Single-line macro, not multiline
+                        break
+                    else:
+                        # Final line of a multiline macro
+                        macro_lines.append(MacroLine(
+                            line_num=cur + 1,
+                            raw=l_raw,
+                            content=l_rstrip.rstrip(),
+                            bs_col=None,
+                            trailing_ws=False
+                        ))
+                        break
+
+            # Only record if it is a multiline macro (has at least one continuation line)
+            continuation_lines = [ml for ml in macro_lines if ml.bs_col is not None]
+            if continuation_lines:
+                macro_end = macro_lines[-1].line_num
+                macros.append(MacroDef(
+                    name=macro_name,
+                    filepath=filepath,
+                    start_line=macro_start,
+                    end_line=macro_end,
+                    lines=macro_lines
+                ))
+            i = cur
+        i += 1
+
+    return macros
+
+
+def is_macro_aligned(macro: MacroDef) -> bool:
+    """A macro is aligned if all continuation lines have backslashes at the same column."""
+    bs_cols = [ml.bs_col for ml in macro.lines if ml.bs_col is not None]
+    if not bs_cols:
+        return True
+    has_trailing_ws = any(ml.trailing_ws for ml in macro.lines)
+    return len(set(bs_cols)) == 1 and not has_trailing_ws
+
+
+def compute_target_column(macro: MacroDef, mode: str, pad: int, target_col: Optional[int]) -> int:
+    """Determine the column where backslashes should be aligned."""
+    max_content_len = max(len(ml.content) for ml in macro.lines)
+    min_needed = max_content_len + pad
+
+    if target_col is not None:
+        return max(target_col, min_needed)
+
+    bs_cols = [ml.bs_col for ml in macro.lines if ml.bs_col is not None]
+    if not bs_cols:
+        return min_needed
+
+    if mode == "min":
+        return min_needed
+    elif mode == "max":
+        return max(max(bs_cols), min_needed)
+    elif mode == "majority":
+        counts = Counter(bs_cols)
+        # Sort by frequency descending, then by column descending
+        majority_col = sorted(counts.items(), key=lambda x: (-x[1], -x[0]))[0][0]
+        return max(majority_col, min_needed)
+    else:
+        return min_needed
+
+
+def realign_macro_lines(macro: MacroDef, target_col: int) -> List[str]:
+    """Format macro lines with backslashes aligned at target_col."""
+    new_lines: List[str] = []
+    for ml in macro.lines:
+        nl = "\r\n" if ml.raw.endswith("\r\n") else "\n"
+        if ml.bs_col is None:
+            # Last line without backslash
+            new_lines.append(ml.raw)
+        else:
+            if not ml.content:
+                spaces = " " * (target_col - 1)
+                new_lines.append(f"{spaces}\\{nl}")
+            else:
+                spaces_needed = max(1, target_col - len(ml.content) - 1)
+                new_lines.append(f"{ml.content}{' ' * spaces_needed}\\{nl}")
+    return new_lines
+
+
+def process_file(filepath: str, args: argparse.Namespace) -> Tuple[int, int, Optional[str]]:
+    macros = parse_macros(filepath)
+    if not macros:
+        return 0, 0, None
+
+    with open(filepath, "r", encoding="utf-8", errors="replace") as f:
+        file_lines = f.readlines()
+
+    misaligned_count = 0
+    modified = False
+    new_file_lines = list(file_lines)
+
+    for macro in macros:
+        aligned = is_macro_aligned(macro)
+        if not aligned or args.all:
+            if not aligned:
+                misaligned_count += 1
+
+            bs_cols = [ml.bs_col for ml in macro.lines if ml.bs_col is not None]
+            max_content = max(len(ml.content) for ml in macro.lines)
+            col_counts = Counter(bs_cols)
+
+            if not args.quiet:
+                logger.info(f"{filepath}:{macro.start_line}-{macro.end_line} [{macro.name}]")
+                logger.info(f"  Max content width: {max_content}, Min needed column (+{args.pad}): {max_content + args.pad}")
+                logger.info(f"  Current backslash columns: {dict(sorted(col_counts.items()))}")
+                trailing_ws_lines = [ml.line_num for ml in macro.lines if ml.trailing_ws]
+                if trailing_ws_lines:
+                    logger.warning(f"  Warning: Trailing whitespace after backslash on line(s): {trailing_ws_lines}")
+
+            target_col = compute_target_column(macro, args.mode, args.pad, args.target_col)
+            if not args.quiet:
+                logger.info(f"  -> Target alignment column: {target_col}")
+
+            realigned = realign_macro_lines(macro, target_col)
+
+            start_idx = macro.start_line - 1
+            end_idx = start_idx + len(macro.lines)
+            if new_file_lines[start_idx:end_idx] != realigned:
+                new_file_lines[start_idx:end_idx] = realigned
+                modified = True
+
+    diff_text = None
+    if modified:
+        diff = difflib.unified_diff(
+            file_lines,
+            new_file_lines,
+            fromfile=f"a/{filepath}",
+            tofile=f"b/{filepath}",
+            lineterm=""
+        )
+        diff_text = "\n".join(diff)
+
+        if args.fix:
+            with open(filepath, "w", encoding="utf-8") as f:
+                f.writelines(new_file_lines)
+            if not args.quiet:
+                logger.info(f"  [FIXED] Updated {filepath}")
+
+    return len(macros), misaligned_count, diff_text
+
+
+def find_source_files(paths: List[str]) -> List[str]:
+    extensions = {".c", ".cpp", ".cc", ".cxx", ".h", ".hpp", ".inl"}
+    result: List[str] = []
+    for p in paths:
+        if os.path.isfile(p):
+            result.append(p)
+        elif os.path.isdir(p):
+            for root, _, files in os.walk(p):
+                for file in sorted(files):
+                    _, ext = os.path.splitext(file)
+                    if ext.lower() in extensions:
+                        result.append(os.path.join(root, file))
+    return sorted(result)
+
+
+def main():
+    logging.basicConfig(level=logging.INFO, format="%(message)s")
+    parser = argparse.ArgumentParser(
+        description="Inspect and align backslashes in multiline C/C++ macros."
+    )
+    parser.add_argument("paths", nargs="*", default=["."], help="Files or directories to scan (default: current dir)")
+    parser.add_argument("--fix", action="store_true", help="Fix misaligned macros in-place")
+    parser.add_argument("--diff", action="store_true", help="Display unified diff of suggested fixes")
+    parser.add_argument("--check", action="store_true", help="Exit with code 1 if misaligned macros exist")
+    parser.add_argument("--mode", choices=["min", "max", "majority"], default="min",
+                        help="Alignment mode: 'min' (max_len + pad), 'max' (max existing col), 'majority' (dominant col)")
+    parser.add_argument("--pad", type=int, default=2, help="Spaces between longest line and backslash (default: 2)")
+    parser.add_argument("--target-col", type=int, default=None, help="Force alignment to an exact column")
+    parser.add_argument("--all", action="store_true", help="Realign all macros even if already aligned (default: only misaligned)")
+    parser.add_argument("-q", "--quiet", action="store_true", help="Only output errors and diffs/summary")
+
+    args = parser.parse_args()
+
+    files = find_source_files(args.paths)
+    if not files:
+        logger.error("No C/C++ source files found.")
+        sys.exit(0)
+
+    total_macros = 0
+    total_misaligned = 0
+    diffs: List[str] = []
+
+    for filepath in files:
+        num_macros, num_misaligned, diff_text = process_file(filepath, args)
+        total_macros += num_macros
+        total_misaligned += num_misaligned
+        if diff_text:
+            diffs.append(diff_text)
+
+    if args.diff and diffs:
+        logger.info("\n--- Proposed Changes ---\n")
+        for d in diffs:
+            logger.info(d)
+
+    logger.info(f"\nSummary: scanned {len(files)} files, {total_macros} multiline macros, {total_misaligned} misaligned.")
+
+    if args.check and total_misaligned > 0:
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

scripts/snapdragon/run.py

@@ -14,6 +14,42 @@ import logging
 logger = logging.getLogger("run")
 
 
+MANAGED_ENV_NAMES = (
+    "GGML_HEXAGON_DEVICES",
+    "GGML_HEXAGON_VERBOSE",
+    "GGML_HEXAGON_PROFILE",
+    "GGML_HEXAGON_NHVX",
+    "GGML_HEXAGON_NHMX",
+    "GGML_HEXAGON_HOSTBUF",
+    "GGML_HEXAGON_OPBATCH",
+    "GGML_HEXAGON_OPQUEUE",
+    "GGML_HEXAGON_OPPOLL",
+    "GGML_HEXAGON_OPFILTER",
+    "GGML_HEXAGON_OPFUSION",
+    "GGML_HEXAGON_VMEM",
+    "GGML_HEXAGON_MBUF",
+    "GGML_HEXAGON_MM_SELECT",
+    "GGML_HEXAGON_FA_SELECT",
+    "GGML_HEXAGON_AR_SELECT",
+    "GGML_HEXAGON_ETM",
+    "GGML_HEXAGON_ARCH",
+    "GGML_HEXAGON_OPTRACE",
+    "GGML_OPENCL_PLATFORM",
+    "GGML_OPENCL_DEVICE",
+    "GGML_OPENCL_OPFILTER",
+    "GGML_OPENCL_KERNEL_CACHE_DIR",
+    "GGML_OPENCL_KERNEL_CACHE_DEBUG",
+    "GGML_OPENCL_FA_TUNE",
+    "GGML_OPENCL_DISABLE_FUSION",
+    "GGML_OPENCL_ADRENO_XMEM_GEMM",
+    "GGML_OPENCL_ADRENO_USE_LARGE_BUFFER",
+    "GGML_SCHED_DEBUG",
+    "MTMD_BACKEND_DEVICE",
+    "D",
+    "DEVICE",
+)
+
+
 def parse_target(target_str):
     if not target_str:
         return None, None
@@ -38,6 +74,57 @@ def shlex_join(args_list):
     return " ".join(pipes.quote(x) for x in args_list)
 
 
+def split_device_list(devices):
+    parts = []
+    curr = []
+    bracket_depth = 0
+
+    for ch in devices:
+        if ch == '[':
+            bracket_depth += 1
+            curr.append(ch)
+        elif ch == ']':
+            if bracket_depth > 0:
+                bracket_depth -= 1
+            curr.append(ch)
+        elif ch == ',' and bracket_depth == 0:
+            part = "".join(curr).strip()
+            if part:
+                parts.append(part)
+            curr = []
+        else:
+            curr.append(ch)
+
+    part = "".join(curr).strip()
+    if part:
+        parts.append(part)
+
+    return parts
+
+
+def device_arg_from_devices(devices):
+    if devices.isdigit():
+        n = int(devices)
+        return ",".join(f"HTP{i}" for i in range(n))
+
+    names = []
+    for part in split_device_list(devices):
+        if "[" in part:
+            part = part.split("[", 1)[0].strip()
+        if part:
+            names.append(part)
+
+    return ",".join(names)
+
+
+def normalize_cmd_device_args(cmd_args):
+    for i, arg in enumerate(cmd_args):
+        if arg == "--device" and i + 1 < len(cmd_args):
+            cmd_args[i + 1] = device_arg_from_devices(cmd_args[i + 1])
+        elif arg.startswith("--device="):
+            cmd_args[i] = "--device=" + device_arg_from_devices(arg.split("=", 1)[1])
+
+
 def main():
     logging.basicConfig(level=logging.INFO, format='%(message)s')
     # Split arguments at '--'
@@ -142,8 +229,6 @@ def main():
     def set_env(env_name, opt_val):
         if opt_val is not None:
             env_vars[env_name] = str(opt_val)
-        elif env_name in os.environ:
-            env_vars[env_name] = os.environ[env_name]
 
     # Resolve and filter devices (HTP vs OpenCL)
     device_in_cmd = None
@@ -166,7 +251,7 @@ def main():
         hex_devices = devices_val
         cl_device = ""
     else:
-        parts = [p.strip() for p in devices_val.split(",")]
+        parts = split_device_list(devices_val)
         # Any device containing "htp" is Hexagon, rest is OpenCL
         hex_parts = [p for p in parts if "htp" in p.lower()]
         cl_parts = [
@@ -181,15 +266,13 @@ def main():
     # Set Hexagon devices
     if hex_devices:
         env_vars["GGML_HEXAGON_DEVICES"] = hex_devices
-    elif "GGML_HEXAGON_DEVICES" in os.environ:
-        env_vars["GGML_HEXAGON_DEVICES"] = os.environ["GGML_HEXAGON_DEVICES"]
+
+    normalize_cmd_device_args(cmd_args)
 
     # Set OpenCL device (unless overridden by --cl-device)
     final_cl_device = args.cl_device if args.cl_device is not None else cl_device
     if final_cl_device:
         env_vars["GGML_OPENCL_DEVICE"] = final_cl_device
-    elif "GGML_OPENCL_DEVICE" in os.environ:
-        env_vars["GGML_OPENCL_DEVICE"] = os.environ["GGML_OPENCL_DEVICE"]
 
     # Map shared & backend-specific parameters with correct overrides
 
@@ -206,8 +289,6 @@ def main():
 
     if args.cl_fa_tune or args.profile is not None:
         env_vars["GGML_OPENCL_FA_TUNE"] = "1"
-    elif "GGML_OPENCL_FA_TUNE" in os.environ:
-        env_vars["GGML_OPENCL_FA_TUNE"] = os.environ["GGML_OPENCL_FA_TUNE"]
 
     # Other Hexagon environment variables
     set_env("GGML_HEXAGON_NHVX", args.hex_nhvx)
@@ -235,18 +316,12 @@ def main():
 
     if args.cl_disable_fusion:
         env_vars["GGML_OPENCL_DISABLE_FUSION"] = "1"
-    elif "GGML_OPENCL_DISABLE_FUSION" in os.environ:
-        env_vars["GGML_OPENCL_DISABLE_FUSION"] = os.environ["GGML_OPENCL_DISABLE_FUSION"]
 
     if args.cl_adreno_xmem:
         env_vars["GGML_OPENCL_ADRENO_XMEM_GEMM"] = "1"
-    elif "GGML_OPENCL_ADRENO_XMEM_GEMM" in os.environ:
-        env_vars["GGML_OPENCL_ADRENO_XMEM_GEMM"] = os.environ["GGML_OPENCL_ADRENO_XMEM_GEMM"]
 
     if args.cl_adreno_large_buffer:
         env_vars["GGML_OPENCL_ADRENO_USE_LARGE_BUFFER"] = "1"
-    elif "GGML_OPENCL_ADRENO_USE_LARGE_BUFFER" in os.environ:
-        env_vars["GGML_OPENCL_ADRENO_USE_LARGE_BUFFER"] = os.environ["GGML_OPENCL_ADRENO_USE_LARGE_BUFFER"]
 
     if args.sched_debug:
         env_vars["GGML_SCHED_DEBUG"] = "2"
@@ -288,15 +363,7 @@ def main():
         has_b = any(arg == "-b" for arg in cmd_args)
         if not has_b:
             if args.devices:
-                if args.devices.isdigit():
-                    n = int(args.devices)
-                    device_val = ",".join(f"HTP{i}" for i in range(n))
-                else:
-                    device_val = args.devices
-            elif "D" in os.environ:
-                device_val = os.environ["D"]
-            elif "DEVICE" in os.environ:
-                device_val = os.environ["DEVICE"]
+                device_val = device_arg_from_devices(args.devices)
             else:
                 device_val = "HTP0"
             if device_val:
@@ -305,17 +372,10 @@ def main():
         has_device = any(arg.startswith("--device") for arg in cmd_args)
         if not has_device:
             if args.devices:
-                if args.devices.isdigit():
-                    n = int(args.devices)
-                    device_val = ",".join(f"HTP{i}" for i in range(n))
-                else:
-                    device_val = args.devices
-            elif "D" in os.environ:
-                device_val = os.environ["D"]
-            elif "DEVICE" in os.environ:
-                device_val = os.environ["DEVICE"]
+                device_val = device_arg_from_devices(args.devices)
             else:
                 device_val = "HTP0"
+
             if device_val:
                 cmd_args += ["--device", device_val]
 
@@ -415,6 +475,8 @@ def main():
         else:
             local_env["LD_LIBRARY_PATH"] = lib_dir + os.path.pathsep + local_env.get("LD_LIBRARY_PATH", "")
 
+        for k in MANAGED_ENV_NAMES:
+            local_env.pop(k, None)
         for k, v in env_vars.items():
             local_env[k] = v