Embed DeepGEMM source (not submodule) for SM100 raw CUDA GEMM primitives
This commit is contained in:
38
third_party/DeepGEMM/deep_gemm/testing/utils.py
vendored
Normal file
38
third_party/DeepGEMM/deep_gemm/testing/utils.py
vendored
Normal file
@@ -0,0 +1,38 @@
|
||||
import functools
|
||||
import os
|
||||
import torch
|
||||
from typing import Callable
|
||||
|
||||
def get_arch_major() -> int:
|
||||
major, minor = torch.cuda.get_device_capability()
|
||||
return major
|
||||
|
||||
|
||||
def test_filter(condition: Callable):
|
||||
def decorator(func):
|
||||
@functools.wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
if condition():
|
||||
func(*args, **kwargs)
|
||||
else:
|
||||
print(f'{func.__name__}:')
|
||||
print(f' > Filtered by {condition}')
|
||||
print()
|
||||
return wrapper
|
||||
return decorator
|
||||
|
||||
|
||||
def ignore_env(name: str, condition: Callable):
|
||||
def decorator(func):
|
||||
@functools.wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
if condition():
|
||||
saved = os.environ.pop(name, None)
|
||||
func(*args, **kwargs)
|
||||
if saved is not None:
|
||||
os.environ[name] = saved
|
||||
else:
|
||||
func(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
return decorator
|
||||
Reference in New Issue
Block a user