diff --git a/docs/triton_plugin.md b/docs/triton_plugin.md new file mode 100644 index 0000000000..1ecaec0118 --- /dev/null +++ b/docs/triton_plugin.md @@ -0,0 +1,39 @@ +# Triton support + +Pyrefly recognizes `tl.constexpr` parameter annotations on functions decorated with +`@triton.jit` or `@triton.jit(...)`. Imported aliases and re-exports also work. + +Triton uses `tl.constexpr` to mark a compile-time argument. The argument can contain +an integer, a Boolean, a string, a tuple, or another compile-time value. +The annotation does not require a Python `constexpr` instance. See the +[Triton constexpr implementation](https://github.com/triton-lang/triton/blob/main/python/triton/language/core.py). + +Pyrefly currently represents these parameter values as `Any`. This permits scalar +operations and ordinary Python defaults without errors about the `constexpr` +class. Pyrefly continues to check other annotations and statements in the function. + +```python +import triton +import triton.language as tl + + +@triton.jit +def next_block(block: tl.constexpr = 128) -> int: + return block + 1 +``` + +Functions without a recognized Triton decorator retain normal Python annotation +rules. A different decorator named `jit` or a different class named `constexpr` +does not enable this behavior. Nested functions require their own Triton decorator. + +This support does not infer constants from kernel calls. It does not yet check +Triton pointer dtypes, tensor shapes, broadcasting, or kernel launch arguments. +It does not introduce a `strict=True` argument or a separate plugin interface. +These features remain outside this first step for +[issue #3884](https://github.com/facebook/pyrefly/issues/3884). + +Run the regression tests without Triton or a GPU: + +```bash +cargo test -p pyrefly --lib test_triton_constexpr +``` diff --git a/pyrefly/lib/alt/solve.rs b/pyrefly/lib/alt/solve.rs index dcb9c0ebb1..31fe4e89d3 100644 --- a/pyrefly/lib/alt/solve.rs +++ b/pyrefly/lib/alt/solve.rs @@ -680,7 +680,13 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { match binding { BindingAnnotation::AnnotateExpr(target, x, class_key) => { let type_form_context = target.type_form_context(); - let mut ann = self.expr_annotation(x, type_form_context, errors); + let mut ann = if matches!(target, AnnotationTarget::TritonConstexprParam(_)) { + // Triton supplies the compile-time value, which need not be a constexpr object. + self.expr_infer(x, errors); + Annotation::new_type(self.heap.mk_any_implicit()) + } else { + self.expr_annotation(x, type_form_context, errors) + }; if let Some(class_key) = class_key && let Some(ty) = &mut ann.ty { diff --git a/pyrefly/lib/binding/binding.rs b/pyrefly/lib/binding/binding.rs index b4e8bf7b30..9bf5e371f6 100644 --- a/pyrefly/lib/binding/binding.rs +++ b/pyrefly/lib/binding/binding.rs @@ -3423,6 +3423,8 @@ impl Display for AnnotationWithTarget { pub enum AnnotationTarget { /// A function parameter with a type annotation Param(Name), + /// A compile-time parameter of a Triton JIT function. + TritonConstexprParam(Name), ArgsParam(Name), KwargsParam(Name), /// A return type annotation on a function. The name is that of the function @@ -3439,7 +3441,7 @@ pub enum AnnotationTarget { impl Display for AnnotationTarget { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match self { - Self::Param(name) => write!(f, "parameter `{name}`"), + Self::Param(name) | Self::TritonConstexprParam(name) => write!(f, "parameter `{name}`"), Self::ArgsParam(name) => write!(f, "args `{name}`"), Self::KwargsParam(name) => write!(f, "kwargs `{name}`"), Self::Return(name) => write!(f, "`{name}` return"), @@ -3453,7 +3455,7 @@ impl Display for AnnotationTarget { impl AnnotationTarget { pub fn type_form_context(&self) -> TypeFormContext<'static> { match self { - Self::Param(_) => TypeFormContext::ParameterAnnotation, + Self::Param(_) | Self::TritonConstexprParam(_) => TypeFormContext::ParameterAnnotation, Self::ArgsParam(_) => TypeFormContext::ParameterArgsAnnotation, Self::KwargsParam(_) => TypeFormContext::ParameterKwargsAnnotation, Self::Return(_) => TypeFormContext::ReturnAnnotation, diff --git a/pyrefly/lib/binding/function.rs b/pyrefly/lib/binding/function.rs index cda4685dbf..75c843d944 100644 --- a/pyrefly/lib/binding/function.rs +++ b/pyrefly/lib/binding/function.rs @@ -39,6 +39,7 @@ use ruff_python_ast::visitor::source_order::walk_stmt; use ruff_text_size::Ranged; use ruff_text_size::TextRange; use starlark_map::small_map::SmallMap; +use starlark_map::small_set::SmallSet; use thin_vec::ThinVec; use crate::binding::binding::AnnotationTarget; @@ -91,6 +92,7 @@ struct Decorators { is_override: bool, is_classmethod: bool, decorator_kind: FunctionDecorators, + is_triton_jit: bool, decorators: Box<[Idx]>, } @@ -291,6 +293,7 @@ impl<'a> BindingsBuilder<'a> { class_key: Option>, method_self_kind: MethodSelfKind, ignore_annotations: bool, + triton_constexpr_parameters: &SmallSet, ) { let mut self_name = None; for x in x.iter_non_variadic_params() { @@ -298,7 +301,11 @@ impl<'a> BindingsBuilder<'a> { self_name = Some(x.parameter.name.clone()); } self.bind_function_param( - AnnotationTarget::Param(x.parameter.name.id.clone()), + if triton_constexpr_parameters.contains(&x.parameter.name.id) { + AnnotationTarget::TritonConstexprParam(x.parameter.name.id.clone()) + } else { + AnnotationTarget::Param(x.parameter.name.id.clone()) + }, &x.parameter, undecorated_idx, class_key, @@ -402,6 +409,7 @@ impl<'a> BindingsBuilder<'a> { class_key: Option>, is_async: bool, method_self_kind: MethodSelfKind, + triton_constexpr_parameters: &SmallSet, ) -> ( YieldsAndReturns, Option, @@ -416,6 +424,7 @@ impl<'a> BindingsBuilder<'a> { class_key, method_self_kind, false, + triton_constexpr_parameters, ); self.init_static_scope(&body, false); self.seed_captured_variables(); @@ -460,6 +469,7 @@ impl<'a> BindingsBuilder<'a> { is_async: bool, method_self_kind: MethodSelfKind, ignore_annotations: bool, + triton_constexpr_parameters: &SmallSet, ) -> Option { // Push a scope to create the parameter keys (but do nothing else with it). self.scopes @@ -470,6 +480,7 @@ impl<'a> BindingsBuilder<'a> { class_key, method_self_kind, ignore_annotations, + triton_constexpr_parameters, ); self.scopes.pop(); // If we are in a class, use a simple visitor to find `self.` assignments. @@ -664,6 +675,7 @@ impl<'a> BindingsBuilder<'a> { let mut is_classmethod = false; let mut decorator_kind = FunctionDecorators::Undecorated; let mut all_receiver_overload = true; + let mut is_triton_jit = false; for d in &decorator_list { let special_export = self.as_special_export(&d.expression); let receiver_special = if self.track_classmethod_uses { @@ -677,6 +689,11 @@ impl<'a> BindingsBuilder<'a> { } else { FunctionDecorators::Other }; + let decorator_export = match &d.expression { + Expr::Call(call) => self.as_special_export(&call.func), + _ => special_export, + }; + is_triton_jit = is_triton_jit || decorator_export == Some(SpecialExport::TritonJit); is_overload = is_overload || matches!(special_export, Some(SpecialExport::Overload)); is_override = is_override || matches!(special_export, Some(SpecialExport::Override)); is_abstract_method = is_abstract_method @@ -709,6 +726,7 @@ impl<'a> BindingsBuilder<'a> { is_override, is_classmethod, decorator_kind, + is_triton_jit, decorators, } } @@ -728,6 +746,20 @@ impl<'a> BindingsBuilder<'a> { undecorated_idx: Idx, class_key: Option>, ) -> (BodyKind, bool, bool, Option) { + // Resolve markers before parameter names can shadow imports from Triton. + let triton_constexpr_parameters = if decorators.is_triton_jit { + parameters + .iter_non_variadic_params() + .filter(|p| { + p.annotation().is_some_and(|annotation| { + self.as_special_export(annotation) == Some(SpecialExport::TritonConstexpr) + }) + }) + .map(|p| p.parameter.name.id.clone()) + .collect::>() + } else { + SmallSet::new() + }; // If the first statement in the body is a docstring, remove it let body_no_docstring = if let Some(s) = body.first() && is_docstring(s) @@ -805,6 +837,7 @@ impl<'a> BindingsBuilder<'a> { is_async, method_self_kind, decorators.has_no_type_check, + &triton_constexpr_parameters, ); (false, self_assignments) } else if is_unannotated { @@ -819,6 +852,7 @@ impl<'a> BindingsBuilder<'a> { class_key, is_async, method_self_kind, + &triton_constexpr_parameters, ); self.analyze_return_type( func_name, @@ -846,6 +880,7 @@ impl<'a> BindingsBuilder<'a> { class_key, is_async, method_self_kind, + &triton_constexpr_parameters, ); if !ignore_unused_parameters { self.record_unused_parameters(unused_parameters); diff --git a/pyrefly/lib/error/context.rs b/pyrefly/lib/error/context.rs index 38f4da018f..c1bc249f03 100644 --- a/pyrefly/lib/error/context.rs +++ b/pyrefly/lib/error/context.rs @@ -226,6 +226,7 @@ impl TypeCheckKind { pub fn from_annotation_target(target: &AnnotationTarget) -> Self { match target { AnnotationTarget::Param(name) + | AnnotationTarget::TritonConstexprParam(name) | AnnotationTarget::ArgsParam(name) | AnnotationTarget::KwargsParam(name) => Self::CallArgument(Some(name.clone()), None), AnnotationTarget::Return(_func) => Self::ExplicitFunctionReturn, diff --git a/pyrefly/lib/export/special.rs b/pyrefly/lib/export/special.rs index e9917dfa9a..45f5017082 100644 --- a/pyrefly/lib/export/special.rs +++ b/pyrefly/lib/export/special.rs @@ -99,6 +99,8 @@ pub enum SpecialExport { AttrsNothing, FunctoolsWraps, FunctoolsUpdateWrapper, + TritonJit, + TritonConstexpr, } impl SpecialExport { @@ -189,6 +191,8 @@ impl SpecialExport { "NOTHING" => Some(Self::AttrsNothing), "wraps" => Some(Self::FunctoolsWraps), "update_wrapper" => Some(Self::FunctoolsUpdateWrapper), + "jit" => Some(Self::TritonJit), + "constexpr" => Some(Self::TritonConstexpr), _ => None, } } @@ -289,6 +293,13 @@ impl SpecialExport { Self::AttrsLegacyAttrib | Self::AttrsNextGenField | Self::AttrsNothing => { matches!(m.as_str(), "attr" | "attrs") } + Self::TritonJit => matches!( + m.as_str(), + "triton" | "triton.runtime" | "triton.runtime.jit" + ), + Self::TritonConstexpr => { + matches!(m.as_str(), "triton.language" | "triton.language.core") + } } } diff --git a/pyrefly/lib/test/mod.rs b/pyrefly/lib/test/mod.rs index 8fa5c2cfaf..a20375a5fb 100644 --- a/pyrefly/lib/test/mod.rs +++ b/pyrefly/lib/test/mod.rs @@ -94,6 +94,7 @@ mod state; mod subscript_narrow; mod suppression; mod sys_info; +mod triton; mod tsp; mod tuple; mod type_alias; diff --git a/pyrefly/lib/test/triton.rs b/pyrefly/lib/test/triton.rs new file mode 100644 index 0000000000..c0175e4d2d --- /dev/null +++ b/pyrefly/lib/test/triton.rs @@ -0,0 +1,256 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +use crate::test::util::TestEnv; +use crate::testcase; + +fn env() -> TestEnv { + let mut env = TestEnv::new(); + env.add_with_path( + "triton.runtime.jit", + "triton/runtime/jit.pyi", + r#" +from typing import Callable, overload + +class _JitDecorator: + def __call__[**P, R](self, fn: Callable[P, R]) -> Callable[P, R]: ... + +@overload +def jit[**P, R](fn: Callable[P, R]) -> Callable[P, R]: ... +@overload +def jit(*, debug: bool = False) -> _JitDecorator: ... +"#, + ); + env.add("triton", "from triton.runtime.jit import jit as jit"); + env.add("triton.language.core", "class constexpr: ..."); + env.add( + "triton.language", + "from triton.language.core import constexpr as constexpr", + ); + env.add( + "kernel_helpers", + r#" +from triton import jit as compile_kernel +from triton.language import constexpr as constant +"#, + ); + env +} + +testcase!( + test_triton_constexpr_parameter, + env(), + r#" +from typing import Any, assert_type +import triton +import triton.language as tl + +@triton.jit +def kernel(block: tl.constexpr) -> int: + assert_type(block, Any) + return block + 1 + +assert_type(kernel(128), int) +"#, +); + +testcase!( + test_triton_constexpr_decorator_factory, + env(), + r#" +from typing import Any, assert_type +import triton +import triton.language as tl + +@triton.jit(debug=True) +def kernel(flag: tl.constexpr = True) -> bool: + assert_type(flag, Any) + return flag + +assert_type(kernel(), bool) +assert_type(kernel(False), bool) +"#, +); + +testcase!( + test_triton_constexpr_aliases, + env(), + r#" +from typing import Any, assert_type +from triton import jit as compile_kernel +from triton.language import constexpr as constant + +@compile_kernel +def kernel(value: constant) -> str: + assert_type(value, Any) + return value + +assert_type(kernel("mode"), str) +"#, +); + +testcase!( + test_triton_constexpr_definition_modules, + env(), + r#" +from typing import Any, assert_type +from triton.runtime.jit import jit +from triton.language.core import constexpr + +@jit +def kernel(shape: constexpr) -> None: + assert_type(shape, Any) + +kernel((16, 32)) +"#, +); + +testcase!( + test_triton_constexpr_parameter_kinds, + env(), + r#" +from typing import Any, assert_type +from triton import jit +from triton.language import constexpr + +@jit +def kernel(block: "constexpr", /, *, mode: constexpr = "fast") -> None: + assert_type(block, Any) + assert_type(mode, Any) + +kernel(128) +kernel(64, mode="slow") +"#, +); + +testcase!( + test_triton_constexpr_reexports, + env(), + r#" +from typing import Any, assert_type +from kernel_helpers import compile_kernel, constant + +@compile_kernel(debug=True) +def kernel(value: constant) -> None: + assert_type(value, Any) + +kernel(128) +kernel() # E: Missing argument `value` +"#, +); + +testcase!( + test_triton_constexpr_parameter_shadows_import, + env(), + r#" +from typing import Any, assert_type +from triton import jit +from triton.language import constexpr + +@jit +def kernel(constexpr: int, block: constexpr) -> None: + assert_type(constexpr, int) + assert_type(block, Any) + +kernel(1, 128) +"#, +); + +testcase!( + test_triton_constexpr_other_parameters_checked, + env(), + r#" +from triton import jit +from triton.language import constexpr + +@jit +def kernel(block: constexpr, count: int) -> int: + return count + +kernel(128, "bad") # E: not assignable to parameter `count` +"#, +); + +testcase!( + test_triton_constexpr_body_checked, + env(), + r#" +from triton import jit +from triton.language import constexpr + +@jit +def kernel(block: constexpr) -> int: + count: int = "bad" # E: not assignable to `int` + return count +"#, +); + +testcase!( + test_triton_constexpr_without_jit, + env(), + r#" +from typing import assert_type +from triton.language import constexpr + +def ordinary(value: constexpr) -> None: + assert_type(value, constexpr) + +ordinary(128) # E: not assignable to parameter `value` +"#, +); + +testcase!( + test_triton_constexpr_unrelated_decorator, + env(), + r#" +from typing import assert_type +from triton.language import constexpr + +def jit[T](fn: T) -> T: + return fn + +@jit +def ordinary(value: constexpr) -> None: + assert_type(value, constexpr) + +ordinary(128) # E: not assignable to parameter `value` +"#, +); + +testcase!( + test_triton_constexpr_unrelated_annotation, + env(), + r#" +from typing import assert_type +from triton import jit + +class constexpr: ... + +@jit +def kernel(value: constexpr) -> None: + assert_type(value, constexpr) + +kernel(128) # E: not assignable to parameter `value` +"#, +); + +testcase!( + test_triton_constexpr_nested_function, + env(), + r#" +from typing import Any, assert_type +from triton import jit +from triton.language import constexpr + +@jit +def kernel(value: constexpr) -> None: + assert_type(value, Any) + def ordinary(inner: constexpr) -> None: + assert_type(inner, constexpr) + ordinary(128) # E: not assignable to parameter `inner` +"#, +);