Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Relax] Ignore non-relax functions in relax.transform.RunCodegen #16586

Merged
merged 2 commits into from
Feb 17, 2024
Merged

[Relax] Ignore non-relax functions in relax.transform.RunCodegen #16586

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e6b5da8
[Relax] Ignore non-relax functions in relax.transform.RunCodegen
Lunderberg Jan 12, 2024
60bdbeb
Remove debug changes
Lunderberg Feb 16, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/relax/transform/run_codegen.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,9 +116,9 @@ class CodeGenRunner : ExprMutator {
auto ret_sinfo = GetStructInfo(call);
if (auto it = extern_funcs_.find(gvar_node); it != extern_funcs_.end()) {
return create_call_dps_packed(it->second, ret_sinfo);
} else {
} else if (auto opt_func = builder_->GetContextIRModule()->Lookup(gvar).as<Function>()) {
// TODO(@sunggg): Is there any better way to get this func?
Function func = Downcast<Function>(builder_->GetContextIRModule()->Lookup(gvar));
Function func = opt_func.value();
Expr new_func = VisitExpr(func);

if (new_func->IsInstance<ExternFuncNode>()) {
Expand Down
32 changes: 31 additions & 1 deletion tests/python/relax/test_transform_codegen_pass.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@
)

# Global variable in pytest that applies markers to all tests.
pytestmark = [has_tensorrt_codegen, has_tensorrt_runtime]
# pytestmark = [has_tensorrt_codegen, has_tensorrt_runtime]
Copy link
Member

Choose a reason for hiding this comment

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

I guess we will need this

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Whoops, thank you for the catch. This was commented out for debugging, and the PR is updated to remove the change.


# Target gpu
target_str = "nvidia/nvidia-t4"
Expand Down Expand Up @@ -350,6 +350,36 @@ def main(

after = relax.transform.RunCodegen()(Before)
tvm.ir.assert_structural_equal(after["main"], Expected["main"])
after.show()
Copy link
Member

Choose a reason for hiding this comment

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

remove this printing

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Thank you, and cleaned up!



def test_no_op_for_call_to_tir():
"""Calls to PrimFunc are ignored

RunCodegen should only update calls to Relax functions annotated
with the `"Codegen"` attribute. Calls to any other function type
should be ignored.

This is a regression test. Previous implementations performed an
unconditional cast from `tvm::BaseFunc` to `tvm::relax::Function`,
which produced an error.
"""

@tvm.script.ir_module
class Before:
@R.function
def main(x: R.Tensor):
R.func_attr({"relax.force_pure": True})
_ = Before.shape_func(x)
return x

@T.prim_func(private=True)
def shape_func(H: T.Buffer(T.int64(4), "int64")):
H[T.int64(0)] = H[T.int64(0)] + T.int64(1)

Expected = Before
After = relax.transform.RunCodegen()(Before)
tvm.ir.assert_structural_equal(Expected, After)


# TODO(@sunggg): test with more complex patterns (e.g., multiple annots, mixed codegens, different ops, const binding)
Expand Down
Loading