diff --git a/lib/component.py b/lib/component.py index 32aaf73f4f387255b6aea98eb708decd83ca0170..bc78f4b37c5a9e4e56aa91cfd9df10f3d44c6dc3 100644 --- a/lib/component.py +++ b/lib/component.py @@ -219,10 +219,29 @@ class GPU(Video): # 00120000 Iluvatar CoreX CLASS_CODE = ['000302', '00120000', '000B40', '000380'] + # CABC : Cambricon + VENDOR_ID = ['CABC'] # add others vendor_id + def __init__(self, logdir, cases, category, args): super(Video, self).__init__(logdir, cases, category, args) self.build_cases(cases) + def build(self, dstree): + for dev in dstree.devices: + if dev.vid and is_sub_vendor_id(self.__class__.VENDOR_ID, dev.vid) and dev.is_pci_endpoint(): + ctr = Controller(dev, self) + self.mapping(ctr, dstree) + if ctr.valid: + self.available.append(ctr) + else: + self.unavailable.append(ctr) + self.sort() + self.display('gpu device') + if not self.available: + raise NonAvaliableController('No available gpu device to test! Please double check ' + 'driver whether attached successful.') + self.build_works() + class CPU(Component): def __init__(self, logdir, cases, category, args): diff --git a/lib/utils.py b/lib/utils.py index 070bbad08180c93ded8c70f6f2b6c1c545c4f99f..662a72286c60ed6ce9bc80b48dccdd908bb09bfd 100644 --- a/lib/utils.py +++ b/lib/utils.py @@ -125,6 +125,11 @@ def is_sub_class_code(class_code, ctr_class_code): return True return False +def is_sub_vendor_id(vendor_id, ctr_vendor_id): + for id in vendor_id: + if ctr_vendor_id.startswith(id): + return True + return False def is_virtual_platform(): cmd = 'virt-what'