[Architecture] Adopt an AST-first TileLang DSL architecture to decouple frontend, backend, and JIT workstreams

[Zhendong404]

Background

We want to converge the current tile-level DSL and vector-level DSL into one unified TileLang DSL stack, instead of maintaining two partially overlapping authoring systems.

Today the split is roughly this:

• pto-dsl carries a substantial amount of tile-level semantics and a JIT-oriented workflow
• tilelang_dsl carries the current vector-level / VPTO-oriented frontend path inside PTOAS and already participates in TileOp expansion / VPTO authoring

This split is increasingly awkward for both users and compiler engineering:

• tile-level and vector-level authoring are not expressed in one coherent DSL architecture
• similar semantics are spread across two different codebases and two different programming models
• the long-term backend direction is unclear because the two DSLs were built on different technical foundations
• frontend, lowering backend, and runtime/JIT concerns are not cleanly separated today

The goal of this proposal is therefore not just “add one more frontend feature” or “add JIT”.
The goal is to define a unified architecture where:

• tile-level and vector-level semantics live in the same DSL stack
• there is one clear public programming model
• frontend work and backend work can proceed independently
• JIT can later be added on top of the same compilation pipeline instead of becoming a separate system

Main difficulties

There are two architectural difficulties that must be resolved first.

1. pto-dsl and tilelang_dsl use different technical routes

The two DSLs were built with different assumptions:

• pto-dsl is fundamentally IR-builder oriented
• tilelang_dsl is fundamentally AST-oriented

This means we do not just have two implementations of similar semantics; we have two competing frontend models.

Before we can merge tile-level and vector-level DSL capabilities, we must make a deliberate technical choice about the unified route:

• Should the long-term public DSL be builder-first?
• Or should it be AST-first, with builder as an internal implementation backend?

Without resolving this, any “merge” would only move code around while preserving the underlying split.

2. current tilelang_dsl still uses a text-emission backend and frontend/backend coupling is too strong

Current tilelang_dsl lowering is still centered on emitting authoring-form VPTO MLIR text.

That creates two related problems:

• there is no proper pybinding / MLIR-builder backend yet
• frontend semantics and backend emission logic are too tightly coupled

In practice, this means:

• frontend and lowering evolution interfere with each other
• changing the backend implementation is expensive because the frontend is not isolated from emission details
• JIT cannot be cleanly layered on top, because there is not yet a stable compile-driver boundary

So even before discussing JIT, we need a clearer split between:

• frontend language / semantics
• backend MLIR construction
• compile driver / runtime launch

Why choose an AST-first route

The proposed direction is:

• keep AST as the only primary public frontend
• use MLIR builder / pybinding as the target lowering backend

This is not a rejection of builder infrastructure. It is a decision about the user-facing language boundary.

Why AST is the better unification route

If we adopt a builder-first public route, then the user-facing DSL becomes tightly coupled to IR-construction mechanics:

• authoring style starts to resemble “writing a Python script that builds compiler IR”
• frontend evolution becomes tied to backend builder APIs
• syntax sugar, template-style authoring, type-directed specialization, and source diagnostics become harder to manage as language features
• merging tile-level and vector-level authoring still leaves us with a compiler-API-shaped user model

An AST-first route gives us a cleaner long-term unification point:

• users write kernel semantics, not IR construction steps
• tile-level and vector-level constructs can be expressed in one source language
• templates, polymorphism, symbolic analysis, and source-level diagnostics stay in the frontend
• the frontend can lower into a backend-neutral semantic representation
• the backend can evolve from text emission to pybinding builder without changing the public programming model

What AST does not mean

Choosing AST does not mean:

• moving optimization logic into Python
• replacing MLIR/PTO/VPTO/LLVM passes with frontend logic
• devaluing builder-based infrastructure

Most real optimization should still remain in:

• PTO / VPTO passes
• MLIR canonicalization / cleanup
• LLVM / Bisheng backend

The value of AST here is mainly:

• one stable public language surface
• frontend/backend decoupling
• a better basis for unifying tile-level and vector-level authoring

Proposal: how to solve the coupling problem

The key architectural move is to introduce a backend-neutral semantic layer between the AST frontend and all backend/runtime work.

Target architecture:

User API
  -> tilelang_dsl AST surface
      -> Frontend Analyzer
          -> Canonical Semantic Module
              -> MLIR Text Emitter (legacy / fallback)
              -> MLIR Builder Emitter (target)
              -> Template Instantiation Service
                  -> Compile Driver
                      -> ptoas VPTO LLVM path
                      -> Bisheng / cce-ld
                      -> Runtime Launcher / JIT

Core rule

The Canonical Semantic Module becomes the only internal contract shared by all workstreams.

That gives us a clean separation of responsibilities:

• Frontend owns parsing, semantic analysis, specialization, source diagnostics, and language sugar
• Emitters own conversion from semantic IR to MLIR
• Compile driver owns artifact generation through the current VPTO LLVM path
• Runtime/JIT owns caching, load, argument bridging, stream handling, and launch

Practical workstream split

A. frontend/backend decoupling first

This is the highest-priority step.

We should explicitly split tilelang_dsl into:

• frontend_ast
• semantic_model
• emitters
• compile_driver/runtime

This is the prerequisite for every later step.

B. import pto-dsl semantics into the AST + semantic layer

pto-dsl should contribute:

• tile-level semantic coverage
• module organization ideas
• entry metadata / multi-function module patterns

But it should not become the primary public builder-style frontend.

C. build a real pybinding backend behind the semantic layer

The current text emitter can remain temporarily as:

• a reference path
• a fallback path
• a debugging aid

But the target direction should be:

• SemanticModule -> MLIR builder / pybinding as the official lowering backend

D. add JIT after the boundary is stable

JIT should be layered on top of the compile driver, not fused into the frontend.

That way:

• JIT reuses the same frontend and lowering stack
• JIT does not become a second architecture
• compile-only and compile-and-run can share the same artifact model

Real compile chain we should design for

The architecture must align with the actual current PTOAS VPTO LLVM path:

AST / semantic frontend
  -> MLIR
  -> ptoas --pto-backend=vpto --vpto-emit-hivm-bc|llvm
  -> LLVM IR / bitcode
  -> Bisheng device object
  -> fatobj / kernel.so packaging
  -> Python load + launch

This proposal does not assume the old kernel.cpp / caller.cpp path as the long-term primary model.

Expected outcome

If we follow this route, we get:

1. one unified DSL architecture for tile-level and vector-level authoring
2. one stable public programming model based on AST
3. one backend-neutral semantic contract for frontend/backend parallel work
4. one clean migration path from text emission to pybinding builder lowering
5. one future JIT path built on top of the same compile driver instead of a separate runtime stack

Non-goals

This proposal is not trying to:

• rewrite everything in one patch
• remove the current text emitter immediately
• promise full public ptodsl API compatibility
• define every runtime ABI detail in this issue
• claim that AST itself is the main optimization engine

Discussion points for the team

1. Do we agree that the real background problem is unifying tile-level DSL and vector-level DSL into one architecture, not merely adding another frontend feature?
2. Do we agree that the first architectural conflict to resolve is IR-builder route vs AST route?
3. Do we agree that the second architectural blocker is the lack of a proper pybinding backend and the current frontend/backend coupling inside tilelang_dsl?
4. Do we agree that AST should remain the only primary public frontend, while builder becomes the internal lowering backend?
5. Do we agree that the semantic module should be the only contract between frontend, backend, and future JIT/runtime work?

Related issues

• [Architecture] TileLang DSL lowering backend 渐进迁移到 pybinding builder #237 [Architecture] TileLang DSL lowering backend 渐进迁移到 pybinding builder
• [Architecture] VPTO IR cache 机制应迁移到 DSL 框架侧 #210 [Architecture] VPTO IR cache 机制应迁移到 DSL 框架侧
• epic(tilelang-dsl): complete cube matrix tileop templates and ST coverage #337 epic(tilelang-dsl): complete cube matrix tileop templates and ST coverage

This issue is intended to act as the top-level architecture discussion that explains how these tracks fit together.

[learning-chip]

Why choose an AST-first route

The proposed direction is:

• keep AST as the only primary public frontend
• use MLIR builder / pybinding as the target lowering backend

This is not a rejection of builder infrastructure. It is a decision about the user-facing language boundary.

AST-heavy frontend has many risks. I learned hard lessons when working on the pypto's ast parser last year. It took me days and nights to debug endless ast-induced bugs, like this bug, this bug, this bug, and many more.

The problems of AST-heavy parsing:

• If the parser is built on ast.NodeVisitor, ast.NodeTransformers, or just plain ast module, the parser code is much longer than just eagerly calling MLIR python bindings (like 10K lines vs 1K lines). This introduces unnecessary complexity for maintaining and debugging, especially when we have such a rush development pace.
• You need to very carefully think about how to interpret Python-native constructs, such as for i in range(), list, dict, class, dataclass, iterators, etc. Should all Python-native objects be allowed, or should some be banned? Should for ... range() be interpreted as Python eager for, or scf.for IR, and how to express loop unrolling and mark dynamic vs constant loop bounds? When should if cond: be lower to scf.if IR, or just compile-time evaluated branch (like #ifdef). Of course you can define clearly the behavior of each corner cases , see CuteDSL control flow notes, but we have much simpler way to avoid this trouble (see below)

I intentionally avoid any ast parsing in current ptodsl. The core ptodsl/api has very very short & thin code on top of MLIR python binding. It just reuses what ptoas already has (python binding auto-generated from mlir tblgen). This reduces the burden for framework maintainers, so that we can focus on performance engineering and feature coverage, which is more important and urgent now. The frontend can still be much improved, before resorting to ast parsing. In this way, we don't need to think about how to deal with Python-native objects -- they just run like normal Python code. All IR builder behaviors are inside the pto.xxx namespace.

I think a light-weight (not heavy) usage of ast is ok, to make the frontend syntax look a bit more Python-native. CuteDSL frontend uses a hybrid ast+tracing approach, first transforming if else for to IR builder calls, and then call the MLIR python bindings. We can use the same approach in the long term.

Anyways, such little "frontend improvement" is not critical point. Just pure tracing without ast can also enable a quite clean frontend. JAX tracing and pyptx are two excellent examples of non-ast Python frontend. Their code style also looks quite decent and clean (just with ptx.if_(tid == 0): instead of if tid == 0:), and requires the framework engineer to think much less about corner cases.

Other framework's frontend for comparison

Pure tracing, no AST: JAX, pyptx as above.

Hybrid ast mutation (parsing) + IR builder (tracing):

• TileLang mutate/DSLMutator, then call TVM IR builder
• CuteDSL DSLPreprocessor, then call MLIR/scf IR builder
• FlyDSL ast_rewriter.py
• eudsl-python-extras (thin wrapper on top of standard MLIR (arith, scf, func, ...) python bindings

Pure AST parsing

• CuTile _ast2ir
• Triton CodeGenerator

[Zhendong404]

PTODSL is taking the tracing routine