[sepcnt]

Fixes #3274.

SIMT im2col becomes visible too late for scheduling.
a6bbc5c is only a minimal PoC for now.

Change

Add opt-in ExpandSIMTIm2Col, enabled by tl.enable_early_simt_im2col=True:

• Expose guarded parallel loads before WS, pipeline planning, and layout inference.
• Select by the actual SIMT implementation.
• Retain tail zero-filling, destination offsets, annotated operations, and manual WS schedules.

Performance

RTX 5090, NVCC 13.2.78, sm_120a. Same convolution kernel and WS setting; two seeds × nine randomized rounds. Median GPU µs; speedup is the ratio of medians.

Case	WS	Off	On	Speedup
aligned	off	14.880	10.614	1.40×
aligned	on	15.505	9.619	1.61×
stride_tail	off	23.081	8.702	2.65×
stride_tail	on	10.583	6.256	1.69×
large	off	159.410	106.176	1.50×
large	on	210.306	106.255	1.98×

Approximately matches explicit loops. WS uses conditional 16-byte async input copies; ptxas registers drop from 240 to 78. Large without WS increases register usage and lowering time, so the pass remains opt-in.

Summary

Adds the opt-in tl.enable_early_simt_im2col pass. When enabled for a CUDA target that selects the SIMT im2col implementation, the pass expands eligible Im2ColOp calls into guarded parallel loops before scheduling.

The rewrite preserves bounds checks and zero-fills invalid elements. It skips annotated operations, manual warp-specialization schedules, and unsupported shapes or buffer regions. Specialized implementations remain unchanged.

The change adds Im2ColUsesSIMT to query the selected implementation. Tests cover pass selection, idempotence, unsupported cases, output correctness, and convolution behavior. Test execution results were not provided.

C++ style / lint notes

The PR changes C++ source and a public header. The rules in docs/developer_guide/cpp_style.md apply. The guide recommends Is/Has/Can prefixes for Boolean functions; Im2ColUsesSIMT does not follow that convention. This is an advisory style observation, not a reported correctness or build issue.

CI includes the C++ API Style Audit (warning only) step. No audit findings for this PR were supplied. No build or test results were supplied.

[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]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: tile-ai/tilelang/.coderabbit.yaml
• Review profile: CHILL
• Plan: Advanced
• Run ID: 4801012e-ee3e-455a-ae69-4e89acf4fc6e

📥 Commits

Reviewing files that changed from the base of the PR and between a6bbc5c and 17ddc0e.

📒 Files selected for processing: 2

• src/op/copy.h
• testing/python/transform/test_tilelang_transform_im2col_fallback.py

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 3 remain after this review.

📝 Walkthrough

The change adds a disabled-by-default CUDA pass that can expand eligible SIMT Im2Col operations into guarded parallel loops before scheduling and pipeline planning. It adds a query for the selected Im2Col implementation and tests target selection, preservation cases, generated source, and execution results.

Changes: CUDA SIMT Im2Col Expansion

Layer / File(s)	Summary
SIMT selection and opt-in configuration src/op/copy.h, src/op/copy.cc, tilelang/transform/pass_config.py, src/cuda/transform/expand_simt_im2col.cc	Adds a query for whether the selected Im2Col implementation uses SIMT and a disabled-by-default tl.enable_early_simt_im2col configuration key.
Eligible operation expansion src/cuda/transform/expand_simt_im2col.cc, tilelang/cuda/transform/__init__.py, testing/python/transform/test_tilelang_transform_expand_simt_im2col.py	Adds a pass that rewrites supported Im2Col calls as guarded parallel loops. The pass leaves unsupported calls and annotated blocks unchanged. Tests cover target selection, idempotence, and preservation cases.
Pipeline integration and execution checks tilelang/cuda/pipeline.py, testing/python/transform/test_tilelang_transform_expand_simt_im2col.py, testing/python/transform/test_tilelang_transform_im2col_fallback.py	The CUDA prologue conditionally runs the expansion and simplifies the IR. Tests compare generated source and execution results. The fallback test helper accepts a configurable pipeline stage count.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Bug fix

Sequence Diagram(s)

sequenceDiagram
  participant CUDAPassPipelineBodyPrologue
  participant UsesSIMTIm2Col
  participant ExpandSIMTIm2Col
  participant Simplify
  CUDAPassPipelineBodyPrologue->>UsesSIMTIm2Col: query target when config is enabled
  UsesSIMTIm2Col-->>CUDAPassPipelineBodyPrologue: return SIMT selection
  CUDAPassPipelineBodyPrologue->>ExpandSIMTIm2Col: expand eligible operations
  CUDAPassPipelineBodyPrologue->>Simplify: simplify expanded IR

Loading

Suggested labels: bug, python

Suggested reviewers: leiwang1999, yongqi-zhuo

Merge Risk: ⚪ Minimal · up to 17ddc

The expansion is disabled by default. When enabled, it selects the SIMT path and applies guarded expansion before scheduling; the reviewed changes present no actionable merge risk.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 22.22% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 27 functions across 8 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly and concisely describes the main CUDA transform change: exposing SIMT im2col loads before scheduling.
Linked Issues check		The PR meets the coding requirements in [#3274]. CUDAPassPipelineBodyPrologue invokes ExpandSIMTIm2Col after simplification and before warp specialization, pipeline planning, and layout inference.…
Out of Scope Changes check		The changed files stay within [#3274]. They add the opt-in configuration, selected-implementation query, early expansion pass, pipeline integration, public bindings, and focused tests. The `num_stages…

Full details: Linked Issues check

Explanation

The PR meets the coding requirements in [#3274]. CUDAPassPipelineBodyPrologue invokes ExpandSIMTIm2Col after simplification and before warp specialization, pipeline planning, and layout inference. Im2ColUsesSIMT checks the selected lowering callback, so specialized implementations remain unchanged. The expander emits parallel guarded loads, preserves zero fill and destination offsets, and skips annotated operations and manual tl.ws_schedule blocks. The added tests cover implementation selection, metadata preservation, SM90 stability, tails, offsets, and convolution execution.

Full details: Out of Scope Changes check

Explanation

The changed files stay within [#3274]. They add the opt-in configuration, selected-implementation query, early expansion pass, pipeline integration, public bindings, and focused tests. The num_stages helper parameter supports the affected pipeline tests and keeps the default behavior unchanged. No unrelated product change is demonstrated.

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

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.}