You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2021/05/06 07:24:20 UTC
[GitHub] [tvm] Hzfengsy opened a new pull request #7987: [TensorIR] CreatePrimFunc from TE
Hzfengsy opened a new pull request #7987:
URL: https://github.com/apache/tvm/pull/7987
This PR introduces a method to create PrimFunc (and further TensorIR schedule) from Tensor Expression. (e.g. `te.compute`)
cc @tqchen @junrushao1994 @comaniac @jcf94 @xqdan
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen edited a comment on pull request #7987: [TensorIR] CreatePrimFunc from TE
Posted by GitBox <gi...@apache.org>.
tqchen edited a comment on pull request #7987:
URL: https://github.com/apache/tvm/pull/7987#issuecomment-833857175
Create prim_func returns a PrimFunc, the name is attached in module creation process, we can likely recompose by
```python
mod = IRModule({"mycoolname": create_prim_func())
```
Note such name is different from the global_symbol(name of generated function), which can be set by `create_prim_func().with_attr("tir.global_symbol", global_symbol)`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 commented on pull request #7987: [TensorIR] CreatePrimFunc from TE
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on pull request #7987:
URL: https://github.com/apache/tvm/pull/7987#issuecomment-834755740
Thanks for the great discussion @tkonolige @tqchen @Hzfengsy! I am going to merge it in given there is no objection :-)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen edited a comment on pull request #7987: [TensorIR] CreatePrimFunc from TE
Posted by GitBox <gi...@apache.org>.
tqchen edited a comment on pull request #7987:
URL: https://github.com/apache/tvm/pull/7987#issuecomment-833897919
If we really want to encourage the name, and given IRModule is the first class item for running exchange between passes. Let us consider make the function return an IRModule instead, in this case, perhaps we should rename to `create_tir` (coming back to another old name)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tkonolige commented on pull request #7987: [TensorIR] CreatePrimFunc from TE
Posted by GitBox <gi...@apache.org>.
tkonolige commented on pull request #7987:
URL: https://github.com/apache/tvm/pull/7987#issuecomment-833683027
Would it be possible to add a name field to the `create_prim_func` function? That way when we print out primfuncs for debugging, they aren't all called "main".
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 commented on pull request #7987: [TensorIR] CreatePrimFunc from TE
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on pull request #7987:
URL: https://github.com/apache/tvm/pull/7987#issuecomment-833997012
@tkonolige The mechanism causing the issue you mentioned is that when constructing a `tir::Schedule` class, and the input is a PrimFunc, then it will put that PrimFunc into an IRModule with the default name `main`, so that the schedule is always IRModule-based.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on pull request #7987: [TensorIR] CreatePrimFunc from TE
Posted by GitBox <gi...@apache.org>.
tqchen commented on pull request #7987:
URL: https://github.com/apache/tvm/pull/7987#issuecomment-833897919
I think we can add `global_symbol` as an optional argument(default to main)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on pull request #7987: [TensorIR] CreatePrimFunc from TE
Posted by GitBox <gi...@apache.org>.
tqchen commented on pull request #7987:
URL: https://github.com/apache/tvm/pull/7987#issuecomment-834469067
I see, i now agree that we can go with the compositional style API. Thanks @junrushao1994 @Hzfengsy
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 commented on pull request #7987: [TensorIR] CreatePrimFunc from TE
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on pull request #7987:
URL: https://github.com/apache/tvm/pull/7987#issuecomment-833997537
@tqchen The problem of creating an IRModule instead is that we might want to an IRModule that contains multiple PrimFuncs, e.g.:
```python
mod = IRModule({
"func_1": create_prim_func(...),
"func_2": create_prim_func(...),
})
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on pull request #7987: [TensorIR] CreatePrimFunc from TE
Posted by GitBox <gi...@apache.org>.
tqchen commented on pull request #7987:
URL: https://github.com/apache/tvm/pull/7987#issuecomment-833857175
Create prim_func returns a PrimFunc, the name is attached in module creation process, we can likely recompose by
```
mod = IRModule({"mycoolname": create_prim_func())
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen commented on a change in pull request #7987: [TensorIR] CreatePrimFunc from TE
Posted by GitBox <gi...@apache.org>.
tqchen commented on a change in pull request #7987:
URL: https://github.com/apache/tvm/pull/7987#discussion_r627758326
##########
File path: python/tvm/te/operation.py
##########
@@ -426,3 +427,51 @@ def reduce_axis(dom, name="rv", thread_tag="", span=None):
An iteration variable representing the value.
"""
return tvm.tir.IterVar(dom, name, 2, thread_tag, span)
+
+
+def create_prim_func(ops: List[_tensor.Tensor]) -> tvm.tir.PrimFunc:
+ """Create a TensorIR PrimFunc from tensor expression
+ Parameters
+ ----------
+ ops : List[Tensor]
+ The source expression.
+
+ Example
+ -------
+ We define a matmul kernel using following code:
+
+ .. code-block:: python
+
+ import tvm
+ from tvm import te
+
+ A = te.placeholder((128, 128), name="A")
+ B = te.placeholder((128, 128), name="B")
+ C = te.compute((128, 128), lambda x, y: te.sum(A[x, k] * B[y, k], axis=k), name="C")
+ func = create_prim_func([A, B, C])
+ print(tvm.script.asscript(func))
+
+ If we want to use TensorIR schedule to do transformations on such kernel,
+ we need to use `create_prim_func([A, B, C])` to create a schedulable PrimFunc.
+ The generated function looks like:
+
+ .. code-block:: python
+
+ def tir_matmul(a: ty.handle, b: ty.handle, c: ty.handle) -> None:
Review comment:
add a @tvm.script.tir decorator
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tkonolige commented on pull request #7987: [TensorIR] CreatePrimFunc from TE
Posted by GitBox <gi...@apache.org>.
tkonolige commented on pull request #7987:
URL: https://github.com/apache/tvm/pull/7987#issuecomment-833871651
Instead of `create_prim_func().with_attr("tir.global_symbol", global_symbol)`, I think we should just make `create_prim_func` take a name parameter. I was working on autoscheduler earlier and all the primfuncs it had were called "main". Having a name parameter encourages the user to name their functions and makes debugging easier.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen edited a comment on pull request #7987: [TensorIR] CreatePrimFunc from TE
Posted by GitBox <gi...@apache.org>.
tqchen edited a comment on pull request #7987:
URL: https://github.com/apache/tvm/pull/7987#issuecomment-834469067
I see, i now agree that we can go with the compositional style API. While I agree with @tkonolige that naming functions is going to be useful. We can do that through the compositional API and when creating example tutorials for scheduling. Thanks @junrushao1994 @Hzfengsy
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] junrushao1994 merged pull request #7987: [TensorIR] CreatePrimFunc from TE
Posted by GitBox <gi...@apache.org>.
junrushao1994 merged pull request #7987:
URL: https://github.com/apache/tvm/pull/7987
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tqchen edited a comment on pull request #7987: [TensorIR] CreatePrimFunc from TE
Posted by GitBox <gi...@apache.org>.
tqchen edited a comment on pull request #7987:
URL: https://github.com/apache/tvm/pull/7987#issuecomment-833857175
Create prim_func returns a PrimFunc, the name is attached in module creation process, we can likely recompose by
```python
mod = IRModule({"mycoolname": create_prim_func())
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org