[muti_backend] #294

Galaxy1458 · 2024-11-16T05:37:22Z

PR Category

Other, support muti_backend

Type of Change

New Feature

Description

After this code is executed, the global device information is obtained for future use import flag_gems
The vendor maintains its own auto_tune config and calls the get_op_tune_config interface directly when passing parameters

@triton.autotune(
    configs=runtime.get_op_tune_config("bmm"),
    key=["M", "N", "K"],
)

Use runtime.device_guard uniformly for device_guard such as torch.cuda.device
with runtime.device_guard(x.device)
The specified operator can be disabled at initialization flag_gems.enable(unused=["add", "mm"])

Issue

Progress

Change is properly reviewed (1 reviewer required, 2 recommended).
Change is responded to an issue.
Change is fully covered by a UT.

Performance

StrongSpoon · 2024-11-21T09:30:05Z

src/flag_gems/runtime/backend/__init__.py

+
+vendors = vendors
+AUTOGRAD = AUTOGRAD
+Autograd = Autograd


what do these mean?

for future use.

StrongSpoon · 2024-11-21T10:07:01Z

src/flag_gems/__init__.py


 __version__ = "2.1"
-
+device = runtime.device.device_instance.device_name


is device used here?

StrongSpoon · 2024-11-22T02:04:01Z

src/flag_gems/runtime/register/__init__.py

+    global register_instance
+    if not register_instance:
+        register_instance = Register(*args, **kargs)
+    return register_instance


what is to_register used for?

for global instance, will be used in next pr.

StrongSpoon · 2024-11-22T02:30:00Z

src/flag_gems/runtime/register/register.py

+    def get_vendor_unused_op(self):
+        if self.device.vendor != backend.vendors.NVIDIA:
+            return backend.get_curent_device_unused_op(self.device.vendor_name)
+        return {}


user might specify unused ops when device is NVIDIA.

The unused_list consists of two parts, one specified by the vendor and the other specified by the user in enable(unused=[xxx]) or use_gems(unused=[xxx])

StrongSpoon · 2024-11-22T02:53:29Z

src/flag_gems/runtime/register/register.py

+        self.lib.impl(key, fn, device_key)
+
+    def close(self):
+        self.lib


I don't understand.

will remove close.

StrongSpoon · 2024-11-22T06:32:46Z

src/flag_gems/runtime/backend/__init__.py

+        return get_vendor_module(vendor_name, query).device.get_vendor_info()
+    global vendor_module
+    get_vendor_module(vendor_name)
+    return vendor_module.device.get_vendor_info()


ditto. I didn't get why set query as an argument.

StrongSpoon · 2024-11-22T06:59:07Z

src/flag_gems/runtime/device/device.py

+from .. import backend, error
+
+
+class device_ctx:


what's it used for?

load the device information.

StrongSpoon · 2024-11-22T07:24:39Z

src/flag_gems/runtime/tune_config/config.py

+                    single_config["META"],
+                    num_warps=current_config["num_warps"],
+                    num_stages=current_config["num_stages"],
+                    num_ctas=current_config["num_ctas"],


it's not necessary to fill all fields in triton.Config. triton.Autotuner will automatically complement the default value.

All are passed in because it is not certain which default values were changed in the above calculation

StrongSpoon · 2024-11-22T07:30:15Z

src/flag_gems/runtime/tune_config/config.py

+    def to_gen_config(self, gen_config):
+        param_config = gen_config["param_map"]
+        meta_config = param_config["META"]
+        iteration_keys = list(meta_config) + list(param_config)


is there overlap between meta_config and param_config?

StrongSpoon · 2024-11-22T07:40:43Z

src/flag_gems/runtime/tune_config/config.py

+                current_config["META"][param_key] = single_value
+            else:
+                current_config[param_key] = single_value
+            config_item = self._gen_impl(


I suggest using loop instead of recursion.

iclementine · 2024-11-26T01:45:05Z

src/flag_gems/runtime/device/device.py

+
+
+class device_ctx:
+    def __init__(self, vendor_name=None):


Add necessary documentation to help understanding the code.

iclementine · 2024-11-26T01:47:10Z

src/flag_gems/runtime/device/device.py

+        self.device_info = self.get_vendor(vendor_name)
+        self.vendor_name, self.device_name, self.cmd, self.vendor = self.device_info
+
+    def get_vendor(self, vendor_name=None) -> tuple:


I suggest use typing to add type hints and use type hints that works for older versions of python, for example typing.Tuple. Also, if it is possible to be clear, Tuple of what?

iclementine · 2024-11-26T01:48:07Z

src/flag_gems/runtime/device/device.py

+    def get_vendor(self, vendor_name=None) -> tuple:
+        if vendor_name is not None:
+            return backend.get_vendor_info(vendor_name)
+        vendor_from_evn = self._get_vendor_from_evn()


If you mean environment, please use env, which is a more accepted abbreviation for it.

iclementine · 2024-11-26T02:00:18Z

src/flag_gems/runtime/backend/__init__.py

+    def get_module(vendor_name):
+        current_file_path = os.path.abspath(__file__)
+        current_dir_path = os.path.dirname(current_file_path)
+        sys.path.append(current_dir_path)


Adding this path to sys.path may increase the risk to cause name conflicts with other packages.
Is there alternative ways to do this?

iclementine · 2024-11-26T02:04:20Z

src/flag_gems/runtime/backend/_nvidia/__init__.py

+class device:
+    @staticmethod
+    def get_vendor_info():
+        return ("nvidia", "cuda", "nvidia-smi")


What does this tuple mean?
If the fields have specific meaning, consider using a namedtuple or a dataclass to represent it.

iclementine · 2024-11-26T02:07:53Z

src/flag_gems/runtime/error/__init__.py

@@ -0,0 +1,3 @@
+from .error import *  # noqa: F403
+
+__all__ = ["*"]


If the module is simple, do not wrap it into a package. A simple module is enough.

iclementine · 2024-11-26T02:10:33Z

src/flag_gems/runtime/backend/_nvidia/__init__.py

+from .ops import *  # noqa: F403
+
+
+class device:


What's the relationship of this device and device_ctx.

If the only purpose of this class is to provide a static method, why bother?

device is stored in each vendor's directory and maintained by each vendor. device_ctx is used to obtain device information of each vendor

iclementine · 2024-11-26T02:18:40Z

src/flag_gems/runtime/device/device.py

+    def _get_vendor_from_lib(self):
+        try:
+            return triton.get_vendor_info()
+        except Exception:
+            return torch.get_vendor_info()


Is this valid code?

I checked triton and torch, neither has get_vendor_info.

Reserve the associated interface for triton or torch, although they are not implemented yet.

iclementine · 2024-11-26T02:23:29Z

src/flag_gems/runtime/device/device.py

+        device_from_evn = os.environ.get("GEMS_VENDOR")
+        return None if device_from_evn not in self.vendor_list else device_from_evn
+
+    def _get_vendor_from_sys(self):


This function simply runs a detection to detect if device is available by try a command and expect return code to be zero. Is this a reasonable way to detect all devices?

There is no other way at the gems level. will push more chip vendors to provide apis at the torch or triton level in the future

iclementine · 2024-12-03T02:26:06Z

src/flag_gems/runtime/__init__.py

+
+global configer, device
+configer = Config()
+device = device_ctx()


If this is intended to be a singleton, make it a real singleton.

See also the device in device.py

iclementine · 2024-12-03T02:27:21Z

src/flag_gems/runtime/error.py

@@ -0,0 +1,15 @@
+def backend_not_support(device_name, backend_list):
+    raise RuntimeError(f"The {device_name} device is not support currently. ")


iclementine · 2024-12-03T02:28:36Z

src/flag_gems/runtime/error.py

+
+def device_not_found():
+    raise RuntimeError(
+        "No devices were detected on your machine ! \n "


No is followed by singular form.

iclementine · 2024-12-03T02:32:14Z

src/flag_gems/runtime/error.py

+
+def register_error(e):
+    raise RuntimeError(
+        e, "An error was encountered while registering the triton operator."


Is this a conventional way to initialize an RuntimeError?

iclementine · 2024-12-03T02:35:29Z

src/flag_gems/runtime/backend/backend_utils.py

+
+
+@dataclass
+class vendor_info_base:


Please follow our code style and use CamaelCase for class name.

Is this class designed to be a base class?
I don't see it bing subclassed?

Also, add documentation to clarify the intention of this class.

iclementine · 2024-12-03T02:42:34Z

src/flag_gems/runtime/backend/_nvidia/__init__.py

+
+from .ops import *  # noqa: F403
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))


Dont's append other path to sys.path since it may cause name conflicts with other packages.

Also, why do you need this?
It may be a sign that the dependencies between modules or packages are not properly organized.

iclementine · 2024-12-03T02:53:49Z

src/flag_gems/runtime/backend/_nvidia/ops/__init__.py

@@ -0,0 +1,3 @@
+from .add import add
+
+__all__ = ["add"]


What is the purpose of this package?

For backend-specific operators?

iclementine · 2024-12-03T02:56:08Z

src/flag_gems/runtime/backend/__init__.py

+    try:
+        exec(compiled_code, globals())
+    except Exception as e:
+        RuntimeError(e)


Why not just raise e?

Is Constructing a RuntimeError the desired behavior?

iclementine · 2024-12-03T03:17:37Z

src/flag_gems/runtime/config.py

+
+class Config:
+    def __init__(self):
+        self.config = self.get_vendor_tune_config()


If it is just some raw config that need to be preprocessed, do not expose it or do not use such a PLAIN name, which is confusing.

iclementine · 2024-12-03T03:17:42Z

src/flag_gems/runtime/config.py

+from .device import device
+
+
+class Config:


I don't see the point of this class.

It is actually not a config. It is a config loader.

src/flag_gems/runtime/backend/_nvidia/ops/add.py

Bowen12992 · 2024-12-03T06:56:45Z

src/flag_gems/runtime/backend/_nvidia/ops/add.py

@@ -0,0 +1,52 @@
+import torch


what's the meaning of _ in _nvidia dir?

nvidia is an existing library

src/flag_gems/runtime/backend/backend_utils.py

src/flag_gems/runtime/backend/_nvidia/tune_configs.json

src/flag_gems/runtime/commom_utils.py

src/flag_gems/runtime/register.py

Bowen12992

I have some same questions with @StrongSpoon and @iclementine, so it would be better if you could answer their questions in the conversation to help me understand~

Galaxy1458 added 20 commits November 16, 2024 05:24

new feature, muti_backend

fd24be2

update auto_tune_module

c8c1bcd

update auto_tune_module

8bb14e3

update auto_tune_module

cca8de5

update __init__

1b8cc65

rebase

d1596dd

rebase

e6bbce4

fix bug

81ee7a8

modifiy auto_tune_config

f6b3e94

fix bug

8b77745

fix bug

9fbb1e3

fix conflict

d55a303

update

0adbfe5

update

389f726

update scatter&gather

c2b03c1

fix auto_tune

51b22cf

add gen_torch_device_fn

8e0617e

fix codestyle

d359b6a

fix codestyle

df73601

fix comflict

3f64fff

StrongSpoon reviewed Nov 22, 2024

View reviewed changes

iclementine reviewed Nov 26, 2024

View reviewed changes

Galaxy1458 added 7 commits November 28, 2024 14:04

Modify code based on comments

1b98131

Adjust the code based on the comments

8c83a46

Modify gen_impl with loops instead of recursion

dbb90d7

Update code structure

14a9a19

Polish code

bcf2d78

update

6854b5b

Polish code

5e46495

iclementine reviewed Dec 3, 2024

View reviewed changes

Modify code based on comments

98fa16a