sunggg opened a new pull request, #15217:
URL: https://github.com/apache/tvm/pull/15217
Currently, `PrimValue` is treated just like a tensors and `FuseOpByPattern`
pass creates composite functions that have `PrimValues` as function parameters.
For example, `param_0` and `param_1` are defined to handle `PrimValues`
```Python
@I.ir_module
class Module:
@R.function(private=True)
def fused_relax_clip(x: R.Tensor((10, 10), dtype="float32"), param_0:
R.Prim("int64"), param_1: R.Prim("int64")) -> R.Tensor((10, 10),
dtype="float32"):
R.func_attr({"Composite": "x.clip", "Primitive": 1})
with R.dataflow():
gv: R.Tensor((10, 10), dtype="float32") = R.clip(x, param_0,
param_1)
R.output(gv)
return gv
@R.function
def main(x: R.Tensor((10, 10), dtype="float32")) -> R.Tensor((10, 10),
dtype="float32"):
cls = Module
with R.dataflow():
gv: R.Tensor((10, 10), dtype="float32") =
cls.fused_relax_clip(x, R.prim_value(0), R.prim_value(4))
R.output(gv)
return gv
```
However, this does not comply with BYOC codegens since they expect each
composite function to contain these information that would have been stored in
operator attribute before the adoption of `PrimValue`.
Therefore, this PR fixes `FuseOpByPattern` to include `PrimValue` inside the
composite function.
With this change, the example above will have the following output.
```Python
@I.ir_module
class Expected1:
@R.function(private=True)
def fused_relax_clip(x: R.Tensor((10, 10), dtype="float32")) ->
R.Tensor((10, 10), dtype="float32"):
R.func_attr({"Composite": "x.clip", "Primitive": 1})
with R.dataflow():
gv: R.Tensor((10, 10), dtype="float32") = R.clip(x,
R.prim_value(0), R.prim_value(4))
R.output(gv)
return gv
@R.function
def main(x: R.Tensor((10, 10), dtype="float32")) -> R.Tensor((10,
10), dtype="float32"):
cls = Expected1
with R.dataflow():
gv: R.Tensor((10, 10), dtype="float32") =
cls.fused_relax_clip(x)
R.output(gv)
return gv
```
cc. @masahi
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
