Add SSA compiler pipeline and unified backends#164
Conversation
ad873c2 to
a8249a9
Compare
a8249a9 to
9a4ba7a
Compare
437f718 to
9a4ba7a
Compare
This comment was marked as resolved.
This comment was marked as resolved.
This comment was marked as resolved.
This comment was marked as resolved.
voltjia
left a comment
voltjia
left a comment
There was a problem hiding this comment.
Submitting pending inline review comments.
|
关于模块和命名，建议把这个 PR 收敛成下面这套确定方案。核心原则是：模块负责承载上下文，public-ish 名称去掉重复的 模块布局 模块单复数建议也固定下来：表示一个领域或命名空间的 package 用单数/概念名，比如 IR 命名
下面这些 legacy / transitional kernel-level op IR 不建议保留在这个 PR 里，统一删掉或并入 SSA 入口： Frontend / render 命名 这里的重点是避免全局函数名重复携带 Pass 命名 对应 helper 也建议收短： 这些名字在 Backend 命名
后端入口建议从 这样外层仍然可以是 |
|
建议把新增 IR/backend/compiler 数据结构的构造方式统一收紧一下，避免这个 PR 还没合入时就固化一批容易错位的位置参数 API。 总体原则：
具体建议： @dataclass(frozen=True, kw_only=True)
class TensorSpec:
ndim: int
shape: tuple[str, ...] = ()
dtype: str | None = None
jagged_dim: int | None = None
constexpr: bool = False
name: str
attrs: Mapping[str, Any] = field(default_factory=dict)
@dataclass(frozen=True, kw_only=True)
class Type:
kind: str
shape: tuple[str, ...] = ()
dtype: str | None = None
attrs: Mapping[str, Any] = field(default_factory=dict)
其它 SSA IR 也建议 keyword-only： @dataclass(frozen=True, kw_only=True)
class Value:
name: str
type: Type
@dataclass(frozen=True, kw_only=True)
class Operation:
opcode: str
operands: tuple[str, ...] = ()
results: tuple[Value, ...] = ()
attrs: Mapping[str, Any] = field(default_factory=dict)
regions: tuple["Block", ...] = ()
@dataclass(frozen=True, kw_only=True)
class Block:
name: str = "entry"
args: tuple[Value, ...] = ()
operations: tuple[Operation, ...] = ()
@dataclass(frozen=True, kw_only=True)
class Program:
kind: str
inputs: tuple[Value, ...] = ()
outputs: tuple[Value, ...] = ()
blocks: tuple[Block, ...] = ()
metadata: Mapping[str, Any] = field(default_factory=dict)
@dataclass(frozen=True, kw_only=True)
class Launch:
name: str
args: tuple[str, ...] = ()
grid: str | None = None
@dataclass(frozen=True, kw_only=True)
class Kernel:
kernel_name: str
source: str
source_path: str | None = None
source_language: str = "triton"
entrypoint: str | None = None
launch: Launch | None = None
tensors: tuple[TensorSpec, ...] = ()
compiler_options: Mapping[str, Any] = field(default_factory=dict)
metadata: Mapping[str, Any] = field(default_factory=dict)
ssa: Program | None = NoneBackend/compiler 侧的 dataclass 也建议同样收紧：
其中 @dataclass(frozen=True, kw_only=True)
class _TensorInfo:
ndim: int = 1
shape: tuple[str, ...] = ()
dtype: str = "float32"
name: str
source_name: str | None = None
source_shape: tuple[str, ...] = ()
source_strides: tuple[str, ...] = ()
view_linear_offset: str | None = None
view_mask: str | None = None
attrs: Mapping[str, Any] | None = None调用点统一改成 keyword 形式，例如： TensorSpec(ndim=1, shape=("n",), dtype="float32", name="x")
ssa.Type(kind="tensor", shape=("n",), dtype="float32")
ssa.Value(name="x", type=type_)
ssa.Operation(opcode="mem.store", operands=(value, "out"))
Kernel(kernel_name="add", source=source, tensors=tensors, ssa=program)
Launch(name="launch_add", args=("x", "out"), grid="lambda meta: (1,)")最后建议补几条测试，防止后续回退：
|
|
这个 PR 目前已经引入了新的 SSA lowering / multi-backend 路径，但从现有主线入口来看，它还没有和原来的用户路径形成良好的集成。 我做了一些临时验证来确认这个集成状态。在 所以这里的主要问题不是“应该在某个函数里改一行切换过去”，而是这套新路径目前还缺少和主线 API 的系统性集成与验证。建议在这个 PR 或后续 PR 中明确补上：
换句话说，这条新路径本身可以继续演进，但目前还不宜认为它已经和原有 |
* Add Triton AOT multi-context design * Fix Triton AOT multi-context launches * Cover fresh Triton AOT context guard * Remove development process documents
* Address remaining SSA review comments * Sort test imports for Ruff * Remove redundant entrypoint assertions
3 participants
Summary
This PR introduces an SSA-first compiler pipeline and a unified multi-backend architecture for NineToothed. The new path is integrated with the existing
jit,make,aot, andbuildentry points, while keeping frontend lowering, compiler passes, backend emission, and runtime materialization as separate layers.Highlights
pytestoutput:All current GitHub Actions checks pass, including pytest, Ruff, distribution build, metadata, and documentation.