diff --git a/KLR/Trace/Builtin.lean b/KLR/Trace/Builtin.lean index 354c2275..404301f5 100644 --- a/KLR/Trace/Builtin.lean +++ b/KLR/Trace/Builtin.lean @@ -34,11 +34,13 @@ def nki_isa : Name := .str nki_ "isa" def nki_stdlib : Name := .str nki_ "stdlib" def nki_lang : Name := .str nki_ "language" def nki_typing : Name := .str nki_ "typing" -def nki_tensor : Name := .str nki_ "tensor" +def nki_meta : Name := .str nki_ "meta" +def nki_tensor : Name := .str nki_meta "tensor" def nl : String -> Name := .str nki_lang def nisa : String -> Name := .str nki_isa def nt : String -> Name := .str nki_typing +def nmeta : String -> Name := .str nki_meta def ntensor : String -> Name := .str nki_tensor def numpy : Name := .str .anonymous "numpy" diff --git a/KLR/Trace/Tensor.lean b/KLR/Trace/Tensor.lean index 4718fbff..aad14413 100644 --- a/KLR/Trace/Tensor.lean +++ b/KLR/Trace/Tensor.lean @@ -63,13 +63,13 @@ nki builtin.tensor.transpose (src : Access) := do warn "transpose is not supported" return .access src -nki builtin.tensor.zeros (shape: Shape) (dtype: Dtype) := do +nki builtin.meta.tensor.zeros (shape: Shape) (dtype: Dtype) := do let tlShape := TensorLib.Shape.mk shape.toList let tlDtype <- dtype.toTensorLibDtype let tensor := TensorLib.Tensor.zeros tlDtype tlShape return .tensor tensor -nki builtin.tensor.arange +nki builtin.meta.tensor.arange (start : Nat) (stop : Nat) (step : Nat := 1) @@ -84,7 +84,7 @@ nki builtin.tensor.arange data := data.append bytes return .tensor {dtype := tlDtype, shape := tlShape, data := data} -nki builtin.tensor.identity +nki builtin.meta.tensor.identity (N : Nat) (dtype : Dtype := .float32) := do let tlDtype <- dtype.toTensorLibDtype diff --git a/KLR/Trace/Term.lean b/KLR/Trace/Term.lean index 64d4eba5..a11cfdaa 100644 --- a/KLR/Trace/Term.lean +++ b/KLR/Trace/Term.lean @@ -716,6 +716,6 @@ def builtinEnv : List (Name × Term) := Id.run do | .str `builtin.isa n => [nisa n, name] | .str `builtin.typing n => [nt n, name] | .str `builtin.lang n => [nl n, name] - | .str `builtin.tensor n => [ntensor n, name] + | .str `builtin.meta.tensor n => [ntensor n, name] | _ => [name] names.map fun n => (n, fn)