Coverage for src/flag_gems/runtime/__init__.py: 89%

1from . import backend, common, error 

2from .backend.device import DeviceDetector 

3from .configs_loader import TunedConfigLoader 

4from .flagtune import flagtune, flagtune_enabled 

5 

6config_loader = TunedConfigLoader() 

7device = DeviceDetector() 

8 

9""" 

10The dependency order of the sub-directory is strict, and changing the order arbitrarily may cause errors. 

11""" 

12 

13# torch_device_fn is like 'torch.cuda' object 

14backend.set_torch_backend_device_fn(device.vendor_name) 

15torch_device_fn = backend.gen_torch_device_object() 

16 

17# torch_backend_device is like 'torch.backend.cuda' object 

18torch_backend_device = backend.get_torch_backend_device_fn() 

19 

20 

21def get_tuned_config(op_name): 

22 return config_loader.get_tuned_config(op_name) 

23 

24 

25def get_heuristic_config(op_name): 

26 return config_loader.get_heuristics_config(op_name) 

27 

28 

29def get_expand_config(op_name, yaml_path=None): 

30 return config_loader.get_expand_config(op_name=op_name, yaml_path=yaml_path) 

31 

32 

33def ops_get_configs(op_name, pre_hook=None, yaml_path=None): 

34 return config_loader.ops_get_configs( 

35 op_name=op_name, 

36 pre_hook=pre_hook, 

37 yaml_path=yaml_path, 

38 ) 

39 

40 

41__all__ = [ 

42 "TunedConfigLoader", 

43 "DeviceDetector", 

44 "backend", 

45 "common", 

46 "config_loader", 

47 "device", 

48 "error", 

49 "flagtune", 

50 "flagtune_enabled", 

51 "get_expand_config", 

52 "get_heuristic_config", 

53 "get_tuned_config", 

54 "ops_get_configs", 

55 "replace_customized_ops", 

56 "torch_backend_device", 

57 "torch_device_fn", 

58]