Spaces:
Running
Running
File size: 8,868 Bytes
e313085 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | """AST checks that run before any candidate code executes.
Goal: reject submissions that compute with PyTorch instead of custom kernels, or that try to
tamper with the evaluator (imports, monkeypatching, introspection escapes).
Tensor arithmetic operators are checked at runtime in worker.kernel_only_mode, where tensors can
be distinguished from legitimate scalar/shape arithmetic.
"""
import ast
from typing import Dict, List, Optional, Tuple
ALLOWED_IMPORTS = {"torch", "triton", "math", "typing", "functools"}
BANNED_CALLS = {
"exec", "eval", "compile", "__import__", "getattr", "setattr", "delattr",
"globals", "locals", "vars", "open", "breakpoint", "input", "memoryview",
}
# torch.* calls allowed in host code (allocation only)
TORCH_ALLOWED = {
f"torch.{n}"
for n in (
"empty", "empty_like", "empty_strided", "zeros", "zeros_like", "ones", "ones_like",
"full", "full_like", "arange", "Size", "device", "is_tensor",
# random generation: harmless, and models often copy the reference's get_inputs()
"rand", "randn", "randint", "rand_like", "randn_like", "manual_seed",
)
}
# tensor methods that do real compute
TENSOR_COMPUTE = {
"matmul", "mm", "bmm", "addmm", "baddbmm", "einsum", "sum", "mean", "max", "min", "amax",
"amin", "prod", "cumsum", "cumprod", "softmax", "log_softmax", "relu", "relu_", "sigmoid",
"tanh", "exp", "log", "sqrt", "rsqrt", "pow", "abs", "clamp", "clamp_", "norm", "var", "std",
"argmax", "argmin", "topk", "sort", "mul", "mul_", "add", "add_", "sub", "sub_", "div",
"div_", "neg", "gelu", "silu", "conv1d", "conv2d", "conv3d", "matrix_power", "tril", "triu",
"masked_fill", "where", "gather", "scatter", "index_select", "logsumexp", "outer", "dot",
}
KERNEL_DECORATORS = {"triton.jit", "triton.autotune", "triton.heuristics"}
def _path(node: ast.AST, aliases: Dict[str, str]) -> Tuple[Optional[str], bool]:
"""Dotted path of a Name/Attribute chain with import aliases expanded; bool = root is an import."""
parts = []
while isinstance(node, ast.Attribute):
parts.append(node.attr)
node = node.value
if not isinstance(node, ast.Name):
return None, False
root = aliases.get(node.id)
return ".".join([root or node.id] + parts[::-1]), root is not None
def _is_self_attr(node: ast.AST) -> bool:
"""self.<name>: a user-chosen attribute name (e.g. self.conv3d), not a tensor method."""
return isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name) and node.value.id == "self"
def _is_kernel(fn: ast.FunctionDef, aliases: Dict[str, str]) -> bool:
for d in fn.decorator_list:
p, _ = _path(d.func if isinstance(d, ast.Call) else d, aliases)
if p in KERNEL_DECORATORS:
return True
return False
def check_source(code: str, language: str = "triton") -> List[str]:
"""Return a list of violations; empty list means the source may be executed."""
try:
tree = ast.parse(code)
except SyntaxError as e:
return [f"SyntaxError: {e.msg} (line {e.lineno})"]
errors: List[str] = []
aliases: Dict[str, str] = {}
for node in ast.walk(tree):
if isinstance(node, ast.Import):
for a in node.names:
if a.name.split(".")[0] not in ALLOWED_IMPORTS:
errors.append(f"import of '{a.name}' not allowed")
aliases[a.asname or a.name.split(".")[0]] = a.name if a.asname else a.name.split(".")[0]
elif isinstance(node, ast.ImportFrom):
mod = node.module or ""
if node.level or mod.split(".")[0] not in ALLOWED_IMPORTS:
errors.append(f"import from '{mod}' not allowed")
for a in node.names:
aliases[a.asname or a.name] = f"{mod}.{a.name}"
# module-level `ext = load_inline(...)`: ext.relu(x) is a custom kernel, not a tensor op
for node in tree.body:
if isinstance(node, ast.Assign) and isinstance(node.value, ast.Call):
if _path(node.value.func, aliases)[0] == "torch.utils.cpp_extension.load_inline":
aliases.update({t.id: f"cuda_ext.{t.id}" for t in node.targets if isinstance(t, ast.Name)})
for node in ast.walk(tree):
if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id in BANNED_CALLS:
errors.append(f"call to '{node.func.id}' not allowed (line {node.lineno})")
elif isinstance(node, ast.Attribute) and node.attr.startswith("__") and node.attr != "__init__":
errors.append(f"dunder attribute '{node.attr}' not allowed (line {node.lineno})")
elif isinstance(node, (ast.Assign, ast.AugAssign, ast.AnnAssign, ast.Delete)):
targets = node.targets if isinstance(node, (ast.Assign, ast.Delete)) else [node.target]
for t in targets:
for sub in ast.walk(t):
if isinstance(sub, ast.Attribute):
root = sub
while isinstance(root, ast.Attribute):
root = root.value
if not (isinstance(root, ast.Name) and root.id == "self"):
errors.append(f"assigning attributes outside self not allowed (line {node.lineno})")
classes = {n.name: n for n in tree.body if isinstance(n, ast.ClassDef)}
if "ModelNew" not in classes:
errors.append("no top-level class ModelNew")
# self attributes holding nn layers (calling them = PyTorch compute)
nn_attrs = set()
for cls in classes.values():
for fn in cls.body:
if isinstance(fn, ast.FunctionDef) and fn.name == "__init__":
for n in ast.walk(fn):
if isinstance(n, ast.Assign) and isinstance(n.value, ast.Call):
p, _ = _path(n.value.func, aliases)
if p and p.startswith("torch.nn.") and p != "torch.nn.Parameter":
nn_attrs.update(
t.attr for t in n.targets
if isinstance(t, ast.Attribute) and isinstance(t.value, ast.Name) and t.value.id == "self"
)
funcs = [n for n in ast.walk(tree) if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef))]
kernels = [f for f in funcs if _is_kernel(f, aliases)]
kernel_ids = {id(n) for k in kernels for n in ast.walk(k)}
init_ids = {id(n) for f in funcs if f.name == "__init__" for n in ast.walk(f)}
call_func_ids = {id(n.func) for n in ast.walk(tree) if isinstance(n, ast.Call)}
for n in ast.walk(tree):
if id(n) in kernel_ids:
continue
# references that smuggle compute through an alias, e.g. `mm = torch.matmul; mm(a, b)`
if isinstance(n, (ast.Name, ast.Attribute)) and id(n) not in call_func_ids:
p, imp = _path(n, aliases)
if p and imp and p.startswith("torch.") and (
p.rsplit(".", 1)[-1] in TENSOR_COMPUTE
or p.startswith(("torch.nn.functional.", "torch.ops", "torch.compile", "torch._"))
):
errors.append(f"reference to PyTorch op '{p}' (line {n.lineno})")
elif not imp and isinstance(n, ast.Attribute) and n.attr in TENSOR_COMPUTE and not _is_self_attr(n):
errors.append(f"reference to tensor op '.{n.attr}' (line {n.lineno})")
if id(n) in init_ids:
continue
if isinstance(n, ast.BinOp) and isinstance(n.op, ast.MatMult):
errors.append(f"'@' matmul in host code (line {n.lineno})")
if not isinstance(n, ast.Call):
continue
p, root_is_import = _path(n.func, aliases)
if p is None:
continue
if p.startswith("torch.") and p not in TORCH_ALLOWED and not p.startswith(
("torch.cuda.", "torch.utils.cpp_extension.")
):
errors.append(f"PyTorch op '{p}' in host code (line {n.lineno}); use a custom kernel")
elif not root_is_import and isinstance(n.func, ast.Attribute):
if _is_self_attr(n.func):
if n.func.attr in nn_attrs:
errors.append(f"calling nn layer '{p}' in host code (line {n.lineno})")
elif n.func.attr in TENSOR_COMPUTE:
errors.append(f"tensor op '.{n.func.attr}()' in host code (line {n.lineno})")
if language == "triton" and not kernels:
errors.append("no @triton.jit kernel defined")
if language == "cuda":
loads = [
n for n in ast.walk(tree)
if isinstance(n, ast.Call) and _path(n.func, aliases)[0] == "torch.utils.cpp_extension.load_inline"
]
if not loads:
errors.append("no torch.utils.cpp_extension.load_inline call")
return list(dict.fromkeys(errors)) # dedupe, keep order
|