search

taichi-dev / taichi

Productive, portable, and performant GPU programming in Python.
https://taichi-lang.org
Apache License 2.0
25.38k stars 2.27k forks source link

[Bug] [AutoDiff] AD system failed to reverse_segments for a grouped ndrange-for #1751

Open archibate opened 4 years ago

archibate commented 4 years ago

Describe the bug Autodiff system failed to reverse_segments for a grouped ndrange-for. But using multi-index ndrange-for and group them later into a ti.Vector is OK..

To Reproduce

import taichi as ti

ti.init(print_preprocessed=True)

@ti.kernel
def func1():
    for i, j in ti.ndrange(2, 2):
        I = ti.Vector([i, j])
        print(I)

@ti.kernel
def func2():
    for I in ti.grouped(ti.ndrange(2, 2)):
        print(I)

func1.grad()  # OK
func2.grad()  # Error!

Log/Screenshots

[Taichi] mode=development
[Taichi] preparing sandbox at /tmp/taichi-z1ukj5ts
[Taichi] <dev mode>, llvm 10.0.0, commit f7eeaef6, linux, python 3.8.3
[Taichi] Starting on arch=x64
[Taichi] materializing...
Initial AST:
def func1():
    for i, j in ti.ndrange(2, 2):
        I = ti.Vector([i, j])
        print(I)

Preprocessed:
def func1():
    import taichi as ti
    if ti.static(1):
        __ndrange = ti.expr_init(ti.ndrange(2, 2))
        if 1:
            __ndrange_I = ti.Expr(ti.core.make_id_expr(''))
            ___begin = ti.Expr(0)
            ___end = ti.Expr(ti.subscript(__ndrange.acc_dimensions, 0))
            ___begin = ti.cast(___begin, ti.i32)
            ___end = ti.cast(___end, ti.i32)
            ti.core.begin_frontend_range_for(__ndrange_I.ptr, ___begin.ptr,
                ___end.ptr)
            __I = ti.expr_init(__ndrange_I)
            __i = ti.expr_init(__I // ti.subscript(__ndrange.acc_dimensions, 1)
                )
            i = ti.expr_init(__i + ti.subscript(ti.subscript(__ndrange.
                bounds, 0), 0))
            __I.assign(__I - __i * ti.subscript(__ndrange.acc_dimensions, 1))
            __j = ti.expr_init(__I)
            j = ti.expr_init(__j + ti.subscript(ti.subscript(__ndrange.
                bounds, 1), 0))
            I = ti.expr_init(ti.Vector([i, j]))
            ti.ti_print(I)
            del I
            del j
            del __j
            del i
            del __i
            del __I
            ti.core.end_frontend_range_for()
            del __ndrange_I
        del __ndrange

Checked:
def func1():
    import taichi as ti
    if ti.static(1):
        __ndrange = ti.expr_init(ti.ndrange(2, 2))
        if 1:
            __ndrange_I = ti.Expr(ti.core.make_id_expr(''))
            ___begin = ti.Expr(0)
            ___end = ti.Expr(ti.subscript(__ndrange.acc_dimensions, 0))
            ___begin = ti.cast(___begin, ti.i32)
            ___end = ti.cast(___end, ti.i32)
            ti.core.begin_frontend_range_for(__ndrange_I.ptr, ___begin.ptr,
                ___end.ptr)
            __I = ti.expr_init(__ndrange_I)
            __i = ti.expr_init(__I // ti.subscript(__ndrange.acc_dimensions, 1)
                )
            i = ti.expr_init(__i + ti.subscript(ti.subscript(__ndrange.
                bounds, 0), 0))
            __I.assign(__I - __i * ti.subscript(__ndrange.acc_dimensions, 1))
            __j = ti.expr_init(__I)
            j = ti.expr_init(__j + ti.subscript(ti.subscript(__ndrange.
                bounds, 1), 0))
            I = ti.expr_init(ti.Vector([i, j]))
            ti.ti_print(I)
            del I
            del j
            del __j
            del i
            del __i
            del __I
            ti.core.end_frontend_range_for()
            del __ndrange_I
        del __ndrange

Final AST:
def func1():
    import taichi as ti
    if ti.static(1):
        __ndrange = ti.expr_init(ti.ndrange(2, 2))
        if 1:
            __ndrange_I = ti.Expr(ti.core.make_id_expr(''))
            ___begin = ti.Expr(0)
            ___end = ti.Expr(ti.subscript(__ndrange.acc_dimensions, 0))
            ___begin = ti.cast(___begin, ti.i32)
            ___end = ti.cast(___end, ti.i32)
            ti.core.begin_frontend_range_for(__ndrange_I.ptr, ___begin.ptr,
                ___end.ptr)
            __I = ti.expr_init(__ndrange_I)
            __i = ti.expr_init(__I // ti.subscript(__ndrange.acc_dimensions, 1)
                )
            i = ti.expr_init(__i + ti.subscript(ti.subscript(__ndrange.
                bounds, 0), 0))
            __I.assign(__I - __i * ti.subscript(__ndrange.acc_dimensions, 1))
            __j = ti.expr_init(__I)
            j = ti.expr_init(__j + ti.subscript(ti.subscript(__ndrange.
                bounds, 1), 0))
            I = ti.expr_init(ti.Vector([i, j]))
            ti.ti_print(I)
            del I
            del j
            del __j
            del i
            del __i
            del __I
            ti.core.end_frontend_range_for()
            del __ndrange_I
        del __ndrange

[0, 0]
[0, 1]
[1, 0]
[1, 1]
Initial AST:
def func2():
    for I in ti.grouped(ti.ndrange(2, 2)):
        print(I)

Preprocessed:
def func2():
    import taichi as ti
    if ti.static(1):
        __ndrange = ti.ndrange(2, 2)
        I = ti.expr_init(ti.Vector([0] * len(__ndrange.dimensions)))
        ___begin = ti.Expr(0)
        ___end = __ndrange.acc_dimensions[0]
        ___begin = ti.cast(___begin, ti.i32)
        ___end = ti.cast(___end, ti.i32)
        __ndrange_I = ti.Expr(ti.core.make_id_expr(''))
        ti.core.begin_frontend_range_for(__ndrange_I.ptr, ___begin.ptr,
            ___end.ptr)
        __I = __ndrange_I
        for __grouped_I in range(len(__ndrange.dimensions)):
            __grouped_I_tmp = 0
            if __grouped_I + 1 < len(__ndrange.dimensions):
                __grouped_I_tmp = __I // __ndrange.acc_dimensions[
                    __grouped_I + 1]
            else:
                __grouped_I_tmp = __I
            ti.subscript(I, __grouped_I).assign(__grouped_I_tmp + __ndrange
                .bounds[__grouped_I][0])
            if __grouped_I + 1 < len(__ndrange.dimensions):
                __I = __I - __grouped_I_tmp * __ndrange.acc_dimensions[
                    __grouped_I + 1]
        ti.ti_print(I)
        ti.core.end_frontend_range_for()

Checked:
def func2():
    import taichi as ti
    if ti.static(1):
        __ndrange = ti.ndrange(2, 2)
        I = ti.expr_init(ti.Vector([0] * len(__ndrange.dimensions)))
        ___begin = ti.Expr(0)
        ___end = __ndrange.acc_dimensions[0]
        ___begin = ti.cast(___begin, ti.i32)
        ___end = ti.cast(___end, ti.i32)
        __ndrange_I = ti.Expr(ti.core.make_id_expr(''))
        ti.core.begin_frontend_range_for(__ndrange_I.ptr, ___begin.ptr,
            ___end.ptr)
        __I = __ndrange_I
        for __grouped_I in ti.func_call_with_check(range, len(__ndrange.
            dimensions)):
            __grouped_I_tmp = 0
            if __grouped_I + 1 < ti.func_call_with_check(len, __ndrange.
                dimensions):
                __grouped_I_tmp = __I // __ndrange.acc_dimensions[
                    __grouped_I + 1]
            else:
                __grouped_I_tmp = __I
            ti.subscript(I, __grouped_I).assign(__grouped_I_tmp + __ndrange
                .bounds[__grouped_I][0])
            if __grouped_I + 1 < ti.func_call_with_check(len, __ndrange.
                dimensions):
                __I = __I - __grouped_I_tmp * __ndrange.acc_dimensions[
                    __grouped_I + 1]
        ti.ti_print(I)
        ti.core.end_frontend_range_for()

Final AST:
def func2():
    import taichi as ti
    if ti.static(1):
        __ndrange = ti.ndrange(2, 2)
        I = ti.expr_init(ti.Vector([0] * len(__ndrange.dimensions)))
        ___begin = ti.Expr(0)
        ___end = __ndrange.acc_dimensions[0]
        ___begin = ti.cast(___begin, ti.i32)
        ___end = ti.cast(___end, ti.i32)
        __ndrange_I = ti.Expr(ti.core.make_id_expr(''))
        ti.core.begin_frontend_range_for(__ndrange_I.ptr, ___begin.ptr,
            ___end.ptr)
        __I = __ndrange_I
        for __grouped_I in ti.func_call_with_check(range, len(__ndrange.
            dimensions)):
            __grouped_I_tmp = 0
            if __grouped_I + 1 < ti.func_call_with_check(len, __ndrange.
                dimensions):
                __grouped_I_tmp = __I // __ndrange.acc_dimensions[
                    __grouped_I + 1]
            else:
                __grouped_I_tmp = __I
            ti.subscript(I, __grouped_I).assign(__grouped_I_tmp + __ndrange
                .bounds[__grouped_I][0])
            if __grouped_I + 1 < ti.func_call_with_check(len, __ndrange.
                dimensions):
                __I = __I - __grouped_I_tmp * __ndrange.acc_dimensions[
                    __grouped_I + 1]
        ti.ti_print(I)
        ti.core.end_frontend_range_for()

[E 08/23/20 09:54:28.265] [reverse_segments.cpp:reverse_segments@71] Invalid program input for autodiff. Please check the documentation for the "Kernel Simplicity Rule":
https://taichi.readthedocs.io/en/stable/autodiff.html#simplicity_rule

***********************************
* Taichi Compiler Stack Traceback *                                                          
***********************************                                                          
/tmp/taichi-z1ukj5ts/taichi_core.so: taichi::Logger::error(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, bool)                                  
/tmp/taichi-z1ukj5ts/taichi_core.so: taichi::lang::irpass::reverse_segments(taichi::lang::IRNode*)                                                                                        
/tmp/taichi-z1ukj5ts/taichi_core.so: taichi::lang::irpass::compile_to_offloads(taichi::lang::IRNode*, taichi::lang::CompileConfig const&, bool, bool, bool, bool)                         
/tmp/taichi-z1ukj5ts/taichi_core.so: taichi::lang::irpass::compile_to_executable(taichi::lang::IRNode*, taichi::lang::CompileConfig const&, bool, bool, bool, bool, bool, bool, bool)     
/tmp/taichi-z1ukj5ts/taichi_core.so: taichi::lang::Kernel::lower(bool)                       
/tmp/taichi-z1ukj5ts/taichi_core.so: taichi::lang::Program::compile(taichi::lang::Kernel&)   
/tmp/taichi-z1ukj5ts/taichi_core.so: taichi::lang::Kernel::compile()                         
/tmp/taichi-z1ukj5ts/taichi_core.so: taichi::lang::Kernel::operator()(taichi::lang::Kernel::LaunchContextBuilder&)                                                                        
/tmp/taichi-z1ukj5ts/taichi_core.so(+0x45ab14) [0x7f8ec1fcbb14]                              
/tmp/taichi-z1ukj5ts/taichi_core.so(+0x3943e2) [0x7f8ec1f053e2]                              
/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(+0x13e7e6) [0x7f8ec60997e6]                                     
/usr/lib/libpython3.8.so.1.0: PyObject_Call                                                  
/usr/lib/libpython3.8.so.1.0(+0x98f6d) [0x7f8ec5ff3f6d]                                      
/usr/lib/libpython3.8.so.1.0: _PyObject_MakeTpCall                                           
/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                                                  
/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: _PyObject_FastCallDict                                         
/usr/lib/libpython3.8.so.1.0: _PyObject_Call_Prepend
/usr/lib/libpython3.8.so.1.0(+0x1f5e09) [0x7f8ec6150e09]
/usr/lib/libpython3.8.so.1.0: _PyObject_MakeTpCall
/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) [0x7f8ec6133248]
/usr/lib/libpython3.8.so.1.0(+0x1d2483) [0x7f8ec612d483]
/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) [0x565183c3b04e]

Internal Error occurred, check this page for possible solutions:
https://taichi.readthedocs.io/en/stable/install.html#troubleshooting
Traceback (most recent call last):
  File "e.py", line 20, in <module>
    func2.grad()  # Error!
  File "/home/bate/Develop/taichi/python/taichi/lang/kernel.py", line 508, in __call__
    return self.compiled_functions[key](*args)
  File "/home/bate/Develop/taichi/python/taichi/lang/kernel.py", line 468, in func__
    t_kernel(launch_ctx)
RuntimeError: [reverse_segments.cpp:reverse_segments@71] Invalid program input for autodiff. Please check the documentation for the "Kernel Simplicity Rule":
https://taichi.readthedocs.io/en/stable/autodiff.html#simplicity_rule

If you have local commits (e.g. compile fixes before you reproduce the bug), please make sure you first make a PR to fix the build errors and then report the bug.

archibate commented 4 years ago

__ndrange = ti.expr_init(ti.ndrange(2, 2))

__ndrange = ti.ndrange(2, 2)

It seems grouped-ndrange's are braced with ti.expr_init while multi-index-ndrange's are not?

Our frontend AST transformer is so ad-hocly single-pass that I don't even dear to modify it, in afraid of functionality regression. As @yuanming-hu said in GAMES 201's final lecture, simplicity matters, a human brain is not possible for handling such complicated single-pass system. Instead, we should use multi-pass system, to separate the functionality for easy-maintain. We should definitely go for multi-pass AST transforming manner at some point, before everything blows up completely.