all_ops.py 1004 B

1234567891011121314151617181920212223242526
  1. """
  2. Copyright 2020 The Microsoft DeepSpeed Team
  3. """
  4. import os
  5. import pkgutil
  6. import importlib
  7. from deepspeed.accelerator import get_accelerator
  8. # List of all available ops
  9. # reflect all builder names into __op_builders__
  10. op_builder_dir = get_accelerator().op_builder_dir()
  11. op_builder_module = importlib.import_module(op_builder_dir)
  12. __op_builders__ = []
  13. for _, module_name, _ in pkgutil.iter_modules([os.path.dirname(op_builder_module.__file__)]):
  14. # avoid self references
  15. if module_name != 'all_ops' and module_name != 'builder' and module_name != 'builder_names':
  16. module = importlib.import_module("{}.{}".format(op_builder_dir, module_name))
  17. for member_name in module.__dir__():
  18. if member_name.endswith('Builder'):
  19. # append builder to __op_builders__ list
  20. builder = get_accelerator().create_op_builder(member_name)
  21. __op_builders__.append(builder)
  22. ALL_OPS = {op.name: op for op in __op_builders__ if op is not None}