Version-Aware Fusion Plan#912
Conversation
… and constraint support
Codex Review该评论由 review 机器人自动更新。
SummaryReview failed at stage Findings未生成结构化 findings，因为 review 过程提前失败。 Log Tail |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a version-aware tile fusion planning pass and a new PTODSL TileLib backend for template expansion. The changes include adding a new VersionAwareFusionPlan pass, a shared metadata contract for template candidates, and a daemon-based RPC mechanism for template rendering. My review identified several security and reliability issues, including insecure socket creation, potential DoS vulnerabilities in the daemon, missing module imports, and missing template files. I have provided actionable feedback to address these concerns.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| super().__init__(socket_path, _Handler) | ||
| os.chmod(socket_path, 0o600) |
There was a problem hiding this comment.
Creating a Unix domain socket in a world-writable directory like /tmp can lead to a race condition where other local users can connect to the socket before os.chmod restricts its permissions. To prevent this, temporarily set the process umask to 0077 before binding the socket, or create the socket in a secure, user-private directory.
| super().__init__(socket_path, _Handler) | |
| os.chmod(socket_path, 0o600) | |
| import os | |
| old_umask = os.umask(0o077) | |
| try: | |
| super().__init__(socket_path, _Handler) | |
| finally: | |
| os.umask(old_umask) | |
| os.chmod(socket_path, 0o600) |
| @@ -18,25 +18,74 @@ | |||
| from pathlib import Path | |||
There was a problem hiding this comment.
The module sys is used in _pythonpath_already_provides_mlir_and_pto (via sys.path), but it is not imported in this file. This will result in a NameError at runtime. Please import sys at the top of the file.
| from pathlib import Path | |
| import sys | |
| from pathlib import Path |
| ("a5", "pto.tpartadd"): ".a5.tpartadd", | ||
| ("a5", "pto.tpartmax"): ".a5.tpartmax", | ||
| ("a5", "pto.tpartmin"): ".a5.tpartmin", | ||
| ("a5", "pto.tpartmul"): ".a5.tpartmul", |
There was a problem hiding this comment.
The template modules tpartadd, tpartmax, tpartmin, and tpartmul are registered in _TEMPLATE_MODULES, but their corresponding files (tpartadd.py, tpartmax.py, tpartmin.py, tpartmul.py) are missing from this pull request. Attempting to load or expand these operations will result in a ModuleNotFoundError at runtime.
| ("a5", "pto.trowexpandsub"): ".a5.trowexpandsub", | ||
| ("a5", "pto.trowmax"): ".a5.trowmax", | ||
| ("a5", "pto.trowmin"): ".a5.trowmin", | ||
| ("a5", "pto.trowprod"): ".a5.trowprod", |
There was a problem hiding this comment.
The template module trowprod is registered in _TEMPLATE_MODULES as ".a5.trowprod", but the file ptodsl/ptodsl/tilelib/templates/a5/trowprod.py is missing from this pull request. Attempting to load or expand pto.trowprod will result in a ModuleNotFoundError at runtime.
| class _Handler(socketserver.BaseRequestHandler): | ||
| def handle(self): | ||
| try: |
There was a problem hiding this comment.
Since TileLibDaemonServer is a sequential (single-threaded, blocking) server, a client that connects and sends a large length prefix but no payload can block the daemon indefinitely, causing a Denial of Service (DoS). Setting a socket timeout on the connection socket (e.g., in setup or at the start of handle) prevents such hangs.
| class _Handler(socketserver.BaseRequestHandler): | |
| def handle(self): | |
| try: | |
| class _Handler(socketserver.BaseRequestHandler): | |
| def setup(self): | |
| self.request.settimeout(10.0) | |
| def handle(self): | |
| try: |
| chunks = [] | ||
| remaining = length | ||
| while remaining: | ||
| chunk = sock.recv(remaining) |
There was a problem hiding this comment.
Calling sock.recv(remaining) with a very large remaining value (up to 64MB) can cause Python to allocate a large buffer in memory before reading. To prevent excessive memory allocation and potential Out-Of-Memory (OOM) issues, read the socket in smaller, capped chunks (e.g., at most 64KB per recv call).
| chunk = sock.recv(remaining) | |
| chunk = sock.recv(min(remaining, 65536)) |
| def __post_init__(self): | ||
| if len(self.shape) != 2: | ||
| raise ValueError("TileSpec currently only supports rank-2 tile shapes") | ||
| if any(not isinstance(dim, int) or dim <= 0 for dim in self.shape): | ||
| raise ValueError("TileSpec.shape must contain positive integers") |
There was a problem hiding this comment.
The valid_shape parameter in TileSpec is not validated in __post_init__. If provided, it should be checked to ensure it has the same rank as shape, contains positive integers, and does not exceed the physical shape dimensions.
def __post_init__(self):
if len(self.shape) != 2:
raise ValueError("TileSpec currently only supports rank-2 tile shapes")
if any(not isinstance(dim, int) or dim <= 0 for dim in self.shape):
raise ValueError("TileSpec.shape must contain positive integers")
if self.valid_shape is not None:
if len(self.valid_shape) != 2:
raise ValueError("TileSpec.valid_shape must be rank-2")
if any(not isinstance(dim, int) or dim <= 0 for dim in self.valid_shape):
raise ValueError("TileSpec.valid_shape must contain positive integers")
if any(v > s for v, s in zip(self.valid_shape, self.shape)):
raise ValueError("TileSpec.valid_shape dimensions cannot exceed physical shape")| with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock: | ||
| sock.connect(self.socket_path) |
There was a problem hiding this comment.
The client socket is created without a timeout. If the daemon hangs or becomes unresponsive, the client (and consequently ptoas) will block indefinitely. Setting a reasonable timeout (e.g., 30 seconds) on the client socket prevents infinite hangs.
| with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock: | |
| sock.connect(self.socket_path) | |
| with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock: | |
| sock.settimeout(30.0) | |
| sock.connect(self.socket_path) |
3 participants
Description
This draft PR starts the version-aware fusion planning work for Tile Fusion.
The main motivation is that a single
TileOp/FusionComputeNodemay havemultiple legal implementation versions. For example, when tile regions are
contiguous, a tile op may be implementable with a simpler 1D loop nest;
otherwise it may need a different loop structure. The fusion planner therefore
should not treat each compute node as a single fixed unit. Instead, it should
eventually reason about:
The long-term goal is to enumerate possible fusion groups per seed across these
implementation versions, score them with a cost model, and select the best
version-aware group.
This PR is currently a foundation/draft PR. It does not yet replace the
existing
PTOFusionPlanbehavior and does not yet wire in the fullversion-aware search algorithm.
Goal
InsertTemplateAttributesPTOVersionAwareFusionPlan(compute node, implementation version).version-aware search.
Current Algorithm Direction
The intended version-aware planning algorithm is:
flowchart TD A[Block compute nodes] --> B[Pick next unassigned seed] B --> C[Read legal versions for seed] C --> D[Create one GroupState per seed version] D --> E[Explore append candidates] E --> F[For each candidate node] F --> G[Read legal versions for candidate] G --> H[Try appending each version] H --> I{Append legal by DAG / domain / boundary rules?} I -- no --> E I -- yes --> J[Create new GroupState] J --> K[Compute base fusion cost] K --> L[Add version trait cost] L --> M[Record valid group states] M --> N{More expansions?} N -- yes --> E N -- no --> O[Select best group by deterministic comparator] O --> P[Assign group_id / order] P --> BEach
GroupStaterepresents one possible version-aware group:The key difference from the current node-only planner is that the search
frontier branches by implementation version:
Cost Model Direction
The planner will compare candidate groups using both existing fusion cost and
version-specific cost.
Base fusion terms:
Version-aware terms added by this PR:
Current version-trait scoring helper models:
The deterministic tie-break order is:
Current State of This PR
Implemented so far:
include/PTO/Transforms/TemplateAttributes.hlib/PTO/Transforms/TemplateAttributes.cppcandidatesidnameloop_depthpostupdatetailInsertTemplateAttributes.cppto build candidate attrs through theshared helper.
PTOVersionAwareFusionPlan.cppto parse candidate attrs through theshared helper.
TileOpImplVersionPlannedFusionMemberPlannedFusionGroupGroupStateGroupStateper legal implementation versionNot implemented yet:
planBlock.formation behavior.
PTOFusionPlan.cppremains unchanged intentionally.