Taichi: [Bug or feature] Can we access the shape of external array?

[archibate]

Backgrounds
Since ti.field must be declared before any kernel invocation. Making passing data dynamic size impossible.
So I'd like to extent its ability for more flexible data structure.
E.g., Taichi THREE could provide such API:

scene.triangles(arr1, arr2, arr3)  # where arr is numpy array of 3d vectors

# inside taichi_three:
@ti.kernel
def triangles(self, arr1: ti.ext_arr(), ...):  # template id is the (shape, dtype) pair
  ti.static_assert(arr1.shape[1] == 3)
  for i in range(arr1.shape[0]):
    ...

However, urrently the ti.ext_arr() is very edged and even buggy according this forum thread.
So I raise about the first step for it:

Concisely describe the proposed feature
I'd like to access the shape of an external array in Taichi-scope, e.g.:

import taichi as ti
import numpy as np

@ti.kernel
def my_kernel(x: ti.ext_arr()):
    for i in range(x.shape[0]):
        print(x[i])

x = np.array([2, 3, 4, 5])
my_kernel(x)

However, I obtain:

[Taichi] mode=development
[Taichi] preparing sandbox at /tmp/taichi-ybvnxt8j
[Taichi] <dev mode>, llvm 10.0.0, commit 2be05fc7, python 3.8.3
[Taichi] materializing...
[E 08/10/20 22:06:34.318] [expr.cpp:snode@83] Cannot get snode of non-global variables.


***********************************
* Taichi Compiler Stack Traceback *
***********************************
/tmp/taichi-ybvnxt8j/taichi_core.so: taichi::Logger::error(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, bool)
/tmp/taichi-ybvnxt8j/taichi_core.so: taichi::lang::Expr::snode() const
/tmp/taichi-ybvnxt8j/taichi_core.so(+0x4c481f) [0x7fc79aec381f]
/tmp/taichi-ybvnxt8j/taichi_core.so(+0x394272) [0x7fc79ad93272]
/usr/lib/libpython3.8.so.1.0: PyCFunction_Call
/usr/lib/libpython3.8.so.1.0: _PyObject_MakeTpCall
/usr/lib/libpython3.8.so.1.0(+0x13e75f) [0x7fc79ef1e75f]
/usr/lib/libpython3.8.so.1.0: _PyEval_EvalFrameDefault
/usr/lib/libpython3.8.so.1.0(+0x13a378) [0x7fc79ef1a378]
/usr/lib/libpython3.8.so.1.0: _PyObject_GenericGetAttrWithDict
/usr/lib/libpython3.8.so.1.0: _PyEval_EvalFrameDefault
/usr/lib/libpython3.8.so.1.0(+0x13a378) [0x7fc79ef1a378]
/usr/lib/libpython3.8.so.1.0: _PyObject_GenericGetAttrWithDict
/usr/lib/libpython3.8.so.1.0: _PyEval_EvalFrameDefault
/usr/lib/libpython3.8.so.1.0: _PyFunction_Vectorcall
/usr/lib/libpython3.8.so.1.0: _PyEval_EvalFrameDefault
/usr/lib/libpython3.8.so.1.0: _PyEval_EvalCodeWithName
/usr/lib/libpython3.8.so.1.0: _PyFunction_Vectorcall
/usr/lib/libpython3.8.so.1.0: PyObject_Call
/tmp/taichi-ybvnxt8j/taichi_core.so(+0x4c6175) [0x7fc79aec5175]
/tmp/taichi-ybvnxt8j/taichi_core.so(+0x4c610f) [0x7fc79aec510f]
/tmp/taichi-ybvnxt8j/taichi_core.so: taichi::lang::Kernel::Kernel(taichi::lang::Program&, std::function<void ()>, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, bool)
/tmp/taichi-ybvnxt8j/taichi_core.so(+0x447f2c) [0x7fc79ae46f2c]
/tmp/taichi-ybvnxt8j/taichi_core.so: taichi::lang::Program::kernel(std::function<void ()> const&, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, bool)
/tmp/taichi-ybvnxt8j/taichi_core.so(+0x456f92) [0x7fc79ae55f92]
/tmp/taichi-ybvnxt8j/taichi_core.so(+0x394272) [0x7fc79ad93272]
/usr/lib/libpython3.8.so.1.0: PyCFunction_Call
/usr/lib/libpython3.8.so.1.0: _PyObject_MakeTpCall
/usr/lib/libpython3.8.so.1.0(+0x13e75f) [0x7fc79ef1e75f]
/usr/lib/libpython3.8.so.1.0: _PyEval_EvalFrameDefault
/usr/lib/libpython3.8.so.1.0: _PyEval_EvalCodeWithName
/usr/lib/libpython3.8.so.1.0(+0x13e442) [0x7fc79ef1e442]
/usr/lib/libpython3.8.so.1.0: _PyEval_EvalFrameDefault
/usr/lib/libpython3.8.so.1.0: _PyEval_EvalCodeWithName
/usr/lib/libpython3.8.so.1.0: _PyFunction_Vectorcall
/usr/lib/libpython3.8.so.1.0: _PyObject_FastCallDict
/usr/lib/libpython3.8.so.1.0: _PyObject_Call_Prepend
/usr/lib/libpython3.8.so.1.0(+0x1f5e09) [0x7fc79efd5e09]
/usr/lib/libpython3.8.so.1.0: PyObject_Call
/usr/lib/libpython3.8.so.1.0: _PyEval_EvalFrameDefault
/usr/lib/libpython3.8.so.1.0: _PyEval_EvalCodeWithName
/usr/lib/libpython3.8.so.1.0: _PyFunction_Vectorcall
/usr/lib/libpython3.8.so.1.0: _PyEval_EvalFrameDefault
/usr/lib/libpython3.8.so.1.0: _PyEval_EvalCodeWithName
/usr/lib/libpython3.8.so.1.0: PyEval_EvalCode
/usr/lib/libpython3.8.so.1.0(+0x1d8248) [0x7fc79efb8248]
/usr/lib/libpython3.8.so.1.0(+0x1d2483) [0x7fc79efb2483]
/usr/lib/libpython3.8.so.1.0: PyRun_FileExFlags
/usr/lib/libpython3.8.so.1.0: PyRun_SimpleFileExFlags
/usr/lib/libpython3.8.so.1.0: Py_RunMain
/usr/lib/libpython3.8.so.1.0: Py_BytesMain
/usr/lib/libc.so.6: __libc_start_main
python(_start+0x2e) [0x561f5791804e]

Internal Error occurred, check this page for possible solutions:
https://taichi.readthedocs.io/en/stable/install.html#troubleshooting
Traceback (most recent call last):
  File "a.py", line 10, in <module>
    my_kernel(x)
  File "/home/bate/Develop/taichi/python/taichi/lang/kernel.py", line 571, in wrapped
    return primal(*args, **kwargs)
  File "/home/bate/Develop/taichi/python/taichi/lang/kernel.py", line 498, in __call__
    self.materialize(key=key, args=args, arg_features=arg_features)
  File "/home/bate/Develop/taichi/python/taichi/lang/kernel.py", line 374, in materialize
    taichi_kernel = taichi_kernel.define(taichi_ast_generator)
  File "/home/bate/Develop/taichi/python/taichi/lang/kernel.py", line 371, in taichi_ast_generator
    compiled()
  File "a.py", line 6, in my_kernel
    for i in range(x.shape[0]):
  File "/home/bate/Develop/taichi/python/taichi/lang/expr.py", line 140, in shape
    return self.snode.shape
  File "/home/bate/Develop/taichi/python/taichi/lang/expr.py", line 133, in snode
    return SNode(self.ptr.snode())
RuntimeError: [expr.cpp:snode@83] Cannot get snode of non-global variables.

instead.

Describe the solution you'd like (if any)
It seems that in Taichi-scope, the x is not of type np.array, but of type ti.Expr? So I suggest:

1. Let x: ti.ext_arr() to be simply a np.array.
2. Do mocks in ti.subscript(x, indices) when x is detected to be np.array, do the original stuffs.
   So accessing shape or dtype would be strightfoward.

Additional comments

[thinking-tower]

@archibate Hey I'd be keen to work on this!

[archibate]

@archibate Hey I'd be keen to work on this!

Great! What's your plan (1 or 2)? Willing to provide more information if you'd like to help :)

[thinking-tower]

@archibate Here's what I'm planning to do:

• [x] Compile from source
• [x] Play a bit with the framework to get a feel
• [x] Make a decision on 1/2

I am however curious what would your preference be and why?

[yuanming-hu]

Thanks for proposing this. You may need to extend the IR to include something like an ExtArrShapeStmt to query their shapes, given the shapes are not compile-time constant :-)

[archibate]

ExtArrShapeStmt

In fact, Taichi stores external array shapes in ctx.extra_args.

For example:

def f(a: ti.ext_arr(), b: ti.ext_arr(), c: ti.ext_arr()):

Then:
ctx.extra_args[0] is storing the value a.shape[0].
ctx.extra_args[1] is storing the value a.shape[1].
...
ctx.extra_args[8 + 0] is storing the value b.shape[0].
ctx.extra_args[8 + 1] is storing the value b.shape[1].
...
ctx.extra_args[16 + 0] is storing the value c.shape[0].
ctx.extra_args[16 + 1] is storing the value c.shape[1].
...

So we may implement ExtArrShapeStmt to be accessing the extra_args in each backends.

[thinking-tower]

Hey @archibate, would it be appropriate to bring "type information" into the ASTTransformerPreprocessor or ASTTransformerChecks? As @yuanming-hu mentioned, we could provide ExtArrShapeStmt to query the shapes and I was hoping to transform x.shape into ti.get_external_array_shape(x) in compile time if x is a ti.ext_arr().

[archibate]

Nice thought! But sadly we can't easily detect the "type of a variable" at ASTTransformer.
But don't worry, ti.ext_arr() is already translated into ti.Expr when called in!
So the x.shape is actually captured as a @property function:
https://github.com/taichi-dev/taichi/blob/a402f76536dd1fddd77b3d7e3e4eefc5819fd087/python/taichi/lang/expr.py#L145-L147
We may add if self.is_external_array(): return ti.get_external_array_shape(self) here.

[thinking-tower]

@archibate the furthest I got was adding a visit_Attribute(self, node) but there doesn't seem to be any existing mechanism that records 'types'.

We may add if self.is_external_array(): return ti.get_external_array_shape(self) here.

I'll use your suggested approach then 😄

[thinking-tower]

Hey @archibate, I'm currently stuck and here's my progress so far.

I have a couple of questions which I hope you can clarify:

1. How do you "return" an array from the LLVM codegen? Here I imitate how snode https://github.com/taichi-dev/taichi/blob/f1bde29bd7ba32712b59426c19b6877cf5e7f0f8/python/taichi/lang/snode.py#L81-L91 gets its shape because I don't know how to "return" an array from LLVM codegen.

1. Using the code in the my current branch, how would I resolve the error ValueError: Bad constant scalar expression: <class 'taichi.lang.expr.Expr.shape.<locals>.callable_tuple'> I'm getting from this kernel?

import numpy as np
import taichi as ti

@ti.kernel
def my_kernel(x: ti.ext_arr()):
    print(x.shape)

x = np.array([2, 3, 4, 5])
my_kernel(x)

[archibate]

Here I imitate how snode

Thanks, this part LGTM!

How do you "return" an array from the LLVM codegen?

What do you mean by "return"? We can't assume all ti.ext_arr() to be np.ndarray, people may pass torch.Tensor as external arrays too. So only the address of that array is available, that is stored in (void *)ctx->args[arg_id].

print(x.shape)

Sorry, please use ti.static_print(x.shape) instead. print is only able to print runtime variables for now.

[thinking-tower]

@archibate

What do you mean by "return"? We can't assume all ti.ext_arr() to be np.ndarray, people may pass torch.Tensor as external arrays too. So only the address of that array is available, that is stored in (void *)ctx->args[arg_id].

I was hoping to loop through the ctx->extra_args[arg_id][i] for 0 < i < dim and return an array constructed from that. Is that possible from the LLVM IR codegen?

Sorry, please use ti.static_print(x.shape) instead. print is only able to print runtime variables for now.

I ran ti.static_print(x.shape) and I got [Expr(...)]. I suspect it's something to do with how I'm 'returning'

llvm_val[stmt] = builder->CreateCall(
      get_runtime_function("Context_get_extra_args"),
      {get_context(), tlctx->get_constant(arg_id), tlctx->get_constant(axis)});

here

[archibate]

Is that possible from the LLVM IR codegen?

Even if it's possible for LLVM, it probably won't work in other backends like OpenGL or Metal.
Let's keep using the get_external_tensor_shape_along_axis approach for getting shape along axes one-by-one.

I ran ti.static_print(x.shape) and I got [Expr(...)]. I suspect it's something to do with how I'm 'returning'

Oh I see, x.shape is not constant at compile-time, that's expected.

Does applying this and print(a.shape) solve the issue?

[thinking-tower]

@archibate after applying the suggested fix and calling print(a.shape) instead of ti.static_print(a.shape) I get

1. ValueError: Bad constant scalar expression: <class 'list'> (taichi/python/taichi/lang/impl.py", line 273, in make_constant_expr)
2. self.ptr = impl.make_constant_expr(arg).ptr (taichi/python/taichi/lang/expr.py", line 33, in __init__)
3. return Expr(var).ptr (taichi/python/taichi/lang/impl.py", line 408, in entry2content)
4. contentries = [entry2content(entry) for entry in entries] (taichi/python/taichi/lang/impl.py", line 441, in ti_print)
...

Here's the final AST btw:

Final AST:
def my_kernel():
    import taichi as ti
    x = ti.decl_ext_arr_arg(ti.i64, 1)
    ti.ti_print(x.shape)

[archibate]

Bad constant scalar expression:

Sorry for the trouble. What about print(x.shape[0])?

[thinking-tower]

@archibate No worries! I enjoy learning and having you help me debug this is awesome!

Sorry for the trouble. What about print(x.shape[0])?

I've made some progress after following your suggestion. As it turns out, I needed to add the return type to the statement here otherwise https://github.com/taichi-dev/taichi/blob/a402f76536dd1fddd77b3d7e3e4eefc5819fd087/taichi/codegen/codegen_llvm.cpp#L646-L647
throws an error.

In particular, these two kernels work:

@ti.kernel
def my_kernel(x: ti.ext_arr()):
    print(x.shape[0]) # 2
    print(x.shape[1]) # 4

x = np.array([[2, 3, 4, 5], [6, 7, 8, 9]])
my_kernel(x)

@ti.kernel
def my_kernel(x: ti.ext_arr()):
    for i in range(x.shape[0]):
        print(x[i]) # 2 3 4 5

x = np.array([2, 3, 4, 5])
my_kernel(x)

But these two won't:

@ti.kernel
def my_kernel(x: ti.ext_arr()):
    for i in range(x.shape[0]):
        print(x[i]) # IndexError: Field with dim 2 accessed with indices of dim 1

x = np.array([[2, 3, 4, 5]])
my_kernel(x)

@ti.kernel
def my_kernel(x: ti.ext_arr()):
    print(x.shape) # ValueError: Bad constant scalar expression: <class 'list'>

x = np.array([[2, 3, 4, 5]])
my_kernel(x)

[archibate]

Glad to know you're having fun!

@ti.kernel
def my_kernel(x: ti.ext_arr()):
    for i in range(x.shape[0]):
        print(x[i]) # IndexError: Field with dim 2 accessed with indices of dim 1

x = np.array([[2, 3, 4, 5]])
my_kernel(x)

This is expected, since x is a 2d array. You may expect x[0] to be x[0, :] therefore print(x[i]) yields [2, 3, 4, 5]. But currently index slicing is not supported in Taichi kernels for now (#777). Use print(ti.Vector([x[i, j] for j in range(x.shape[1])])) for explicitly slice them into a vector.

@ti.kernel
def my_kernel(x: ti.ext_arr()):
    print(x.shape) # ValueError: Bad constant scalar expression: <class 'list'>

x = np.array([[2, 3, 4, 5]])
my_kernel(x)

This is probably a feature request on print, which doesn't support printing list for now. It should be dealt in a separate issue.

[thinking-tower]

@archibate I'm currently writing some simple tests for this feature. Is it expected behaviour for this to fail?

```@ti.all_archs
def test_get_external_tensor_shape():

@ti.kernel
def func(x: ti.ext_arr(), index: int) -> ti.i32:
    return x.shape[index] 

size = np.random.randint(10, size=(8,))
x = np.random.uniform(size=size)
for idx, i in enumerate(size):
    assert func(x, idx) == i, "Axis {} should equal {}.".format(idx, i)

```

Error: TypeError: list indices must be integers or slices, not Expr

[archibate]

TypeError: list indices must be integers or slices, not Expr

It's expected, since list indices must be compile-time constant, however index is not constant but a taichi variable (Expr).

Solution, use ti.template() as type-hint:

def func(x: ti.ext_arr(), index: ti.template()) -> ti.i32:

This creates a new kernel instance for func on each different index, so that index will be compile-time constant.

[thinking-tower]

@archibate I can't seem to get this simple kernel to work on arch=ti.opengl:

# @ti.archs_excluding(ti.opengl)
@ti.all_archs
def test_get_external_tensor_shape_sum():

    @ti.kernel
    def func(x: ti.ext_arr()) -> ti.i32:
        y = 0
        for i in range(x.shape[0]):
            for j in range(x.shape[1]):
                y += x[i, j]
        return y

    size = np.random.randint(low=1, high=10, size=(2,))
    x = np.random.randint(low=0, high=np.prod(size), size=size)
    out = func(x) 
    assert out == x.sum(), "{} should equal output {}.".format(out, y)

# 0(6) : error C7531: global type int64_t requires "#extension GL_NV_gpu_shader5 : enable" before use
# 0(6) : error C0000: ... or #extension GL_ARB_gpu_shader_int64 : enable

I tried adding TI_OPENGL_REQUIRE(used, GL_ARB_gpu_shader_int64) here
https://github.com/taichi-dev/taichi/blob/c35172a6404aca9d69e00c6e6f9476ce79174abc/taichi/backends/opengl/codegen_opengl.cpp#L567-L572
but it didn't work.

I've pushed the my latest changes on to the same pull request here

[archibate]

x = np.random.randint(low=0, high=np.prod(size), size=size)

Does
x = np.random.randint(low=0, high=np.prod(size), size=size, dtype=np.int32)
solve the issue?

[thinking-tower]

Does
x = np.random.randint(low=0, high=np.prod(size), size=size, dtype=np.int32)
solve the issue?

Yes!