impl part of Feature: Native Triton GPU Kernel Strict Typing Plugin #3884

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
+```

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
                 {

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,

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<KeyDecorator>]>,
 }
 
@@ -291,14 +293,19 @@ impl<'a> BindingsBuilder<'a> {
         class_key: Option<Idx<KeyClass>>,
         method_self_kind: MethodSelfKind,
         ignore_annotations: bool,
+        triton_constexpr_parameters: &SmallSet<Name>,
     ) {
         let mut self_name = None;
         for x in x.iter_non_variadic_params() {
             if class_key.is_some() && self_name.is_none() {
                 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<Idx<KeyClass>>,
         is_async: bool,
         method_self_kind: MethodSelfKind,
+        triton_constexpr_parameters: &SmallSet<Name>,
     ) -> (
         YieldsAndReturns,
         Option<SelfAssignments>,
@@ -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<Name>,
     ) -> Option<SelfAssignments> {
         // 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.<attr>` 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<KeyUndecoratedFunction>,
         class_key: Option<Idx<KeyClass>>,
     ) -> (BodyKind, bool, bool, Option<SelfAssignments>) {
+        // 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::<SmallSet<_>>()
+        } 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);

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,

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")
+            }
         }
     }
 

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;