Skip to content

Conversation

@masahi
Copy link
Member

@masahi masahi commented May 16, 2022

I came across a use case where being able to call a python function as a sugar to construct AST fragments is extremely handy.

The attached test case illustrates two ways of describing the same TVMScript function, one of which is using the proposed syntax sugar. Here, the called function does some complicated index transformation - Calling a function to construct such mapping, rather than spelling them out by hand, significantly cleans up the TVMScript description.

Moreover, since the index map is just a normal python function, it can also be used in the layout transformation primitive in a schedule:

def index_map(i, j):
   # The same index map shared between TVMScript and Python schedule
    return (i // 16, j // 16, *shared_16x16_to_ldmatrix_32x8_layout(i % 16, j % 16) ) 

sch.transform_layout(C_warp, 0, "read", index_map)

This way, we can be sure that tensorize pattern matching will not fail, since the intrinsic description (written in TVMScript) and the schedule are using exactly the same index map. Without this capability, we need to keep the index map used in a schedule and its explicit description in TVMScript always in sync manually.

Note:

  • The outputs of the called function are inlined into the following body, i.e. no variable or let binding is created.
  • The called function cannot be arbitrary - it must be a "pure" function consuming and producing TIR expressions. In particular, I don't think we can (or should) construct TIR stmt with this approach. Besides index transformation, describing various activations could be a good use case.

@Hzfengsy @spectrometerHBH @vinx13 @junrushao1994 @Lunderberg @csullivan

@wrongtest-intellif
Copy link
Contributor

That seems so cool. Is this functionality under the scope of meta-programming?

@masahi
Copy link
Member Author

masahi commented May 16, 2022

Is this functionality under the scope of meta-programming?

yeah, I'd say this is the simplest kind of a macro. I'm literally calling the python function inside the parser.

Copy link
Member

@Hzfengsy Hzfengsy left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What will happen if the function body is not supported by TVM? like

def f(x):
     return np.sum(x)

Can we check it and report error?

@masahi masahi force-pushed the tvmscript-func-call branch from c57e26d to 7c188e0 Compare May 18, 2022 06:16
@masahi
Copy link
Member Author

masahi commented May 18, 2022

@Hzfengsy Thanks for the suggestion. I wrapped the func call with try / catch, and added error report dump via report_error.

I couldn't make a good test case, but this is an example of an err msg when a user tries a funky function call:

error: Error occured when invoking the function sqrt: 
loop of ufunc does not support argument 0 of type Var which has no callable sqrt method
 --> test_tvmscript_syntax_sugar.py:344:19
     |  
 344 |              ind = sqrt(i)
     |                    ^^^^^^^
note: run with `TVM_BACKTRACE=1` environment variable to display a backtrace.

@masahi masahi force-pushed the tvmscript-func-call branch from 7c188e0 to 7036f59 Compare May 18, 2022 06:23
@masahi masahi force-pushed the tvmscript-func-call branch from 7036f59 to 5fdd335 Compare May 18, 2022 06:25
Copy link
Member

@Hzfengsy Hzfengsy left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@masahi masahi merged commit a4be2ed into apache:main May 18, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants