refactor: Update inference ops refactored upstream

build_ffi.py

@@ -56,70 +56,119 @@ def handle_atomic_declaration(line: str) -> str | None:
     return None
 
 
-def handle_deprecated_function(line: str, *, in_function: bool) -> (bool, bool):
-    """Handle removal of deprecated function declarations."""
-    deprecated_re = re.compile(r"__attribute__\(\(deprecated\b")
-
-    if in_function:
-        if line.strip().endswith("}"):
-            in_function = False
-        return (in_function, True)
-    if deprecated_re.search(line):
-        return (True, True)
-
-    return (in_function, False)
-
-
-def handle_inline_function(
-    line: str, brace_depth: int, *, in_function: bool
-) -> (bool, int, str):
-    """Handle removal of inline function declarations."""
-    inline_re = re.compile(r"\bstatic\s+inline\b")
-
-    if in_function:
-        brace_depth += line.count("{") - line.count("}")
-        if brace_depth <= 0:
-            in_function = False
-            brace_depth = 0
-        return in_function, brace_depth, None
-
-    match = inline_re.search(line)
-    if match:
-        prefix = line[: match.start()].rstrip()
-        if prefix:
-            # Still check for atomics in that prefix
-            atomic_result = handle_atomic_declaration(prefix)
-            result_line = atomic_result if atomic_result else prefix
-        brace_depth = line.count("{") - line.count("}")
-        if brace_depth > 0:
-            in_function = True
-        return in_function, brace_depth, result_line
-
-    return in_function, brace_depth, line
+def remove_deprecated_functions(src: str) -> str:
+    """Remove deprecated function declarations."""
+    out, i, n = [], 0, len(src)
+    pat = re.compile(r'__attribute__\s*\(\(deprecated\b')
+
+    while i < n:
+        m = pat.search(src, i)
+        if not m:
+            out.append(src[i:])
+            break
+
+        # Go backward to find the start of the function declaration
+        start = src.rfind('\n', 0, m.start())
+        start = start + 1 if start != -1 else 0
+        out.append(src[i:start])
+
+        # Move forward to skip the function body
+        j = src.find('{', m.end())
+        if j == -1:
+            i = m.end()
+            # Malformed, skip the attribute
+            continue
+
+        brace = 1
+        j += 1
+        while j < n and brace:
+            if src[j] == '{':
+                brace += 1
+            elif src[j] == '}':
+                brace -= 1
+            j += 1
+
+        # Skip trailing characters (semicolons, whitespace)
+        while j < n and src[j] in " \t\n;":
+            j += 1
+
+        i = j
+
+    return ''.join(out)
+
+
+def remove_static_inline_functions(source: str) -> str:
+    """Remove 'static inline' functions."""
+    static_inline_re = re.compile(r'\bstatic\s+inline\b')
+    result = []
+    i = 0
+    n = len(source)
+
+    while i < n:
+        match = static_inline_re.search(source, i)
+        if not match:
+            # No more static inlines, copy the rest
+            result.append(source[i:])
+            break
+
+        start = match.start()
+
+        # Copy anything before the match
+        if start > i:
+            result.append(source[i:start])
+
+        # Find the function body using brace matching
+        brace_count = 0
+        body_start = source.find('{', match.end())
+        if body_start == -1:
+            # Malformed inline without a body, treat as normal text
+            result.append(source[match.start():match.end()])
+            i = match.end()
+            continue
+
+        i = body_start
+        brace_count = 1
+        i += 1
+
+        while i < n and brace_count > 0:
+            if source[i] == '{':
+                brace_count += 1
+            elif source[i] == '}':
+                brace_count -= 1
+            i += 1
+
+        # If there is a semicolon mid-line, skip it
+        while i < n and source[i] in ' \t\n;':
+            i += 1
+
+    return ''.join(result)
 
 
 def sanitize_cdef(cdef: str, pkg: str) -> str:
     """Sanitize cdef by removing unsupported declarations."""
+    # Parse multi-line patterns first
+    cdef = remove_static_inline_functions(cdef)
+    cdef = remove_deprecated_functions(cdef)
+
     output_lines = cdef.splitlines()
     filtered_lines = []
 
     line_re = re.compile(r'# \d+ "(.*?)"')
 
     current_file = None
     pkg_lines = True
-    in_deprecated_function = False
-    in_inline_function = False
-    brace_depth = 0
 
+    # Parse single-line patterns next
     for line in output_lines:
+        # Parse and ignore GCC comments
         if line.startswith("#"):
             match = line_re.match(line)
             if match:
                 current_file = match.group(1)
                 pkg_lines = pkg in current_file
             continue
 
-        # Remove non-package declarations
+        # Ignore non-package declarations
         if not pkg_lines:
             continue
 
@@ -133,21 +182,6 @@ def sanitize_cdef(cdef: str, pkg: str) -> str:
             filtered_lines.append(atomic_result)
             continue
 
-        # Handle deprecated functions
-        (in_deprecated_function, skip) = handle_deprecated_function(
-            line, in_function=in_deprecated_function
-        )
-        if skip:
-            continue
-
-        # Handle inline functions
-        (in_inline_function, brace_depth, result_line) = handle_inline_function(
-            line, brace_depth, in_function=in_inline_function
-        )
-        if result_line is not None:
-            filtered_lines.append(result_line)
-            continue
-
         filtered_lines.append(line)
 
     return "\n".join(filtered_lines)

tests/test_tf.py

@@ -100,6 +100,8 @@ def test_tf(test_nodes, test_tensor, test_model):
     assert out_tensors[0].data == in_tensors[0].data
     assert out_tensors[0].to_bytes() == in_tensors[0].to_bytes()
 
+    session.tf_model_unload(model)
+
 
 def test_tf_with_run_options(test_nodes, test_tensor, test_model):
     session = Session()
@@ -127,6 +129,8 @@ def test_tf_with_run_options(test_nodes, test_tensor, test_model):
     assert out_tensors[0].data == in_tensors[0].data
     assert out_tensors[0].to_bytes() == in_tensors[0].to_bytes()
 
+    session.tf_model_unload(model)
+
 
 def test_tf_from_buffer(test_nodes, test_tensor, test_model):
     session = Session()
@@ -154,6 +158,8 @@ def test_tf_from_buffer(test_nodes, test_tensor, test_model):
     assert out_tensors[0].data == in_tensors[0].data
     assert out_tensors[0].to_bytes() == in_tensors[0].to_bytes()
 
+    session.tf_model_unload(model)
+
 
 def test_tf_from_numpy(test_nodes, test_tensor, test_model):
     session = Session()
@@ -176,3 +182,5 @@ def test_tf_from_numpy(test_nodes, test_tensor, test_model):
     assert out_tensors[0].data_type == in_tensors[0].data_type
     assert out_tensors[0].data == in_tensors[0].data
     assert out_tensors[0].to_bytes() == in_tensors[0].to_bytes()
+
+    session.tf_model_unload(model)

tests/test_tflite.py

@@ -81,6 +81,8 @@ def test_tflite(test_tensor, test_model):
     assert out_tensors[0].data == in_tensors[0].data
     assert out_tensors[0].to_bytes() == in_tensors[0].to_bytes()
 
+    session.tflite_model_unload(model)
+
 
 def test_tflite_from_buffer(test_tensor, test_model):
     session = Session()
@@ -103,6 +105,8 @@ def test_tflite_from_buffer(test_tensor, test_model):
     assert out_tensors[0].data == in_tensors[0].data
     assert out_tensors[0].to_bytes() == in_tensors[0].to_bytes()
 
+    session.tflite_model_unload(model)
+
 
 def test_tflite_from_numpy(test_tensor, test_model):
     session = Session()
@@ -120,3 +124,5 @@ def test_tflite_from_numpy(test_tensor, test_model):
     assert out_tensors[0].data_type == in_tensors[0].data_type
     assert out_tensors[0].data == in_tensors[0].data
     assert out_tensors[0].to_bytes() == in_tensors[0].to_bytes()
+
+    session.tflite_model_unload(model)

tests/test_torch.py

@@ -99,11 +99,13 @@ def test_torch(test_tensor, test_model):
     model = Resource(test_model, ResourceType.MODEL)
     model.register(session)
 
+    session.torch_model_load(model)
+
     in_tensors = [
         Tensor(test_tensor["dims"], test_tensor["type"], test_tensor["data"])
     ]
 
-    out_tensors = session.torch_jitload_forward(model, in_tensors)
+    out_tensors = session.torch_model_run(model, in_tensors)
     assert out_tensors[0].dims == in_tensors[0].dims
     assert out_tensors[0].data_type == in_tensors[0].data_type
     assert out_tensors[0].data == in_tensors[0].data
@@ -119,13 +121,15 @@ def test_torch_with_run_options(test_tensor, test_model):
     model = Resource(test_model, ResourceType.MODEL)
     model.register(session)
 
+    session.torch_model_load(model)
+
     in_tensors = [
         Tensor(test_tensor["dims"], test_tensor["type"], test_tensor["data"])
     ]
 
     run_options = Buffer(b"none")
 
-    out_tensors = session.torch_jitload_forward(
+    out_tensors = session.torch_model_run(
         model, in_tensors, run_options=run_options
     )
     assert out_tensors[0].dims == in_tensors[0].dims
@@ -143,13 +147,15 @@ def test_torch_from_buffer(test_tensor, test_model):
     model = Resource(test_model, ResourceType.MODEL)
     model.register(session)
 
+    session.torch_model_load(model)
+
     in_tensors = [
         Tensor.from_buffer(
             test_tensor["dims"], test_tensor["type"], test_tensor["data_bytes"]
         )
     ]
 
-    out_tensors = session.torch_jitload_forward(model, in_tensors)
+    out_tensors = session.torch_model_run(model, in_tensors)
     assert out_tensors[0].dims == in_tensors[0].dims
     assert out_tensors[0].data_type == in_tensors[0].data_type
     assert out_tensors[0].data == in_tensors[0].data
@@ -165,9 +171,11 @@ def test_torch_from_numpy(test_tensor, test_model):
     model = Resource(test_model, ResourceType.MODEL)
     model.register(session)
 
+    session.torch_model_load(model)
+
     in_tensors = [Tensor.from_numpy(test_tensor["data_np"])]
 
-    out_tensors = session.torch_jitload_forward(model, in_tensors)
+    out_tensors = session.torch_model_run(model, in_tensors)
     assert out_tensors[0].dims == in_tensors[0].dims
     assert out_tensors[0].data_type == in_tensors[0].data_type
     assert out_tensors[0].data == in_tensors[0].data
@@ -184,9 +192,11 @@ def test_torch_from_torch(test_tensor, test_model):
     model = Resource(test_model, ResourceType.MODEL)
     model.register(session)
 
+    session.torch_model_load(model)
+
     in_tensors = [Tensor.from_torch(test_tensor["data_torch"])]
 
-    out_tensors = session.torch_jitload_forward(model, in_tensors)
+    out_tensors = session.torch_model_run(model, in_tensors)
     assert out_tensors[0].dims == in_tensors[0].dims
     assert out_tensors[0].data_type == in_tensors[0].data_type
     assert out_tensors[0].data == in_tensors[0].data