897 lines
30 KiB
Python
897 lines
30 KiB
Python
import torch
|
|
import copy
|
|
from torch.testing._internal.common_methods_invocations import op_db
|
|
from functorch_additional_op_db import additional_op_db
|
|
from enum import Enum
|
|
import torch._functorch.top_operators_github_usage as top_ops
|
|
import pprint
|
|
import unittest
|
|
import enum
|
|
from torch.testing._internal.common_device_type import toleranceOverride
|
|
|
|
# Importing these files make modifications to the op_db that we need
|
|
import test_ops # noqa: F401
|
|
import test_vmap # noqa: F401
|
|
|
|
all_overridable = list(torch.overrides.get_testing_overrides().keys())
|
|
|
|
public_docs = [
|
|
(torch.nn.functional, 'torch.nn.functional', 'docs/source/nn.functional.rst'),
|
|
(torch.fft, 'torch.fft', 'docs/source/fft.rst'),
|
|
(torch.special, 'torch.special', 'docs/source/special.rst'),
|
|
(torch.linalg, 'torch.linalg', 'docs/source/linalg.rst'),
|
|
(torch, 'torch', 'docs/source/torch.rst'),
|
|
(torch.Tensor, 'torch.Tensor', 'docs/source/tensors.rst'),
|
|
]
|
|
|
|
# torch.abs, Tensor.abs, Tensor.abs_ are all considered to be different
|
|
|
|
|
|
def get_public_overridable_apis(pytorch_root='/raid/rzou/pt/debug-cpu'):
|
|
results = {}
|
|
all_overridable_apis = set(torch.overrides.get_testing_overrides().keys())
|
|
for module, module_name, src in public_docs:
|
|
with open(f'{pytorch_root}/{src}') as f:
|
|
lines = f.readlines()
|
|
# APIs eitehr begin with 4 spaces or ".. autofunction::"
|
|
api_lines1 = [line.strip() for line in lines if line.startswith(' ' * 4)]
|
|
api_lines2 = [line.strip()[len('.. autofunction:: '):]
|
|
for line in lines if line.startswith('.. autofunction::')]
|
|
lines = api_lines1 + api_lines2
|
|
lines = [line[7:] if line.startswith('Tensor.') else line for line in lines]
|
|
lines = [line for line in lines if hasattr(module, line)]
|
|
for line in lines:
|
|
api = getattr(module, line)
|
|
if api in all_overridable_apis:
|
|
results[f'{module_name}.{line}'] = api
|
|
return results
|
|
|
|
|
|
denylist = {
|
|
'torch.Tensor.data_ptr',
|
|
'torch.Tensor.dim',
|
|
'torch.Tensor.element_size',
|
|
'torch.Tensor.backward',
|
|
'torch.Tensor.as_strided',
|
|
'torch.Tensor.register_hook',
|
|
'torch.Tensor.record_stream',
|
|
'torch.Tensor.qscheme',
|
|
'torch.Tensor.ndimension',
|
|
'torch.Tensor.smm',
|
|
'torch.Tensor.sspaddmm',
|
|
'torch.Tensor.retain_grad',
|
|
'torch.Tensor.sparse_mask',
|
|
'torch.Tensor.sparse_dim',
|
|
'torch.Tensor.dense_dim',
|
|
'torch.Tensor.values',
|
|
'torch.Tensor.indices',
|
|
'torch.Tensor.numel',
|
|
'torch.Tensor.size',
|
|
'torch.Tensor.nelement',
|
|
'torch.Tensor.q_scale',
|
|
'torch.Tensor.q_zero_point',
|
|
'torch.Tensor.q_per_channel_scales',
|
|
'torch.Tensor.q_per_channel_zero_points',
|
|
'torch.Tensor.q_per_channel_axis',
|
|
'torch.Tensor.int_repr',
|
|
'torch.Tensor.to_sparse',
|
|
'torch.Tensor.is_inference',
|
|
'torch.Tensor.storage',
|
|
'torch.Tensor.storage_type',
|
|
}
|
|
|
|
|
|
def get_method_only_ops_we_care_about():
|
|
apis = get_public_overridable_apis()
|
|
result = []
|
|
for key in apis.keys():
|
|
if not key.startswith('torch.Tensor'):
|
|
continue
|
|
if key in denylist:
|
|
continue
|
|
api = key.split('.')[2]
|
|
# filter out in-place
|
|
if api.endswith('_'):
|
|
continue
|
|
if f'torch.{api}' not in apis.keys():
|
|
result.append(api)
|
|
return result
|
|
|
|
# Deduplicates torch.abs and Tensor.abs
|
|
|
|
|
|
def get_public_overridable_ops():
|
|
results = get_public_overridable_apis()
|
|
cpy = copy.deepcopy(results)
|
|
for key in cpy.keys():
|
|
if not key.startswith('torch.Tensor'):
|
|
continue
|
|
api = key.split('.')[2]
|
|
if f'torch.{api}' in results.keys():
|
|
del results[key]
|
|
return results
|
|
|
|
|
|
def get_public_overridable_outplace_ops():
|
|
results = get_public_overridable_ops()
|
|
cpy = copy.deepcopy(results)
|
|
for key in cpy.keys():
|
|
# NB: there are no dunder methods bcs we don't document those
|
|
if key.endswith('_'):
|
|
del results[key]
|
|
return results
|
|
|
|
|
|
def get_public_overridable_outplace_we_care_about():
|
|
results = get_public_overridable_outplace_ops()
|
|
cpy = copy.deepcopy(results)
|
|
for key in cpy.keys():
|
|
# quantization
|
|
if 'quant' in key or '.q_' in key:
|
|
del results[key]
|
|
|
|
# is_cpu, etc. It doesn't make sense to have OpInfos for these
|
|
if '.is_' in key:
|
|
del results[key]
|
|
|
|
if key in denylist and key in results:
|
|
del results[key]
|
|
return results
|
|
|
|
# e.g. nn.functional.softmax
|
|
|
|
|
|
def get_op(dotted_name):
|
|
names = dotted_name.split('.')
|
|
mod = torch
|
|
for name in names:
|
|
if not hasattr(mod, name):
|
|
return None
|
|
mod = getattr(mod, name)
|
|
return mod
|
|
|
|
# Maps function -> [OpInfo]
|
|
|
|
|
|
def get_ops_covered_by_opinfos():
|
|
ops = {}
|
|
|
|
def safe_append(dct, key, val):
|
|
if key in dct:
|
|
dct[key].append(val)
|
|
else:
|
|
dct[key] = [val]
|
|
|
|
for opinfo in op_db:
|
|
func_op = get_op(opinfo.name)
|
|
if func_op:
|
|
safe_append(ops, func_op, opinfo)
|
|
if opinfo.method_variant:
|
|
safe_append(ops, opinfo.method_variant, opinfo)
|
|
if opinfo.inplace_variant:
|
|
safe_append(ops, opinfo.inplace_variant, opinfo)
|
|
for alias in opinfo.aliases:
|
|
safe_append(ops, alias.op, opinfo)
|
|
return ops
|
|
|
|
|
|
factory_fns = {
|
|
'tensor', 'zeros', 'ones', 'randn', 'arange', 'rand', 'empty', 'randperm',
|
|
'linspace', 'logspace', 'hann_window', 'full', 'eye', 'blackman_window',
|
|
'barlett_window', 'randint', 'range',
|
|
}
|
|
|
|
|
|
def get_top_ops(torch_threshold, nn_fn_threshold, with_counts=False):
|
|
denylist = set({
|
|
# These are either not real "operators", factory functions
|
|
# that trivially work, or not-documented ops.
|
|
'load', 'no_grad', 'save', 'from_numpy',
|
|
'manual_seed', 'set_grad_enabled',
|
|
'set_default_tensor_type', 'set_num_threads',
|
|
'set_printoptions', 'numel',
|
|
'set_default_dtype', 'sparse_coo_tensor', 'set_rng_state',
|
|
'get_rng_state', 'get_default_dtype', 'initial_seed',
|
|
'get_num_threads', 'quantize_per_tensor',
|
|
'hann_window', 'is_tensor', 'as_tensor',
|
|
'equal', 'enable_grad', 'seed', 'is_storage',
|
|
'is_floating_point', 'nn.functional.torch',
|
|
'set_flush_denormal', 'set_num_interop_threads', 'dequantize',
|
|
'get_num_interop_threads', 'nn.functional.math',
|
|
'nn.functional.threshold_',
|
|
'nn.functional.selu_',
|
|
'nn.functional.elu_',
|
|
'nn.functional.rrelu_',
|
|
'nn.functional.leaky_relu_',
|
|
'nn.functional.hardtanh_',
|
|
'nn.functional.has_torch_function',
|
|
'nn.functional.has_torch_function_unary',
|
|
'nn.functional.has_torch_function_variadic',
|
|
'nn.functional.handle_torch_function',
|
|
'nn.functional.adaptive_max_pool1d_with_indices',
|
|
'nn.functional.adaptive_max_pool2d_with_indices',
|
|
'nn.functional.adaptive_max_pool3d_with_indices',
|
|
'nn.functional.fractional_max_pool2d_with_indices',
|
|
'nn.functional.fractional_max_pool3d_with_indices',
|
|
'is_complex',
|
|
'grad',
|
|
'quantize_per_channel',
|
|
'nn.functional.max_pool2d_with_indices',
|
|
'nn.functional.max_pool3d_with_indices',
|
|
'nn.functional.max_pool1d_with_indices',
|
|
'nn.functional.celu_',
|
|
'nn.functional.grad',
|
|
'nn.functional.relu_',
|
|
'nn.functional.boolean_dispatch',
|
|
'nn.functional.assert_int_or_pair',
|
|
'fft', # is namespace
|
|
})
|
|
|
|
torch_ops = top_ops.top_torch
|
|
nn_fn_ops = top_ops.get_nn_functional_top_list()
|
|
torch_ops = [op for op in torch_ops if op[0] not in denylist]
|
|
nn_fn_ops = [op for op in nn_fn_ops if op[0] not in denylist]
|
|
|
|
ops = torch_ops[:torch_threshold] + nn_fn_ops[:nn_fn_threshold]
|
|
|
|
# Now, sort by priority
|
|
ops.sort(reverse=True, key=lambda op: op[1])
|
|
if not with_counts:
|
|
ops = [op[0] for op in ops]
|
|
return ops
|
|
|
|
|
|
def get_ops_percentage(torch_threshold, nn_fn_threshold):
|
|
data = top_ops.top_torch + top_ops.get_nn_functional_top_list()
|
|
|
|
def get_num_usages(opname):
|
|
# Ignore this, this is heavily inflated
|
|
if opname == 't':
|
|
return 0
|
|
result = [op[1] for op in data if op[0] == opname]
|
|
assert len(result) == 1
|
|
return result[0]
|
|
|
|
# get all operators that are not in the denylist
|
|
all_ops = get_top_ops(999999, 999999)
|
|
total_op_usages = sum([get_num_usages(op) for op in all_ops])
|
|
|
|
# get subset of all operators
|
|
subset_ops = get_top_ops(torch_threshold, nn_fn_threshold)
|
|
subset_op_usages = sum([get_num_usages(op) for op in subset_ops])
|
|
return subset_op_usages / total_op_usages
|
|
|
|
|
|
def get_top_ops_not_covered_by_opinfo(torch_threshold=0, nn_fn_threshold=0):
|
|
ops = get_top_ops(torch_threshold, nn_fn_threshold)
|
|
|
|
ops_with_opinfo = []
|
|
for op in op_db:
|
|
ops_with_opinfo.append(op.name)
|
|
ops_with_opinfo.extend([op.name for op in op.aliases])
|
|
ops_with_opinfo = set(ops_with_opinfo)
|
|
|
|
result = [op for op in ops if op not in ops_with_opinfo]
|
|
result = [op for op in result if op not in denylist]
|
|
result = [op for op in result if op not in factory_fns]
|
|
return result
|
|
|
|
|
|
def get_covered_ops(ops_list, invert=False):
|
|
ops_covered_by_opinfo = get_ops_covered_by_opinfos()
|
|
overridable_outplace_ops = ops_list
|
|
results = {}
|
|
for key, op in overridable_outplace_ops.items():
|
|
cond = op in ops_covered_by_opinfo
|
|
if invert:
|
|
cond = not cond
|
|
if cond:
|
|
results[key] = op
|
|
return results
|
|
|
|
|
|
class Status(Enum):
|
|
Correct = 0
|
|
Fast = 1
|
|
|
|
|
|
tests = {
|
|
'test_vmap_exhaustive',
|
|
'test_op_has_batch_rule',
|
|
'test_vjp',
|
|
'test_vmapvjp',
|
|
'test_vmapvjp_has_batch_rule',
|
|
'test_jvp',
|
|
'test_vmapjvp',
|
|
}
|
|
|
|
|
|
def is_decorateinfo_skip_or_xfail(decorateinfo):
|
|
assert len(decorateinfo.decorators) == 1
|
|
actual_decorator = decorateinfo.decorators[0]
|
|
if isinstance(actual_decorator, toleranceOverride):
|
|
return False
|
|
if actual_decorator == unittest.expectedFailure:
|
|
return True
|
|
# Assume the rest are skips
|
|
return True
|
|
|
|
|
|
def get_all_tested_ops():
|
|
overridable_outplace_we_care_about = get_public_overridable_outplace_we_care_about()
|
|
op_to_opinfo = get_ops_covered_by_opinfos()
|
|
result = set({})
|
|
for op in get_covered_ops(overridable_outplace_we_care_about).values():
|
|
opinfos = op_to_opinfo[op]
|
|
for opinfo in opinfos:
|
|
result.add(opinfo.name)
|
|
return result
|
|
|
|
|
|
def get_skipped_or_xfailed_ops_for(test_name):
|
|
overridable_outplace_we_care_about = get_public_overridable_outplace_we_care_about()
|
|
op_to_opinfo = get_ops_covered_by_opinfos()
|
|
result = set({})
|
|
for op in get_covered_ops(overridable_outplace_we_care_about).values():
|
|
opinfos = op_to_opinfo[op]
|
|
for opinfo in opinfos:
|
|
for decorator in opinfo.decorators:
|
|
if not hasattr(decorator, 'test_name'):
|
|
continue
|
|
if decorator.test_name != test_name:
|
|
continue
|
|
if is_decorateinfo_skip_or_xfail(decorator):
|
|
result.add(opinfo.name)
|
|
return result
|
|
|
|
|
|
def get_statuses(for_subset=None, invert=False):
|
|
overridable_outplace_we_care_about = get_public_overridable_outplace_we_care_about()
|
|
if for_subset is not None:
|
|
overridable_outplace_we_care_about = {
|
|
k: v
|
|
for k, v in overridable_outplace_we_care_about.items()
|
|
# Removes "torch."
|
|
if k[6:] in for_subset
|
|
}
|
|
op_to_opinfo = get_ops_covered_by_opinfos()
|
|
result = {}
|
|
_ = get_covered_ops(overridable_outplace_we_care_about)
|
|
|
|
def get_covered_tests(op):
|
|
opinfos = op_to_opinfo[op]
|
|
result = copy.deepcopy(tests)
|
|
for opinfo in opinfos:
|
|
for decorator in opinfo.decorators:
|
|
if not hasattr(decorator, 'test_name'):
|
|
continue
|
|
if decorator.test_name in tests and decorator.test_name in result:
|
|
result.remove(decorator.test_name)
|
|
return result
|
|
|
|
def get_all_aliases(op):
|
|
opinfos = op_to_opinfo[op]
|
|
result = []
|
|
for opinfo in opinfos:
|
|
result.append(opinfo.name)
|
|
result.extend(opinfo.aliases)
|
|
return set(result)
|
|
|
|
for name, op in get_covered_ops(overridable_outplace_we_care_about).items():
|
|
successful_tests = get_covered_tests(op)
|
|
failed_tests = tests - successful_tests
|
|
result[name] = failed_tests if invert else successful_tests
|
|
return result
|
|
|
|
|
|
def transpose_statuses(for_subset=None, invert=False):
|
|
statuses = get_statuses(for_subset, invert=invert)
|
|
result = {}
|
|
for test in tests:
|
|
result[test] = set({})
|
|
for op, supported in statuses.items():
|
|
for test in supported:
|
|
result[test].add(op)
|
|
return result
|
|
|
|
|
|
overridable_apis = get_public_overridable_apis()
|
|
|
|
overridable_ops = get_public_overridable_ops()
|
|
|
|
overridable_outplace_ops = get_public_overridable_outplace_ops()
|
|
|
|
overridable_outplace_we_care_about = get_public_overridable_outplace_we_care_about()
|
|
|
|
tested_overridable_outplace_ops = get_covered_ops(overridable_outplace_we_care_about)
|
|
untested_overridable_outplace_ops = get_covered_ops(overridable_outplace_we_care_about, invert=True)
|
|
|
|
# print("List of OpInfos we need:")
|
|
# for key in untested_overridable_outplace_ops.keys():
|
|
# print(key)
|
|
# print("-" * 80)
|
|
# print("")
|
|
|
|
print(f'Overridable public APIs: {len(overridable_apis)}')
|
|
print(f'Overridable public ops: {len(overridable_ops)}')
|
|
print(f'Overridable public outplace ops: {len(overridable_outplace_ops)}')
|
|
print(f'Overridable public outplace ops we care about: {len(overridable_outplace_we_care_about)}')
|
|
print(f'OpInfo-tested overridable public outplace ops: {len(tested_overridable_outplace_ops)}')
|
|
|
|
|
|
def remove_torch(name):
|
|
assert name[:6] == 'torch.'
|
|
return name[6:]
|
|
|
|
|
|
def get_list_of_all_tests():
|
|
all_tests = list(tested_overridable_outplace_ops.keys())
|
|
return {remove_torch(test) for test in all_tests}
|
|
|
|
|
|
mytest = {
|
|
'test_vmap_exhaustive',
|
|
'test_op_has_batch_rule',
|
|
'test_vjp',
|
|
'test_vmapvjp',
|
|
'test_vmapvjp_has_batch_rule',
|
|
}
|
|
|
|
print('*' * 80)
|
|
all_tests = get_list_of_all_tests()
|
|
for test in mytest:
|
|
result = get_skipped_or_xfailed_ops_for(test)
|
|
diff = len(all_tests - result)
|
|
print(f'{test}: {diff}')
|
|
|
|
|
|
def get_jvp_coverage(subset=None):
|
|
# - number that support autograd
|
|
# - number that support forward_ad (in pytorch core)
|
|
# - number that support functorch.jvp
|
|
op_to_opinfo = get_ops_covered_by_opinfos()
|
|
ops_dct = tested_overridable_outplace_ops
|
|
if subset is not None:
|
|
ops_dct = {name: op for name, op in ops_dct.items()
|
|
if remove_torch(name) in subset}
|
|
supports_autograd_ops_dct = {name: op_to_opinfo[fn] for name, fn in ops_dct.items()
|
|
if op_to_opinfo[fn][0].supports_autograd}
|
|
supports_forwardad_ops_dct = {name: op_to_opinfo[fn] for name, fn in ops_dct.items()
|
|
if op_to_opinfo[fn][0].supports_forward_ad}
|
|
|
|
ops = {remove_torch(test) for test in list(ops_dct.keys())}
|
|
supports_autograd = {remove_torch(test)
|
|
for test in list(supports_autograd_ops_dct.keys())}
|
|
supports_forward_ad = {remove_torch(test)
|
|
for test in list(supports_forwardad_ops_dct.keys())}
|
|
assert supports_forward_ad.issubset(supports_autograd)
|
|
assert supports_autograd.issubset(ops)
|
|
|
|
failed_ops = get_skipped_or_xfailed_ops_for('test_jvp')
|
|
|
|
coverage = len(supports_forward_ad - failed_ops)
|
|
no_forward_ad = len(supports_autograd) - len(supports_forward_ad)
|
|
print(f'test_jvp, {coverage}, {no_forward_ad}, {len(ops)}')
|
|
|
|
|
|
get_jvp_coverage()
|
|
get_jvp_coverage(get_top_ops(100, 25))
|
|
for op in get_top_ops(100, 25):
|
|
print(op)
|
|
print('*' * 80)
|
|
|
|
# result = get_skipped_or_xfailed_ops_for('test_vmap_exhaustive')
|
|
# result = get_skipped_or_xfailed_ops_for('test_op_has_batch_rule')
|
|
# result = get_skipped_or_xfailed_ops_for('test_vjp')
|
|
# result = get_skipped_or_xfailed_ops_for('test_vmapvjp')
|
|
# result = get_skipped_or_xfailed_ops_for('test_vmapvjp_has_batch_rule')
|
|
# import pdb; pdb.set_trace()
|
|
|
|
statuses = transpose_statuses()
|
|
for test in tests:
|
|
print(f'{test} coverage {len(statuses[test])}')
|
|
|
|
method_only_ops = get_method_only_ops_we_care_about()
|
|
# for op in method_only_ops:
|
|
# print(f' {op},')
|
|
|
|
top_ops_not_covered_by_opinfo = get_top_ops_not_covered_by_opinfo(100, 25)
|
|
print('=' * 80)
|
|
for op in top_ops_not_covered_by_opinfo:
|
|
print(f'{op}, {top_ops.usage_count[op]}')
|
|
|
|
# print("top ops not covered by opinfo: ")
|
|
# top_ops_not_covered_by_opinfo = get_top_ops_not_covered_by_opinfo(200, 50)
|
|
# for op in top_ops_not_covered_by_opinfo:
|
|
# print(f'{op}, {top_ops.usage_count[op]}')
|
|
|
|
# print("top ops not covered by opinfo: ")
|
|
# top_ops_not_covered_by_opinfo = get_top_ops_not_covered_by_opinfo(220, 92)
|
|
# for op in top_ops_not_covered_by_opinfo:
|
|
# print(f'{op}, {top_ops.usage_count[op]}')
|
|
|
|
# print("top ops not covered by opinfo: ")
|
|
# top_ops_not_covered_by_opinfo = get_top_ops_not_covered_by_opinfo(999, 999)
|
|
# for op in top_ops_not_covered_by_opinfo:
|
|
# print(f'{op}, {top_ops.usage_count[op]}')
|
|
|
|
|
|
def remove_from_set(parent, to_remove):
|
|
for to_remove_elt in to_remove:
|
|
if to_remove_elt in parent:
|
|
parent.remove(to_remove_elt)
|
|
|
|
|
|
def print_coverage_info(th=100, nn=25):
|
|
print('=' * 80)
|
|
print(f"top {th}, {nn} coverage")
|
|
statuses = transpose_statuses(get_top_ops(th, nn), invert=True)
|
|
top_ops_not_covered_by_opinfo = get_top_ops_not_covered_by_opinfo(th, nn)
|
|
|
|
# testing problems
|
|
exemptions = {
|
|
'torch.nn.functional.dropout', # randomness
|
|
}
|
|
|
|
# Allowed exemptions
|
|
vmap_exemptions = {
|
|
'torch.randn_like', # randomness
|
|
'torch.rand_like', # randomness
|
|
'torch.allclose', # number output
|
|
'torch.unique', # dynamic
|
|
'torch.nonzero', # dynamic
|
|
'torch.masked_select', # dynamic
|
|
'torch.prod', # dynamic (backward)
|
|
'torch.norm', # norm with nuc is not commonly used; we support the other cases.
|
|
'torch.svd', # There isn't a bug, it is just nondeterministic so we can't test it.
|
|
'torch.nn.functional.embedding', # We support everything except the sparse option.
|
|
}
|
|
remove_from_set(statuses['test_vmap_exhaustive'], vmap_exemptions)
|
|
remove_from_set(statuses['test_vmapvjp'], vmap_exemptions)
|
|
remove_from_set(statuses['test_vmapvjp_has_batch_rule'], vmap_exemptions)
|
|
remove_from_set(statuses['test_op_has_batch_rule'], vmap_exemptions)
|
|
remove_from_set(statuses['test_vmapjvp'], vmap_exemptions)
|
|
for test in tests:
|
|
remove_from_set(statuses[test], exemptions)
|
|
|
|
print(f"total ops in set: {th + nn}")
|
|
print(f"tested by OpInfo: {th + nn - len(top_ops_not_covered_by_opinfo)}")
|
|
for test in tests:
|
|
if test in {'test_jvp', 'test_vmapjvp'}:
|
|
continue
|
|
print(f'{test} failing coverage {len(statuses[test])}')
|
|
|
|
# We don't care about these yet
|
|
del statuses['test_jvp']
|
|
del statuses['test_vmapjvp']
|
|
|
|
pprint.pprint(statuses)
|
|
|
|
|
|
def get_name_to_opinfo_map():
|
|
dct = {}
|
|
for op in (op_db + additional_op_db):
|
|
def add(name, op):
|
|
if name not in dct:
|
|
dct[name] = []
|
|
dct[name].append(op)
|
|
add(op.name, op)
|
|
for alias in op.aliases:
|
|
add(alias.name, op)
|
|
return dct
|
|
|
|
|
|
NAME_TO_OPINFO = get_name_to_opinfo_map()
|
|
|
|
|
|
class Support(enum.Enum):
|
|
NO = 0
|
|
YES = 1
|
|
UNKNOWN = 2
|
|
|
|
|
|
FACTORY_FNS = {
|
|
'tensor', 'zeros', 'ones', 'randn', 'arange', 'rand', 'empty', 'range',
|
|
'full', 'randperm', 'eye', 'randint', 'linspace', 'logspace',
|
|
}
|
|
|
|
VJP_EXEMPTIONS = {
|
|
'nn.functional.dropout', # not actually problem, randomness testing artifact
|
|
'nn.functional.dropout2d', # not actually problem, randomness testing artifact
|
|
'nn.functional.rrelu', # not actually problem, randomness testing artifact
|
|
'bernoulli', # not actually problem, randomness testing artifact
|
|
'normal', # not actually problem, randomness testing artifact
|
|
}
|
|
|
|
VMAP_EXEMPTIONS = {
|
|
'randn_like', # randomness
|
|
'rand_like', # randomness
|
|
'allclose', # number output
|
|
'unique', # dynamic
|
|
'nonzero', # dynamic
|
|
'masked_select', # dynamic
|
|
'prod', # dynamic (backward)
|
|
'norm', # norm with nuc is not commonly used; we support the other cases.
|
|
'svd', # There isn't a bug, it is just nondeterministic so we can't test it.
|
|
'nn.functional.embedding', # We support everything except the sparse option.
|
|
'nn.functional.dropout', # randomness
|
|
'nn.functional.dropout2d', # randomness
|
|
'bernoulli', # randomness
|
|
'multinomial', # randomness
|
|
'normal', # randomness
|
|
}
|
|
|
|
JVP_EXEMPTIONS = {
|
|
'nn.functional.dropout', # not actually problem, randomness testing artifact
|
|
'nn.functional.dropout2d', # not actually problem, randomness testing artifact
|
|
'nn.functional.rrelu', # not actually problem, randomness testing artifact
|
|
'normal', # not actually problem, randomness testing artifact
|
|
'bernoulli', # not actually problem, randomness testing artifact
|
|
}
|
|
|
|
|
|
class Operator:
|
|
def __init__(self, name):
|
|
self.name = name
|
|
self.opinfos = NAME_TO_OPINFO.get(name, None)
|
|
assert self.opinfos is None or len(self.opinfos) > 0
|
|
|
|
def has_opinfo(self):
|
|
return self.opinfos is not None
|
|
|
|
def __repr__(self):
|
|
return f'Operator("{self.name}")'
|
|
|
|
def __hash__(self):
|
|
return hash(self.name)
|
|
|
|
def no_opinfos_skip_test(self, test_name):
|
|
"""Returns NO if any opinfos have a skip or xfail for the test"""
|
|
if not self.has_opinfo():
|
|
return Support.UNKNOWN
|
|
for opinfo in self.opinfos:
|
|
for decorator in opinfo.decorators:
|
|
if not hasattr(decorator, 'test_name'):
|
|
continue
|
|
if decorator.test_name != test_name:
|
|
continue
|
|
if is_decorateinfo_skip_or_xfail(decorator):
|
|
return Support.NO
|
|
return Support.YES
|
|
|
|
def any_opinfo_attr(self, attr):
|
|
if not self.has_opinfo():
|
|
raise RuntimeError()
|
|
return any(getattr(opinfo, attr) for opinfo in self.opinfos)
|
|
|
|
def all_opinfo_attr(self, attr):
|
|
if not self.has_opinfo():
|
|
raise RuntimeError()
|
|
return all(getattr(opinfo, attr) for opinfo in self.opinfos)
|
|
|
|
def supports_vjp(self):
|
|
if self.name in FACTORY_FNS:
|
|
return Support.YES
|
|
if self.name in VJP_EXEMPTIONS:
|
|
return Support.YES
|
|
return self.no_opinfos_skip_test('test_vjp')
|
|
|
|
def supports_vmap(self):
|
|
if self.name in FACTORY_FNS:
|
|
return Support.YES
|
|
if self.name in VMAP_EXEMPTIONS:
|
|
return Support.YES
|
|
return self.no_opinfos_skip_test('test_vmap_exhaustive')
|
|
|
|
def supports_fast_vmap(self):
|
|
if self.name in FACTORY_FNS:
|
|
return Support.YES
|
|
if self.name in VMAP_EXEMPTIONS:
|
|
return Support.YES
|
|
return self.no_opinfos_skip_test('test_op_has_batch_rule')
|
|
|
|
def supports_vmapvjp(self):
|
|
if self.name in FACTORY_FNS:
|
|
return Support.YES
|
|
if self.name in VMAP_EXEMPTIONS:
|
|
return Support.YES
|
|
return self.no_opinfos_skip_test('test_vmapvjp')
|
|
|
|
def supports_fast_vmapvjp(self):
|
|
if self.name in FACTORY_FNS:
|
|
return Support.YES
|
|
if self.name in VMAP_EXEMPTIONS:
|
|
return Support.YES
|
|
return self.no_opinfos_skip_test('test_vmapvjp_has_batch_rule')
|
|
|
|
def supports_jvp(self):
|
|
if self.name in FACTORY_FNS:
|
|
return Support.YES
|
|
if self.name in JVP_EXEMPTIONS:
|
|
return Support.YES
|
|
if not self.has_opinfo():
|
|
return Support.UNKNOWN
|
|
if self.any_opinfo_attr('supports_autograd') and \
|
|
not self.all_opinfo_attr('supports_forward_ad'):
|
|
return Support.NO
|
|
return self.no_opinfos_skip_test('test_jvp')
|
|
|
|
def supports_jvpvjp(self):
|
|
if self.name in FACTORY_FNS:
|
|
return Support.YES
|
|
exemptions = {
|
|
# we have support (see OpInfo), testing artifact
|
|
'nn.functional.dropout2d',
|
|
'nn.functional.dropout',
|
|
# exception: we dont even support double backward for this
|
|
'nn.functional.hardswish',
|
|
'bernoulli', # this isn't differentiable
|
|
'normal', # not differentiable
|
|
}
|
|
if self.name in exemptions:
|
|
return Support.YES
|
|
return self.no_opinfos_skip_test('test_jvpvjp')
|
|
|
|
def _supports_vmapjvp_base(self, test):
|
|
if self.name in FACTORY_FNS:
|
|
return Support.YES
|
|
VMAPJVP_EXEMPTIONS = {
|
|
'prod', # dynamic (backward)
|
|
'nn.functional.batch_norm', # testing problem
|
|
'normal', # not actually problem, randomness testing artifact
|
|
'bernoulli', # not actually problem, randomness testing artifact
|
|
'nn.functional.dropout2d', # not actually problem, randomness testing artifact
|
|
'nn.functional.dropout', # not actually problem, randomness testing artifact
|
|
# Not a problem.
|
|
# It's just that the max_norm testing mutates inputs...
|
|
# (we have our own functorch variant of the OpInfo without max_norm)
|
|
'nn.functional.embedding',
|
|
}
|
|
if self.name in VMAPJVP_EXEMPTIONS:
|
|
return Support.YES
|
|
if not self.has_opinfo():
|
|
return Support.UNKNOWN
|
|
if self.any_opinfo_attr('supports_autograd') and \
|
|
not self.all_opinfo_attr('supports_forward_ad'):
|
|
return Support.NO
|
|
return self.no_opinfos_skip_test(test)
|
|
|
|
def supports_vmapjvp(self):
|
|
return self._supports_vmapjvp_base('test_vmapjvpall')
|
|
|
|
def supports_fast_vmapjvp(self):
|
|
return self._supports_vmapjvp_base('test_vmapjvpall_has_batch_rule')
|
|
|
|
|
|
class OperatorSet:
|
|
def __init__(self, operators):
|
|
self.data = set(operators)
|
|
|
|
@classmethod
|
|
def from_names(cls, names):
|
|
return OperatorSet([Operator(name) for name in names])
|
|
|
|
@classmethod
|
|
def from_top_ops_threshold(cls, torch_threshold, nn_fn_threshold):
|
|
names = get_top_ops(torch_threshold, nn_fn_threshold)
|
|
return cls.from_names(names)
|
|
|
|
@classmethod
|
|
def from_top125(cls):
|
|
return cls.from_top_ops_threshold(100, 25)
|
|
|
|
@classmethod
|
|
def from_top160(cls):
|
|
return cls.from_top_ops_threshold(107, 53)
|
|
|
|
@classmethod
|
|
def all(cls):
|
|
dct = get_public_overridable_outplace_we_care_about()
|
|
names = dct.keys()
|
|
names_sanitized = []
|
|
for n in names:
|
|
torch_tensor = 'torch.Tensor.'
|
|
torch_dot = 'torch.'
|
|
if n.startswith(torch_tensor):
|
|
names_sanitized.append(n[len(torch_tensor):])
|
|
elif n.startswith(torch_dot):
|
|
names_sanitized.append(n[len(torch_dot):])
|
|
else:
|
|
raise AssertionError()
|
|
return cls.from_names(names_sanitized)
|
|
|
|
def query(self, operator_method, filter=(Support.NO, Support.YES, Support.UNKNOWN)):
|
|
result = {}
|
|
for key in filter:
|
|
result[key] = set()
|
|
for op in self.data:
|
|
support_status = operator_method(op)
|
|
if support_status in filter:
|
|
result[support_status].add(op)
|
|
return result
|
|
|
|
def summary(self):
|
|
checks = [
|
|
'supports_vjp',
|
|
'supports_vmap',
|
|
'supports_fast_vmap',
|
|
'supports_vmapvjp',
|
|
'supports_fast_vmapvjp',
|
|
'supports_jvp',
|
|
'supports_vmapjvp',
|
|
'supports_fast_vmapjvp',
|
|
'supports_jvpvjp',
|
|
]
|
|
result = ['test, yes, no, unknown']
|
|
for check in checks:
|
|
accessor = getattr(Operator, check)
|
|
all_results = self.query(accessor)
|
|
yes_amt = len(all_results[Support.YES])
|
|
no_amt = len(all_results[Support.NO])
|
|
unknown_amt = len(all_results[Support.UNKNOWN])
|
|
result.append(f'{check}, {yes_amt}, {no_amt}, {unknown_amt}')
|
|
return '\n'.join(result)
|
|
|
|
|
|
opset = OperatorSet.all()
|
|
has_no_opinfo = opset.query(Operator.has_opinfo, (False,))
|
|
|
|
print("=" * 30 + " Summary " + "=" * 30)
|
|
print(f'% of usages on github: {get_ops_percentage(99999, 99999)}')
|
|
print(opset.summary())
|
|
|
|
# sanity checks
|
|
result = opset.query(Operator.supports_vjp, (Support.NO, Support.UNKNOWN))
|
|
# pprint.pprint(result)
|
|
|
|
print("=" * 30 + " Top 60 Summary " + "=" * 30)
|
|
print(f'% of usages on github: {get_ops_percentage(35, 25)}')
|
|
opset = OperatorSet.from_top_ops_threshold(35, 25)
|
|
# result = opset.query(Operator.supports_vmapjvp, (Support.NO, Support.UNKNOWN))
|
|
# pprint.pprint(result)
|
|
# result = opset.query(Operator.supports_jvp, (Support.NO, Support.UNKNOWN))
|
|
# pprint.pprint(result)
|
|
# kresult = opset.query(Operator.supports_jvpvjp, (Support.NO, Support.UNKNOWN))
|
|
# kpprint.pprint(result)
|
|
# result = opset.query(Operator.supports_vmapjvp, (Support.NO, Support.UNKNOWN))
|
|
# pprint.pprint(result)
|
|
# result = opset.query(Operator.supports_fast_vmapjvp, (Support.NO, Support.UNKNOWN))
|
|
# pprint.pprint(result)
|
|
# pprint.pprint(result)
|
|
print(opset.summary())
|
|
|
|
print("=" * 30 + " Top 125 Summary " + "=" * 30)
|
|
print(f'% of usages on github: {get_ops_percentage(100, 25)}')
|
|
opset = OperatorSet.from_top125()
|
|
# result = opset.query(Operator.supports_vmap, (Support.NO, Support.UNKNOWN))
|
|
# pprint.pprint(result)
|
|
# result = opset.query(Operator.supports_jvpvjp, (Support.NO, Support.UNKNOWN))
|
|
# pprint.pprint(result)
|
|
print("supports_vjp")
|
|
result = opset.query(Operator.supports_vjp, (Support.NO, Support.UNKNOWN))
|
|
pprint.pprint(result)
|
|
print("supports_jvp")
|
|
result = opset.query(Operator.supports_jvp, (Support.NO, Support.UNKNOWN))
|
|
pprint.pprint(result)
|
|
print("supports_vmapjvp")
|
|
result = opset.query(Operator.supports_vmapjvp, (Support.NO, Support.UNKNOWN))
|
|
pprint.pprint(result)
|
|
print("supports_jvpvjp")
|
|
result = opset.query(Operator.supports_jvpvjp, (Support.NO, Support.UNKNOWN))
|
|
pprint.pprint(result)
|
|
# result = opset.query(Operator.supports_fast_vmapjvp, (Support.NO, Support.UNKNOWN))
|
|
# pprint.pprint(result)
|
|
# pprint.pprint(result)
|
|
print(opset.summary())
|
|
|
|
# print("=" * 30 + " Top 160 Summary " + "=" * 30)
|
|
# opset = OperatorSet.from_top160()
|
|
# result = opset.query(Operator.supports_jvpvjp, (Support.NO, Support.UNKNOWN))
|
|
# pprint.pprint(result)
|
|
# print(opset.summary())
|
|
|
|
# Print list of everything in order
|
|
# all_ops = get_top_ops(999999, 999999, with_counts=True)
|
|
# for op, count in all_ops:
|
|
# print(f'{op}, {count}')
|