AMD GPU extra-language ops import error

[kxyk99]

Describe the bug

As using the AMD GPU backend, it can't support some tl.extra.shim ops in triton compiler. e.g. div_rn, div_ru and so on. they are only supported by nvidia now. and i think maybe there is a better way to support different tl.extra.language ops in trition for different backend, but not those hardcode w/o any conditions. thanks

issue log:
Traceback (most recent call last):
File "/rocm_perf/FlagGems/src/flag_gems/testing/test_amd_backend.py", line 3, in
import flag_gems
File "/rocm_perf/FlagGems/src/flag_gems/init.py", line 15, in
from .fused import * # noqa: F403
File "/rocm_perf/FlagGems/src/flag_gems/fused/init.py", line 6, in
from .outer import outer
File "/rocm_perf/FlagGems/src/flag_gems/fused/outer.py", line 5, in
from ..ops import mul, mv
File "/rocm_perf/FlagGems/src/flag_gems/ops/init.py", line 43, in
from .div import (
File "/rocm_perf/FlagGems/src/flag_gems/ops/div.py", line 9, in
div_rn = tl_extra_shim.div_rn
AttributeError: module 'triton.language.extra.hip.libdevice' has no attribute 'div_rn'

and you can find them in the src/flaggems/ops/div.py

[tongxin]

Hi, that's indeed a limitation of our current API. Is there a more portable design that you could suggest? @Galaxy1458 could you also comment on this issue?

[kxyk99]

those extra ops should be implemented by the backend. as default, you can provide the cpu software implement at high level APIs. it's similar to the pytorch dispatcher.

and i design a syntactic sugar to do it. they are in triton_language_helper.py

def use_backend(module):
"""using backend module impl"""
def decorator(func):
func_name = func.name
if hasattr(module, func_name):
try:
return getattr(module, func_name)
except Exception:
pass
return func
return decorator

def use_tl_extra(func):
"""backend function shim"""
return use_backend(tl_extra_shim)(func)

Usage:
@use_tl_extra
@triton.jit
def div_rn(x, y):
"""div_rn default - round to nearest"""
result = x / y
rounded_result = tl.floor(result + 0.5)
return rounded_result

the test pass, if agree with it, i can submit them for code review

[Galaxy1458]

@kxyk99 Thank you for your attention. We sincerely regret that the current test for the AMD GPU backend lacks robustness and apologize for any inconvenience this may have caused. We consider your proposed solution to be highly valuable. If you are willing, we would be looking forward to submitting a PR of your method to us.
And the root cause of this issue has been identified as incomplete implementations of certain libdevice methods (e.g., div_rn) within Triton's AMD backend. To advance the functionality and reliability of the AMD backend, we cordially invite you to contribute enhancements to its libdevice implementation if you are willing. Your expertise may be instrumental in supporting these critical capabilities.

[kxyk99]

Sure, I'm very pleased to join the FlagGems open-source development, especially for the AMD backend.

it seems that i can't submit codes to repo.
git push origin master
remote: Permission to FlagOpen/FlagGems.git denied to kxyk99.
fatal: unable to access 'https://github.com/FlagOpen/FlagGems.git/': The requested URL returned error: 403

[iclementine]

You can fork the repo and create a pull request. Pushing to master branch is not allowed since the branch is protected.

[kxyk99]

okay. Done. pls review. not sure that if adding @triton.jit before div_rn. maybe you can give me some tips. thanks d4e490e

[Galaxy1458]

The PR has been accepted. Thank you for your contribution!