From c2b98df63af87edd5537721515a8b02f6871a307 Mon Sep 17 00:00:00 2001
From: richard-byron <richard.byron@amd.com>
Date: Wed, 2 Sep 2026 12:03:50 -0700
Subject: [PATCH] Add unsupported_layer options to skip specific layers

---
 src/mldebug/layer_info.py  | 29 ++++++++++++++++-------------
 src/mldebug/mldebug_cli.py | 34 +++++++++++++++++++++++++++++++++-
 2 files changed, 49 insertions(+), 14 deletions(-)

diff --git a/src/mldebug/layer_info.py b/src/mldebug/layer_info.py
index f94e9a3..d19cac9 100644
--- a/src/mldebug/layer_info.py
+++ b/src/mldebug/layer_info.py
@@ -576,7 +576,7 @@ def __init__(self, args):
       )
     # 4. Initialize Layers
     if not args.aie_only:
-      self._init_layers(data, args.aie_iface, num_stamps, num_batches)
+      self._init_layers(data, args.aie_iface, num_stamps, num_batches, args.unsupported_layers)
     # 5: Parse work dir
     if self.x2:
       for layer in self.layers:
@@ -898,7 +898,7 @@ def _read_buffer_info(self, buffer_info_file):
     self.x2 = data[".meta"].get("flow") == "x2"
     return data
 
-  def _init_layers(self, raw_info, aie_iface, num_stamps, num_batches=1):
+  def _init_layers(self, raw_info, aie_iface, num_stamps, num_batches = 1, skip_layers=None):
     """
     Parse all layer entries from metadata and populate self.layers.
 
@@ -907,6 +907,7 @@ def _init_layers(self, raw_info, aie_iface, num_stamps, num_batches=1):
         aie_iface: AIE interface object.
         num_stamps (int): Stamps per batch (S from BxSxCxR).
         num_batches (int): Number of batches (B from BxSxCxR).
+        skip_layers (set[int]): layer_order values to skip (from --unsupported_layer).
     """
     version = Version.from_string(raw_info[".meta"]["version"])
     size_shift = raw_info[".meta"].get("size_shift")
@@ -924,18 +925,20 @@ def _init_layers(self, raw_info, aie_iface, num_stamps, num_batches=1):
     for entry in raw_layers:
       info = entry[1]
       self._warn_if_scheduled_in_chunks(info)
-      self.layers.append(
-        Layer(
-          info,
-          size_shift,
-          version,
-          aie_iface,
-          num_stamps,
-          self.mladf_report,
-          num_batches=num_batches,
-          device_batch_size=self.layout[0],
-        )
+      layer = Layer(
+        info,
+        size_shift,
+        version,
+        aie_iface,
+        num_stamps,
+        self.mladf_report,
+        num_batches=num_batches,
+        device_batch_size=self.layout[0],
       )
+      if skip_layers and layer.layer_order in skip_layers:
+        LOGGER.verbose_print(f"[WARNING] unsupported layer {layer.layer_order} will be skipped.")
+        layer.is_unsupported = True
+      self.layers.append(layer)
     self._reorder_layers_by_execution()
 
   def _warn_if_scheduled_in_chunks(self, info):
diff --git a/src/mldebug/mldebug_cli.py b/src/mldebug/mldebug_cli.py
index eae3ed5..8ffaeea 100644
--- a/src/mldebug/mldebug_cli.py
+++ b/src/mldebug/mldebug_cli.py
@@ -58,7 +58,28 @@ def _apply_unsupported_kernels_from_args(args):
       if token:
         layer_info.unsupported_superkernels.append(token.lower())
 
-
+def _apply_unsupported_layers_from_args(args):
+  """
+  Parse --unsupported_layer into args.unsupported_layers (set of layer_order ints).
+  Supports both:
+    --unsupported_layer 285 286
+  and:
+    --unsupported_layer 285,286
+  """
+  values = args.unsupported_layers
+  if not values:
+    args.unsupported_layers = set()
+    return
+  skip_layers = set()
+  for v in values:
+    if v is None:
+      continue
+    for token in str(v).split(","):
+      token = token.strip()
+      if token:
+        skip_layers.add(int(token))
+  args.unsupported_layers = skip_layers
+  
 def check_args(args):
   """
   Check argument rules
@@ -170,6 +191,7 @@ def launch_debug(args, output_dir):
     context_id, pid = check_hw_context(args)
   # Top debug handle
   _apply_unsupported_kernels_from_args(args)
+  _apply_unsupported_layers_from_args(args)
   handle = ClientDebug(args, context_id, pid, output_dir)
   if args.dump_layers:
     handle.dump_layers(None if args.dump_layers == "-" else args.dump_layers)
@@ -380,6 +402,16 @@ def app():
     # help="Additional kernel names to treat as unsupported and skip during execution.\n"
     # "Example: --unsupported_kernels conv2d_maxpool superkernel_clip1d\n",
   )
+  p.add_argument(
+    "--unsupported_layer",
+    dest="unsupported_layers",
+    nargs="*",
+    default=None,
+    metavar="LAYER",
+    help=argparse.SUPPRESS,
+    #help="Layer numbers to skip during execution.\n"
+    #"Example: --unsupported_layer 285 286\n",
+  )
   p.add_argument(
     "-f",
     "--run_flags",