I have a function foo
that is part of a kernel. It’s called with a LayoutTensor
. I want to use a separate helper function bar
that works on scalars coming from the LayoutTensor
. How do I do that?
from layout.layout_tensor import (
Layout,
LayoutTensor,
)
from layout.layout import UNKNOWN_VALUE
from layout.tensor_builder import LayoutTensorBuild as tb
from layout.tensor_builder import dynamic, static
fn bar[dtype: DType](one: Scalar[dtype]):
print(one)
fn foo[dtype: DType, corners_layout: Layout](
corners: LayoutTensor[
dtype, corners_layout, MutableAnyOrigin
], # x1, y1, x2, y2
):
bar(corners[0,0])
def main():
alias N = 16
alias corners_layout = Layout.row_major(N, 4)
var t_2d_static = tb[DType.float32]().row_major(
static[16](), static[4]()
).alloc()
foo[DType.float32, corners_layout](t_2d_static)
The code above fails with
/Users/mseritan/dev/hackathon-may-25/nms.mojo/crash.mojo:17:8: error: invalid call to 'bar': argument #0 cannot be converted from 'SIMD[dtype, __init__[::Origin[::Bool(IntTuple(1), IntTuple(1)).size()]' to 'SIMD[dtype, 1]'
bar(corners[0,0])
~~~^~~~~~~~~~~~~~
/Users/mseritan/dev/hackathon-may-25/nms.mojo/crash.mojo:17:16: note: types parameters include unfolded expression at parser time; try rebinding to a consistent type?
bar(corners[0,0])
^
/Users/mseritan/dev/hackathon-may-25/nms.mojo/crash.mojo:9:4: note: function declared here
fn bar[dtype: DType](one: Scalar[dtype]):
^
mojo: error: failed to parse the provided Mojo source module