You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by "akaashrp (via GitHub)" <gi...@apache.org> on 2025/10/08 02:44:05 UTC
[PR] [WebLLM] Replace int64s with int32s in WebGPU kernels [tvm]
akaashrp opened a new pull request, #18361:
URL: https://github.com/apache/tvm/pull/18361
(no comment)
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [WebLLM] Replace int64s with int32s in WebGPU kernels [tvm]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on code in PR #18361:
URL: https://github.com/apache/tvm/pull/18361#discussion_r2412367199
##########
python/tvm/relax/backend/gpu_generic/sampling.py:
##########
@@ -282,15 +283,15 @@ def parallel_sampling_from_prob(
# at least one iteration
while T.tvm_thread_invariant(
(step_iter[()] == 0 or aggregate[()] < u - eps)
- and T.Cast("int64", step_iter[()]) < T.ceildiv(vocab_size, block_elem)
+ and T.Cast("int32" if "webgpu" in str(tvm.target.Target.current()) else "int64", step_iter[()]) < T.Cast("int32" if "webgpu" in str(tvm.target.Target.current()) else "int64", T.ceildiv(vocab_size, block_elem))
Review Comment:
The expression `"int32" if "webgpu" in str(tvm.target.Target.current()) else "int64"` is used multiple times in this function (here and on line 294). This harms readability and maintainability. Consider defining a Python variable for this string outside of the `@T.prim_func` definition, inside `gpu_multinomial_from_uniform`, and using that variable here. This would make the code cleaner and easier to modify.
##########
python/tvm/topi/gpu/sort.py:
##########
@@ -359,8 +376,13 @@ def merge(source, dest, source_idx, dest_idx):
def mergesort(source, dest, source_idx, dest_idx, size, width, even):
# calculate the start, mid, and end points of this section
start = width * bz
- middle = cast(tvm.te.min(start + tvm.tir.indexdiv(width, 2), size), "int64")
- end = cast(tvm.te.min(start + width, size), "int64")
+ target = tvm.target.Target.current()
+ if "webgpu" in str(target):
+ middle = cast(tvm.te.min(start + tvm.tir.indexdiv(width, 2), size), "int32")
+ end = cast(tvm.te.min(start + width, size), "int32")
+ else:
+ middle = cast(tvm.te.min(start + tvm.tir.indexdiv(width, 2), size), "int64")
+ end = cast(tvm.te.min(start + width, size), "int64")
Review Comment:
This conditional logic can be simplified by defining the `dtype` based on the target and then using it for casting. This avoids repeating the `cast` and `tvm.te.min` calls.
```python
target = tvm.target.Target.current()
dtype = "int32" if "webgpu" in str(target) else "int64"
middle = cast(tvm.te.min(start + tvm.tir.indexdiv(width, 2), size), dtype)
end = cast(tvm.te.min(start + width, size), dtype)
```
##########
python/tvm/topi/gpu/sort.py:
##########
@@ -250,10 +258,17 @@ def serial_merge(
first,
last,
):
- i = ib.allocate("int64", (1,), name="i", scope="local")
- j = ib.allocate("int64", (1,), name="j", scope="local")
- i[0] = aStart + first
- j[0] = bStart + diag - last
+ target = tvm.target.Target.current()
+ if "webgpu" in str(target):
+ i = ib.allocate("int32", (1,), name="i", scope="local")
+ j = ib.allocate("int32", (1,), name="j", scope="local")
+ i[0] = cast(aStart + first, "int32")
+ j[0] = cast(bStart + diag - last, "int32")
+ else:
+ i = ib.allocate("int64", (1,), name="i", scope="local")
+ j = ib.allocate("int64", (1,), name="j", scope="local")
+ i[0] = aStart + first
+ j[0] = bStart + diag - last
Review Comment:
Similar to the previous comment, this block has duplicated code. You can refactor it by determining the `dtype` first and then using it.
```suggestion
target = tvm.target.Target.current()
is_webgpu = "webgpu" in str(target)
dtype = "int32" if is_webgpu else "int64"
i = ib.allocate(dtype, (1,), name="i", scope="local")
j = ib.allocate(dtype, (1,), name="j", scope="local")
i_val = aStart + first
j_val = bStart + diag - last
if is_webgpu:
i[0] = cast(i_val, dtype)
j[0] = cast(j_val, dtype)
else:
i[0] = i_val
j[0] = j_val
```
##########
python/tvm/topi/gpu/sort.py:
##########
@@ -219,11 +219,19 @@ def compare(a, b):
upper_lim = ceil_log2(size)
def get_merge_begin(source, base_idx, aCount, bCount, aStart, bStart, diag, step_count):
- first = ib.allocate("int64", (1,), name="first", scope="local")
- mid = ib.allocate("int64", (1,), name="mid", scope="local")
- last = ib.allocate("int64", (1,), name="last", scope="local")
- first[0] = tvm.te.max(0, diag - bCount)
- last[0] = tvm.te.min(diag, aCount)
+ target = tvm.target.Target.current()
+ if "webgpu" in str(target):
+ first = ib.allocate("int32", (1,), name="first", scope="local")
+ mid = ib.allocate("int32", (1,), name="mid", scope="local")
+ last = ib.allocate("int32", (1,), name="last", scope="local")
+ first[0] = cast(tvm.te.max(0, diag - bCount), "int32")
+ last[0] = cast(tvm.te.min(diag, aCount), "int32")
+ else:
+ first = ib.allocate("int64", (1,), name="first", scope="local")
+ mid = ib.allocate("int64", (1,), name="mid", scope="local")
+ last = ib.allocate("int64", (1,), name="last", scope="local")
+ first[0] = tvm.te.max(0, diag - bCount)
+ last[0] = tvm.te.min(diag, aCount)
Review Comment:
There's quite a bit of duplicated code between the `if` and `else` branches. You can refactor this to improve readability and reduce redundancy by defining the `dtype` first and then using it for allocation and casting. This also makes the logic clearer.
```suggestion
target = tvm.target.Target.current()
is_webgpu = "webgpu" in str(target)
dtype = "int32" if is_webgpu else "int64"
first = ib.allocate(dtype, (1,), name="first", scope="local")
mid = ib.allocate(dtype, (1,), name="mid", scope="local")
last = ib.allocate(dtype, (1,), name="last", scope="local")
max_val = tvm.te.max(0, diag - bCount)
min_val = tvm.te.min(diag, aCount)
if is_webgpu:
first[0] = cast(max_val, dtype)
last[0] = cast(min_val, dtype)
else:
first[0] = max_val
last[0] = min_val
```
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [WebLLM] Replace int64s with int32s in WebGPU kernels [tvm]
Posted by "gemini-code-assist[bot] (via GitHub)" <gi...@apache.org>.
gemini-code-assist[bot] commented on PR #18361:
URL: https://github.com/apache/tvm/pull/18361#issuecomment-3379366358
## Summary of Changes
Hello @akaashrp, I'm Gemini Code Assist[^1]! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request refactors several GPU kernels, specifically in sampling and sorting routines, to conditionally utilize `int32` data types instead of `int64` when compiling for WebGPU targets. This modification is crucial for enhancing performance and ensuring broader compatibility on platforms where `int64` operations might be less efficient or not fully supported by WebGPU. The changes are implemented with conditional logic, preserving `int64` usage for other backend targets.
### Highlights
* **Conditional Integer Type Casting**: The pull request introduces conditional logic to cast `int64` types to `int32` specifically when the compilation target is WebGPU. This affects variables and operations within sampling and sorting kernels.
* **WebGPU Performance and Compatibility**: By switching to `int32` for WebGPU, the changes aim to improve performance and ensure broader compatibility, as `int64` support can be less optimized or problematic on some WebGPU implementations.
* **Target-Specific Code Generation**: The modifications ensure that `int64` types are preserved for other non-WebGPU targets, maintaining existing behavior and performance where `int64` is preferred or necessary.
<details>
<summary><b>Using Gemini Code Assist</b></summary>
<br>
The full guide for Gemini Code Assist can be found on our [documentation page](https://developers.google.com/gemini-code-assist/docs/review-github-code), here are some quick tips.
<b>Invoking Gemini</b>
You can request assistance from Gemini at any point by creating a comment using either `/gemini <command>` or `@gemini-code-assist <command>`. Below is a summary of the supported commands on the current page.
Feature | Command | Description
--- | --- | ---
Code Review | `/gemini review` | Performs a code review for the current pull request in its current state.
Pull Request Summary | `/gemini summary` | Provides a summary of the current pull request in its current state.
Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help | `/gemini help` | Displays a list of available commands.
<b>Customization</b>
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a `.gemini/` folder in the base of the repository. Detailed instructions can be found [here](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github).
<b>Limitations & Feedback</b>
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with :thumbsup: and :thumbsdown: on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up [here](https://google.qualtrics.com/jfe/form/SV_2cyuGuTWsEw84yG).
<b>You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the [Gemini Code Assist IDE Extension](https://cloud.google.com/products/gemini/code-assist).</b>
</details>
[^1]: Review the [Privacy Notices](https://policies.google.com/privacy), [Generative AI Prohibited Use Policy](https://policies.google.com/terms/generative-ai/use-policy), [Terms of Service](https://policies.google.com/terms), and learn how to configure Gemini Code Assist in GitHub [here](https://developers.google.com/gemini-code-assist/docs/customize-gemini-behavior-github). Gemini can make mistakes, so double check it and [use code with caution](https://support.google.com/legal/answer/13505487).
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [WebLLM] Replace int64s with int32s in WebGPU kernels [tvm]
Posted by "MasterJH5574 (via GitHub)" <gi...@apache.org>.
MasterJH5574 merged PR #18361:
URL: https://github.com/apache/tvm/pull/18361
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org