[feat] Smart magi-register-op

[themistbeforedawn]

🗂️ PR Category

• ✨ New Feature
• 🚀 Optimization (performance, memory, etc.)
• 💥 Breaking Change
• 🐛 Bug Fix
• 🛠️ Development / Refactoring
• 📚 Documentation
• 🧹 Chore (Dependencies, CI/CD, Configuration, etc.)
• 🧪 Testing

📝 Description

Summary

This PR introduces major enhancements to the @magi_register_custom_op decorator, enabling seamless support for structured data (nested frozen Dataclasses, Optionals, Literals) and ensuring that internal Triton kernels become fully transparent to the Inductor compiler (torch.compile).

The implementation acts as a robust bridge between the user-friendly Pythonic space and torch.library's restrictive flat-schema requirements without compromising performance.

Key Features & Changes

1. Deep Dataclass Flattening & Schema Bridging

• Nested structures: Fully supports arbitrarily nested @dataclass(frozen=True) as inputs. The decorator recursively flattens them into primitives for the torch.library schema and seamlessly reconstitutes the original dataclass struct during runtime execution.
• Enhanced type annotations: Handled graceful resolution and fallbacks. Added seamless mappings for Optional[Tensor/scalar], PEP-604 (Tensor | None), and lists (list[int/bool/Tensor]).
• Enums & Literals: Auto-downgrades Literal[str, ...] and string-valued Enum annotations to str to appease the PyTorch schema while passing the exact user string to the underlying op implementation.

2. Dataclass-Aware Autograd

• setup_context_fn and backward_fn are now bridged to directly support the original dataclass signature rather than flat internal parameters.
• Allows returning per-field gradients: Users can return None (for dropping gradients across entire dataclasses), dictionary mappings, or reconstructed dataclass shapes for accurate partial gradients out of backward_fn.
• Built-in support for mixed tuple/tensor backward returns alongside dataclass kwargs.

3. Advanced Triton Kernel Auto-Detection (triton_op)

• Inductor See-Through: Op boundaries dynamically intercept bare kernel[grid](...) calls using AST introspection, overriding the module __globals__ to emit wrap_triton(kernel) under the hood. Inductor correctly traces into the kernels rather than treating the op as opaque.
• Double-Wrap Prevention (Idempotency): Built an AST visitor pass that detects cases where users manually invoke wrap_triton(...). Excludes these IDs from the shadow-rewriter, avoiding wrap_triton(wrap_triton(kernel)) exceptions.
• Supported retrieving complex retrieval patterns: module-level, closure-wrapped, static/class methods, and externally imported third-party helpers.
• Automatically surfaces friendly, actionable errors if a user puts @triton.heuristics at the top level of a kernel (which PyTorch currently cannot wrap).

4. Friendly Error Boundaries & Fallbacks

• Drastically improved debugger experiences by catching schema violations at decoration time with highly actionable error logs. Examples include explicit blocks on tuple/dict primitives, un-trackable nested mutable containers, Dataclass return types, missing namespaces (namespace::name), and duplicate operation names.

5. Massive Testing Refactor

• Completely re-architected tests/api_tests/test_register_custom_op.py (79 tests) and test_register_triton_op.py (46 tests) away from procedural naming into grouped, theme-driven logic (Core Registration, Type Support, Autograd Bridge, Compositions). Added a detailed Markdown Coverage Matrix bounding the top of each file.