Skip to content

[Feature] Support Python native and and or expressions in PTODSL #1332

Description

@liggest

Summary

PTODSL RuntimeValue 目前不能直接用于 Python 原生 andor 布尔表达式
编译时会尝试把左侧条件使用 bool(...) 转换,引发报错

希望能够支持在 PTODSL 中正常使用 andor 布尔表达式,并保持短路行为:

  • x and y => x 为假时不对 y 求值
  • x or y => x 为真时不对 y 求值

and 简单样例

from ptodsl import pto, scalar


@pto.jit(name="and_kernel", kernel_kind="vector", target="a5")
def kernel(value: pto.i32, divisor: pto.i32, out: pto.ptr(pto.i8, "gm")):
    pred = (divisor != 0) and ((value // divisor) > 0)
    scalar.store(pred, out, 0)


kernel.compile()

or 简单样例

from ptodsl import pto, scalar


@pto.jit(name="or_kernel", kernel_kind="vector", target="a5")
def kernel(value: pto.i32, divisor: pto.i32, out: pto.ptr(pto.i8, "gm")):
    pred = (divisor == 0) or ((value // divisor) > 0)
    scalar.store(pred, out, 0)


kernel.compile()

两种情况都会在编译阶段报错:

ptodsl._diagnostics.PTODSLTracingMisuseError: native Python if/while condition cannot consume a PTODSL runtime value during tracing. This value is a device-side SSA/runtime-metadata value, not a Python bool/int. Use pto.if_(...) or pto.for_(...) for device-side control flow, or keep the bound/condition in pto.const_expr.

Motivation / use case

支持 Python 原生 andor 布尔表达式,能帮助用户更方便地在 PTODSL 中表达复杂一些的条件语句,这在实际算子逻辑中时不时会出现

同时,布尔表达式的短路行为能在不满足左侧(x)条件时减少对右侧(y)无意义的求值,在右侧计算开销较大,或右侧只在左侧满足时才有效的场景中有助于提升计算效率

Proposed API / behavior

No response

Alternatives considered

No response

Additional context

No response

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or requesttilelangIssues from tilelang

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions