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.

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

[CODEGEN] Vector-Codegen support for llvm-pure-intrin #16985

Merged
merged 3 commits into from
Jun 4, 2024

Conversation

rutkoor
Copy link
Contributor

@rutkoor rutkoor commented May 10, 2024

If we have a loop with llvm intrinsic call like below,

Before,

for i in range T.serial(32):
    A[i] = T.call_llvm_pure_intrin(_ , _ , _ , B[i])

using sch.vectorize, we cannot generate ramp node for the above code.

In this pull request, I enabled vectorization support for the llvm_pure_intrin builtin. During the vectorize pass, I explicitly vectorized the compute expressions that serve as input to the intrinsic. As a result of these changes, I was able to generate the ramp instructions for the expressions inside the call_llvm_pure_intrin function.

After,

for i in range T.serial(...):
    A[T.Ramp( ... )] = T.call_llvm_pure_intrin(_ , _ , _ , B[T.Ramp( ... )])

Note: If the intrinsic doesn't support vector arguments as input, then the error will be thrown during llvm codegen.

@rutkoor rutkoor force-pushed the vector_codegen_llvm_pure_intrin branch from deefe06 to d3be9ab Compare May 10, 2024 10:49
@rutkoor
Copy link
Contributor Author

rutkoor commented May 11, 2024

@tvm-bot rerun

@rutkoor
Copy link
Contributor Author

rutkoor commented May 14, 2024

cc: @tqchen @Lunderberg

@rutkoor
Copy link
Contributor Author

rutkoor commented May 16, 2024

cc: @cbalint13 @ekalda

@@ -139,7 +139,8 @@ TIR_DEFINE_BUILTIN_FUNC(call_llvm_intrin)
TIR_DEFINE_BUILTIN_FUNC(call_llvm_pure_intrin)
.set_attr<TCallEffectKind>("TCallEffectKind", Integer(CallEffectKind::kPure))
.set_attr<TScriptDtypePrintLocation>("TScriptDtypePrintLocation",
Integer(ScriptDtypePrintLocation::kFirst));
Integer(ScriptDtypePrintLocation::kFirst))
.set_attr<TVectorizable>("TVectorizable", true);
Copy link
Contributor

Choose a reason for hiding this comment

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

Can add a small testcase for the IR printer ?
tests/python/unittest/test_tvmscript_printer_tir.py

@cbalint13
Copy link
Contributor

cbalint13 commented May 16, 2024

@rutkoor
Implementation LGTM (I did't go after the functionality).

@rutkoor
Copy link
Contributor Author

rutkoor commented Jun 3, 2024

@rutkoor Implementation LGTM (I did't go after the functionality).

Hi @cbalint13, thank you for your quick reply. And apologize for not replying back sooner.

I have added a testcase inside test_tvmscript_printer_tir.py and also explained what I am doing in the description of this PR. Please let me know if you want me to add more details.

Thanks.

@cbalint13
Copy link
Contributor

@rutkoor Implementation LGTM (I did't go after the functionality).

Hi @cbalint13, thank you for your quick reply. And apologize for not replying back sooner.

I have added a testcase inside test_tvmscript_printer_tir.py and also explained what I am doing in the description of this PR. Please let me know if you want me to add more details.

Thank you very much @rutkoor !

Thanks.

@rutkoor
Copy link
Contributor Author

rutkoor commented Jun 3, 2024

@tvm-bot rerun

@cbalint13 cbalint13 merged commit 78a1f80 into apache:main Jun 4, 2024
24 of 26 checks passed
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.

None yet

2 participants