Taichi: [Lang] [IR] Enhanced `print`

[yuanming-hu]

Concisely describe the proposed feature

A lot of users complain that print works for only scalars in Taichi kernels.

• Let print support Vectors/Matrices.
• Let print support strings as well.

Describe the solution you'd like (if any)

• Modify the frontend AST transformer
• Maybe modify the Matrix class in Python (you decide how :-)
• Extend PrintStmt to let it support a vector of std::string/Stmt *.

Desired API

@ti.kernel
def test():
    a = 1.0
    print(a) 
    # "1.0"      instead of "a = 1.0"

    print("a = ", a) 
    # "a = 1.0"

    b = ti.Matrix([[1, 2]])
    print("b = ", b, "b[0, 0] =", b[0, 0]) 
    # "b = [[1, 2]] b[0, 0] = 1"
    ...

Additional comments

In the future, we might want to support f-strings (https://www.python.org/dev/peps/pep-0498/), but for now supporting a list of arguments separated by commas is good enough.

[archibate]

Thank for purposing this! Here are my thoughts:

1. To support print for matrix/vector:

class Matrix:
  is_taichi_class = True
  ...
  def print(self):
    ... # implementation here

def ti_print(x):
  if ti.is_taichi_class(x): # May also apply to other taichi classes (if we add later)
    return x.print()
  else:
    ...

1. To support multiple arguments for print, translate print("a = ", a) into:

PrintStringStmt "a = " (eol=false) # no '\n'
PrintNumberStmt $a (eol=true) # add '\n'

[yuanming-hu]

Thanks for the ideas. One thing to keep in mind: we want a single PrintStmt for each print in Taichi kernels since on parallel hardware breaking down a single print will probably result in messy outputs due to parallelism.

[archibate]

I'm thinking about:

class FrontendPrintStmt {
  std::vector<std::eithier<std::string, Expr>> content;

or:

class FrontendPrintStmt {
  std::vector<std::pair<std::string, Expr>> content;
  std::string end;

wdyt?

[yuanming-hu]

Sure, please feel free to make use of https://en.cppreference.com/w/cpp/utility/variant

[archibate]

Cool C++17! Does this also work for pybind11?

[archibate]

How do I deal with type_hint? print(i32, f32), we abandon that?
Btw, is load_if_ptr on the refactor list?

[archibate]

I found std::variant support is WIP in pybind11: https://github.com/pybind/pybind11/pull/1210.
The PR seems to be stale and we need some walk around before they take this done.

[archibate]

[I 05/22/20 23:24:47.041] [compile_to_offloads.cpp:operator()@21] Access lowered:
kernel {
  <i32 x1> $0 = arg[0]
  <i32 x1> $1 = const [2]
  <i32 x1> $2 = mul $0 $1
  print $0, $2
}
[I 05/22/20 23:24:47.041] [compile_to_offloads.cpp:operator()@21] DIE:
kernel {
  print $0, $2
}

[E 05/22/20 23:14:09.200] [codegen_llvm.cpp:finalize_offloaded_task_function@1237] !llvm::verifyFunction(*func, &errs())
Why?
What does TI_STMT_DEF_FIELDS do with std::variant?

[yuanming-hu]

I found std::varient support is WIP in pybind11: pybind/pybind11#1210.
The PR seems to be stale and we need some walk around before they take this done.

Actually we don't really need to expose std::variant to pybind11. We can create a FrontendPrintStmt builder class with add_expr and add_string methods and expose these to pybind11.

[yuanming-hu]

[I 05/22/20 23:24:47.041] [compile_to_offloads.cpp:operator()@21] Access lowered:
kernel {
  <i32 x1> $0 = arg[0]
  <i32 x1> $1 = const [2]
  <i32 x1> $2 = mul $0 $1
  print $0, $2
}
[I 05/22/20 23:24:47.041] [compile_to_offloads.cpp:operator()@21] DIE:
kernel {
  print $0, $2
}

[E 05/22/20 23:14:09.200] [codegen_llvm.cpp:finalize_offloaded_task_function@1237] !llvm::verifyFunction(*func, &errs())
Why?
What does TI_STMT_DEF_FIELDS do with std::variant?

Please see
https://github.com/taichi-dev/taichi/blob/2a610a8b7aeb1d44c358426e5d6ab32e9e3c9067/taichi/ir/ir.h#L1246-L1273

[archibate]

Please see

So I should add an } else if constexpr (std::is_same<std::variant<Expr, std::string>>::value) {?

[yuanming-hu]

exactly!!

[archibate]

exactly!!

I failed to make it work. So I used another vector to store Stmt * and pass to TI_STMT_FIELDS, hope this works.