[EagerJIT] Lazy Evaluation of Kernel Body in Eager JIT (#1690)

[kurisu6912]

This function modify the elaboration process of EagerJIT

In the previous version of EagerJIT:

1. Elaborate the function, and obtain signature
2. Parse const arguments from signature
3. Substitute the function with the constexprs

However, #1690 gives an bad case: the constexprs are used in T.alloc_xxx. In this case, the shape of the allocated buffer is a unknown variable, thus compilation error arises.

In this pr, we modify the elaboration process of EagerJIT:

1. Elaborate the function, but stop at the first T.Kernel, only retrieve the signature
2. Parse const arguments from the signature
3. Re-elaborate the function with the constexprs

Summary by CodeRabbit

• Bug Fixes

  • Improved error reporting by using more appropriate exception types.

• Refactor

  • Enhanced eager JIT execution model with phase-based compilation and template improvements.
  • Added return type annotations for clearer APIs and IDE support.

• Tests

  • Added a new unit test covering a TileLang reduction/jitting scenario.
  • Removed redundant test code and pre-compilation steps.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

[coderabbitai]

Actionable comments posted: 1

🤖 Fix all issues with AI agents

In `@tilelang/language/eager/builder.py`:
- Around line 892-901: The phase2 branch in the eager JIT substitution (when
builder.eager_jit == "phase2") currently indexes builder.eager_jit_subs[name]
(and for comma/space-separated lists) which can raise an unhelpful KeyError for
unknown constexpr names; update the code in that branch to validate each name
exists in builder.eager_jit_subs before returning substitutions and raise a
clear, descriptive error (e.g., raise KeyError or ValueError with a message that
lists the missing names and the original requested name(s) and mentions phase2)
so callers can see which constexpr(s) were not registered in phase1; apply this
check for the single-name path and the tuple paths that iterate over names.

♻️ Duplicate comments (1)

tilelang/language/eager/builder.py (1)

977-990: Add defensive null check for ir_gen when is_lazy_style is false.

The type annotation allows ir_gen: IRGenerator[_P, _T] | None = None, but line 986 calls self.ir_gen.gen(builder) without checking for None. Add a guard after the lazy style check:

🐛 Proposed fix

     def get_tir(self, tensor_args, given_tensor_args, kwargs):
         if self.is_lazy_style:
             return self.prim_func
+        if self.ir_gen is None:
+            raise RuntimeError(f"ir_gen not available for eager-style template '{self.name}'")
         values = self._parse_phase2_key(**given_tensor_args, **kwargs)

🧹 Nitpick comments (1)

tilelang/language/eager/builder.py (1)

917-917: Type annotation mismatch for constexprs field.

The field is typed as set[Var] but defaults to None. For type correctness, annotate as set[Var] | None = None.

✏️ Suggested fix

-    constexprs: set[Var] = None
+    constexprs: set[Var] | None = None