Integrate native determinism with canonical contacts and constraint allocation

.pre-commit-config.yaml

@@ -19,5 +19,6 @@ repos:
       - id: kernel-analyzer
         name: kernel-analyzer
         entry: uv run python contrib/kernel_analyzer/kernel_analyzer/cli.py
+        args: [--check_atomic=True]
         language: system
         types: [python]

contrib/kernel_analyzer/.vscodeignore

@@ -1,3 +1,3 @@
 # vsce package will skip packaging these file patterns
-__pycache__
-client/node_modules
+**/__pycache__/**
+client/node_modules/**

contrib/kernel_analyzer/kernel_analyzer/ast_analyzer.py

@@ -155,6 +155,15 @@ def __str__(self):
     return f'"wp.{self.node.func.attr}" not allowed inside wp.static(); use pure Python or math.*'
 
 
+@dataclasses.dataclass
+class MissingDeterministicFactory(Issue):
+  """Report an atomic operation outside a configurable deterministic factory."""
+
+  def __str__(self):
+    """Describe the missing configurable factory for this atomic operation."""
+    return f'"{self.kernel}" contains atomic operations but is not enclosed in a kernel factory with a "deterministic" argument'
+
+
 # TODO(team): add argument order analyzer.
 # this one is tricky because just verifying order does not tell you if the arguments
 # match the parameter signature.
@@ -266,7 +275,12 @@ def _is_wp_call(node: ast.AST) -> bool:
   )
 
 
-def analyze(source: str, filename: str, type_source: str) -> List[Issue]:
+def analyze(
+  source: str,
+  filename: str,
+  type_source: str,
+  check_atomic: bool = False,
+) -> List[Issue]:
   """Parses Python code and finds functions with unsorted simple parameters."""
   logging.info(f"Analyzing {filename}...")
 
@@ -316,6 +330,7 @@ def analyze(source: str, filename: str, type_source: str) -> List[Issue]:
         if attr_type.startswith("wp.array"):
           dummy_node = ast.AST()
           dummy_node.lineno = 1
+          dummy_node.end_lineno = 1
           issues.append(DirectWpArrayForbidden(dummy_node, "", attr_name, cls_name))
   source_lines = source.splitlines()
 
@@ -329,12 +344,12 @@ def _is_func(name: str) -> bool:
 
   top_kernels: dict[str, bool] = {}
 
-  def _analyze_function(node: ast.FunctionDef, is_nested: bool):
+  def _analyze_function(node: ast.FunctionDef, is_nested: bool, enclosing_func: Optional[ast.FunctionDef] = None):
     """Analyze a function definition for kernel issues."""
     # Recursively check nested functions first
     for child in ast.iter_child_nodes(node):
       if isinstance(child, ast.FunctionDef):
-        _analyze_function(child, is_nested=True)
+        _analyze_function(child, is_nested=True, enclosing_func=node)
 
     # Find wp.kernel or wp.func decorator
     decorator = None
@@ -354,6 +369,23 @@ def _analyze_function(node: ast.FunctionDef, is_nested: bool):
       else:
         top_kernels[node.name] = has_unique
 
+    # Kernels and functions with atomics must be inside a deterministic factory.
+    if check_atomic and (_is_kernel(decorator) or _is_func(decorator)):
+      has_det = False
+      if is_nested and enclosing_func is not None:
+        has_det = any(a.arg == "deterministic" for a in enclosing_func.args.args + enclosing_func.args.kwonlyargs)
+      if not has_det:
+        for sub in ast.walk(node):
+          if (
+            isinstance(sub, ast.Call)
+            and isinstance(sub.func, ast.Attribute)
+            and isinstance(sub.func.value, ast.Name)
+            and sub.func.value.id == "wp"
+            and sub.func.attr.startswith("atomic_")
+            and sub.func.attr not in ("atomic_or", "atomic_and", "atomic_xor")
+          ):
+            issues.append(MissingDeterministicFactory(sub, node.name))
+
     _analyze_kernel(node)
 
   def _analyze_kernel(node: ast.FunctionDef):
@@ -628,17 +660,49 @@ def _is_modulo_shape0(index_node: ast.AST, param_name: str) -> bool:
               issues.append(SharedModuleCustomBlockDim(sub_node, target_name))
 
   # skip issues in ignored lines
-  ignore_lines = set()
-  ignoring = False
+  # Supports # kernel_analyzer: ignore, ignore[rule], off/on, off[rule]/on[rule]
+  ignore_lines: Dict[int, Set[str]] = {}
+  active_off_rules: Set[str] = set()
+
   for lineno, line in enumerate(source_lines, 1):  # lineno is 1-indexed
-    if "# kernel_analyzer: off" in line:
-      ignoring = True
-    elif "# kernel_analyzer: on" in line:
-      ignoring = False
-    if "# kernel_analyzer: ignore" in line or ignoring:
-      ignore_lines.add(lineno)
-
-  filtered_issues = [i for i in issues if i.node.lineno not in ignore_lines]
+    m_off = re.search(r"#\s*kernel_analyzer:\s*off(?:\[(.*?)\])?", line)
+    if m_off:
+      active_off_rules.add(m_off.group(1).strip() if m_off.group(1) else "*")
+    m_on = re.search(r"#\s*kernel_analyzer:\s*on(?:\[(.*?)\])?", line)
+    if m_on:
+      rule = m_on.group(1).strip() if m_on.group(1) else "*"
+      active_off_rules.discard(rule)
+      if rule == "*":
+        active_off_rules.clear()
+
+    m_ign = re.search(r"#\s*kernel_analyzer:\s*ignore(?:\[(.*?)\])?", line)
+    rules_to_ignore = set(active_off_rules)
+    if m_ign:
+      rules_to_ignore.add(m_ign.group(1).strip() if m_ign.group(1) else "*")
+
+    if rules_to_ignore:
+      ignore_lines[lineno] = rules_to_ignore
+
+  def _is_ignored(iss: Issue) -> bool:
+    """Check suppression rules without extending a declaration into its body."""
+    start_line = iss.node.lineno
+    # Function-level diagnostics belong to the declaration, not its entire body.
+    end_line = (
+      start_line if isinstance(iss.node, (ast.FunctionDef, ast.AsyncFunctionDef)) else (iss.node.end_lineno or start_line)
+    )
+    matched_rules: Set[str] = set()
+    for l in range(start_line, end_line + 1):
+      if l in ignore_lines:
+        matched_rules.update(ignore_lines[l])
+    if not matched_rules:
+      return False
+    if "*" in matched_rules:
+      return True
+    if isinstance(iss, MissingDeterministicFactory):
+      return any("atomic" in r.lower() for r in matched_rules)
+    return False
+
+  filtered_issues = [i for i in issues if not _is_ignored(i)]
   filter_count, ignore_count = len(filtered_issues), len(issues) - len(filtered_issues)
 
   logging.info(f"Finished analyzing {filename}. Found {filter_count} issues, ignoring {ignore_count} issues.")

contrib/kernel_analyzer/kernel_analyzer/ast_analyzer_test.py

@@ -598,5 +598,181 @@ def launch_it():
     self.assertEqual(len(shared_issues), 0, shared_issues)
 
 
+class TestDeterministicFactory(absltest.TestCase):
+  """Check deterministic factory detection and local suppression rules."""
+
+  def test_inline_ignore_does_not_suppress_function_diagnostic(self):
+    """Keep an inline atomic ignore from hiding a function-level error."""
+    code = """
+import warp as wp
+
+def factory():
+  @wp.kernel
+  def kernel(arr_out: wp.array[float]):
+    arr_out[0] = 1.0  # kernel_analyzer: ignore
+"""
+    issues = ast_analyzer.analyze(code, "test.py", "", check_atomic=True)
+    self.assertTrue(any(isinstance(issue, ast_analyzer.MissingModuleUnique) for issue in issues))
+
+  def test_multiline_atomic_ignore_remains_local(self):
+    """Apply an ignore to its multiline atomic call only."""
+    code = """
+import warp as wp
+
+def factory():
+  @wp.kernel
+  def kernel(arr_out: wp.array[float]):
+    wp.atomic_add(
+      arr_out, 0, 1.0  # kernel_analyzer: ignore[atomic]
+    )
+"""
+    issues = ast_analyzer.analyze(code, "test.py", "", check_atomic=True)
+    self.assertTrue(any(isinstance(issue, ast_analyzer.MissingModuleUnique) for issue in issues))
+    self.assertFalse(any(isinstance(issue, ast_analyzer.MissingDeterministicFactory) for issue in issues))
+
+  def test_factory_missing_deterministic_arg_raises_issue(self):
+    """Reject an atomic kernel factory without a deterministic argument."""
+    bad_code = """
+import warp as wp
+
+def test_factory(n: int):
+  @wp.kernel(module="unique")
+  def kernel(arr_out: wp.array[int]):
+    wp.atomic_add(arr_out, 0, 1)
+"""
+    issues = ast_analyzer.analyze(bad_code, "test.py", "", check_atomic=True)
+    det_issues = [i for i in issues if isinstance(i, ast_analyzer.MissingDeterministicFactory)]
+    self.assertEqual(len(det_issues), 1)
+
+  def test_top_level_kernel_with_atomics_raises_issue(self):
+    """Reject a top-level kernel with an unconfigured atomic reduction."""
+    bad_code = """
+import warp as wp
+
+@wp.kernel
+def kernel(arr_out: wp.array[int]):
+  wp.atomic_add(arr_out, 0, 1)
+"""
+    issues = ast_analyzer.analyze(bad_code, "test.py", "", check_atomic=True)
+    det_issues = [i for i in issues if isinstance(i, ast_analyzer.MissingDeterministicFactory)]
+    self.assertEqual(len(det_issues), 1)
+
+  def test_kernel_with_atomics_in_deterministic_factory_passes(self):
+    """Accept atomic kernels inside a deterministic factory."""
+    good_code = """
+import warp as wp
+
+def test_factory(n: int, deterministic: bool = False):
+  @wp.kernel(module="unique")
+  def kernel(arr_out: wp.array[int]):
+    wp.atomic_add(arr_out, 0, 1)
+"""
+    issues = ast_analyzer.analyze(good_code, "test.py", "", check_atomic=True)
+    det_issues = [i for i in issues if isinstance(i, ast_analyzer.MissingDeterministicFactory)]
+    self.assertEqual(len(det_issues), 0)
+
+  def test_kernel_without_atomics_passes(self):
+    """Allow ordinary kernels without atomic operations."""
+    good_code = """
+import warp as wp
+
+def test_factory(n: int):
+  @wp.kernel(module="unique")
+  def kernel(arr_out: wp.array[int]):
+    arr_out[0] = 1
+"""
+    issues = ast_analyzer.analyze(good_code, "test.py", "", check_atomic=True)
+    det_issues = [i for i in issues if isinstance(i, ast_analyzer.MissingDeterministicFactory)]
+    self.assertEqual(len(det_issues), 0)
+
+  def test_ignore_suppresses_atomic_issue(self):
+    """Honor an explicit atomic diagnostic suppression."""
+    ignored_code = """
+import warp as wp
+
+@wp.kernel
+def kernel(arr_out: wp.array[int]):
+  wp.atomic_add(arr_out, 0, 1)  # kernel_analyzer: ignore
+"""
+    issues = ast_analyzer.analyze(ignored_code, "test.py", "", check_atomic=True)
+    det_issues = [i for i in issues if isinstance(i, ast_analyzer.MissingDeterministicFactory)]
+    self.assertEqual(len(det_issues), 0)
+
+  def test_func_with_atomics_missing_deterministic_factory(self):
+    """Reject an unconfigured device function containing atomics."""
+    bad_code = """
+import warp as wp
+
+@wp.func
+def helper(arr_out: wp.array[int]):
+  wp.atomic_add(arr_out, 0, 1)
+"""
+    issues = ast_analyzer.analyze(bad_code, "test.py", "", check_atomic=True)
+    det_issues = [i for i in issues if isinstance(i, ast_analyzer.MissingDeterministicFactory)]
+    self.assertEqual(len(det_issues), 1)
+
+  def test_func_with_atomics_in_deterministic_factory_passes(self):
+    """Accept an atomic device function inside a deterministic factory."""
+    good_code = """
+import warp as wp
+
+def test_factory(n: int, deterministic: bool = False):
+  @wp.func
+  def helper(arr_out: wp.array[int]):
+    wp.atomic_add(arr_out, 0, 1)
+"""
+    issues = ast_analyzer.analyze(good_code, "test.py", "", check_atomic=True)
+    det_issues = [i for i in issues if isinstance(i, ast_analyzer.MissingDeterministicFactory)]
+    self.assertEqual(len(det_issues), 0)
+
+  def test_bitwise_atomics_natively_ignored(self):
+    """Allow order-independent bitwise atomics without a factory."""
+    code = """
+import warp as wp
+
+@wp.kernel
+def kernel(flags_out: wp.array[int]):
+  wp.atomic_or(flags_out, 0, 1)
+  wp.atomic_and(flags_out, 0, 2)
+  wp.atomic_xor(flags_out, 0, 4)
+"""
+    issues = ast_analyzer.analyze(code, "test.py", "", check_atomic=True)
+    det_issues = [i for i in issues if isinstance(i, ast_analyzer.MissingDeterministicFactory)]
+    self.assertEqual(len(det_issues), 0)
+
+  def test_ignore_atomic_with_reason_comment(self):
+    """Honor an atomic ignore followed by its explanation."""
+    code = """
+import warp as wp
+
+@wp.kernel
+def kernel(arr_out: wp.array[int]):
+  wp.atomic_add(arr_out, 0, 1)  # kernel_analyzer: ignore[atomic]
+"""
+    issues = ast_analyzer.analyze(code, "test.py", "", check_atomic=True)
+    det_issues = [i for i in issues if isinstance(i, ast_analyzer.MissingDeterministicFactory)]
+    self.assertEqual(len(det_issues), 0)
+
+  def test_block_off_determinism(self):
+    """Limit atomic diagnostic suppression to its off/on block."""
+    code = """
+import warp as wp
+
+# kernel_analyzer: off[atomic]
+@wp.kernel
+def kernel1(arr_out: wp.array[int]):
+  wp.atomic_add(arr_out, 0, 1)
+# kernel_analyzer: on[atomic]
+
+@wp.kernel
+def kernel2(arr_out: wp.array[int]):
+  wp.atomic_add(arr_out, 0, 1)
+"""
+    issues = ast_analyzer.analyze(code, "test.py", "", check_atomic=True)
+    det_issues = [i for i in issues if isinstance(i, ast_analyzer.MissingDeterministicFactory)]
+    self.assertEqual(len(det_issues), 1)
+    self.assertEqual(det_issues[0].kernel, "kernel2")
+
+
 if __name__ == "__main__":
   absltest.main()

contrib/kernel_analyzer/kernel_analyzer/cli.py

@@ -24,9 +24,13 @@
 _VERBOSE = flags.DEFINE_bool("verbose", False, "Enable debug logging.")
 _OUTPUT = flags.DEFINE_enum("output", "console", ["console", "github"], "Analyzer output format.")
 _TYPES_PATH = flags.DEFINE_string("types", "", "Path to mujoco_warp types.py.")
+_CHECK_ATOMIC = flags.DEFINE_bool(
+  "check_atomic", True, "Enforce that kernels with atomic operations have a deterministic factory."
+)
 
 
 def main(argv):
+  """Analyze kernel sources and report convention or atomic-factory violations."""
   log_level = logging.DEBUG if _VERBOSE.value else logging.WARNING
   logging.set_verbosity(log_level)
 
@@ -65,7 +69,7 @@ def err_github(iss):
       types_source = types_path.read_text(encoding="utf-8")
       content = filepath.read_text(encoding="utf-8")
 
-      file_issues = ast_analyzer.analyze(content, str(filepath), types_source)
+      file_issues = ast_analyzer.analyze(content, str(filepath), types_source, check_atomic=_CHECK_ATOMIC.value)
       issues.extend(file_issues)
 
       for issue in file_issues:

contrib/kernel_analyzer/kernel_analyzer/lsp_server.py

@@ -43,7 +43,7 @@ def __init__(self, *args, **kwargs):
     self.issues: Dict[str, List[Any]] = {}
 
 
-_server = KernelAnalyzerLanguageServer("kernel-analyzer", "v0.8.2")
+_server = KernelAnalyzerLanguageServer("kernel-analyzer", "v0.9.0")
 
 
 @_server.feature(types.INITIALIZE)
@@ -63,7 +63,7 @@ async def initialize(client: LanguageServer, params: InitializeParams):
     ),
     server_info=InitializeResultServerInfoType(
       name="kernel-analyzer",
-      version="v0.8.2",
+      version="v0.9.0",
     ),
   )
 
@@ -74,9 +74,15 @@ async def initialize(client: LanguageServer, params: InitializeParams):
 async def validate(ls: KernelAnalyzerLanguageServer, params):
   """Validate the document using core_logic.ast_analyzer."""
   config = await ls.get_configuration_async(
-    WorkspaceConfigurationParams(items=[ConfigurationItem(scope_uri="", section="kernelAnalyzer.typesPath")])
+    WorkspaceConfigurationParams(
+      items=[
+        ConfigurationItem(scope_uri="", section="kernelAnalyzer.typesPath"),
+        ConfigurationItem(scope_uri="", section="kernelAnalyzer.checkAtomic"),
+      ]
+    )
   )
   type_source = Path(config[0]).read_text()
+  check_atomic = bool(config[1]) if len(config) > 1 and config[1] is not None else True
 
   text_doc = ls.workspace.get_text_document(params.text_document.uri)
   source = text_doc.source
@@ -85,7 +91,7 @@ async def validate(ls: KernelAnalyzerLanguageServer, params):
   logging.info(f"Validating document: {text_doc.uri}")
 
   try:
-    issues = ast_analyzer.analyze(source, text_doc.uri, type_source)
+    issues = ast_analyzer.analyze(source, text_doc.uri, type_source, check_atomic=check_atomic)
     logging.info(f"Analyzer found {len(issues)} issues in {text_doc.uri}")
     # store for potential future use (code actions)
     ls.issues[text_doc.uri] = issues