[Bugfix] Copy pass_configs dict to prevent mutation across multiple JIT compilations

[LeiWang1999]

Summary

• Fix a bug where pass_configs dictionary was shared and mutated across multiple JIT kernel compilations
• When calling a JIT-decorated function multiple times, the compile_flags would accumulate (e.g., ['-O3'] → ['-O3', '-O3'])
• Solution: Create a copy of pass_configs before modifying it in _compile_and_create_adapter

Root Cause Analysis

This is a classic Python mutable object sharing issue.

How Python handles objects

In Python, variables store references (pointers) to objects, not the objects themselves:

a = {"key": 1}
b = a          # b and a point to the same dict object
b["key"] = 2
print(a)       # {'key': 2} — a is also modified!

What happened in this bug

When using the @tilelang.jit decorator:

@tilelang.jit(
    pass_configs={...},   # This dict object is created only ONCE
    compile_flags=["-O3"],
)
def matmul(...):
    ...

The pass_configs dictionary is created once when the decorator is applied. Every subsequent call to matmul() reuses the same dictionary object.

When the code modifies this dictionary:

pass_configs[key] = compile_flags  # Mutates the shared dict

All future calls see the accumulated modifications:

• 1st call: compile_flags_cfg = None → sets ['-O3']
• 2nd call: compile_flags_cfg = ['-O3'] → sets ['-O3', '-O3']
• 3rd call: compile_flags_cfg = ['-O3', '-O3'] → sets ['-O3', '-O3', '-O3']

The fix

Copy the dictionary before modifying:

# Before (shared dict)
pass_configs = self.pass_configs or {}

# After (new copy each time)
pass_configs = dict(self.pass_configs) if self.pass_configs else {}

Test plan

• Run the reproduction script to verify flags no longer accumulate across multiple compilations

Summary by CodeRabbit

• Bug Fixes
  • Improved object handling to prevent unintended mutations during configuration processing.

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[coderabbitai]

Caution

Review failed

The pull request is closed.

📝 Walkthrough

Walkthrough

The PR modifies the _compile_and_create_adapter method in tilelang/jit/kernel.py to create a shallow copy of pass_configs when it exists, replacing the original conditional assignment. This prevents in-place mutations to the original dictionary while maintaining fallback behavior for missing configs.

Changes

Cohort / File(s)	Summary
Pass configs dictionary handling tilelang/jit/kernel.py	Changed pass_configs assignment to create a shallow copy of the existing dictionary using dict() constructor, preventing in-place mutations to the original object.

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~3 minutes

Possibly related PRs

• [Bugfix] Convey compile_flags to ffi compilation path with pass_configs #1434: Also modifies tilelang/jit/kernel.py's pass_configs handling, adding TL_DEVICE_COMPILE_FLAGS merging in the same method.

Poem

🐰 A copy made with gentle care,
No mutations in the air,
The dict stays pure, untouched, serene—
The safest pass_configs ever seen! ✨

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch flag_0203

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}