| """Explicit native-math primitives; CPU remains a mathematical diagnostic."""
|
| from functools import lru_cache
|
| from pathlib import Path
|
| import torch
|
| from .kernel_build import stable_source
|
|
|
|
|
| @lru_cache(None)
|
| def extension():
|
| from torch.utils.cpp_extension import load
|
| return load(name='fidel_compatibility_math_reference_v1',
|
| sources=[str(stable_source(Path(__file__).with_name('cuda_math_reference.cu')))],
|
| extra_cuda_cflags=['-O3','--use_fast_math'],verbose=False)
|
|
|
|
|
| def unary(value,operation):
|
| value=value.float().contiguous()
|
| if value.is_cuda:
|
| if torch.cuda.get_device_capability(value.device)<(7,5):
|
| raise RuntimeError('Native tanh.approx reference currently requires SM75+; older-SM emulation remains unvalidated')
|
| return extension().unary(value,operation)
|
| functions=(torch.tanh,torch.tanh,torch.cos,torch.sin,torch.exp,torch.sigmoid)
|
| return functions[operation](value)
|
|
|
|
|
| def bf16_product(a,b,accumulator=None):
|
| """BF16 operands with FP32 accumulation; no intermediate BF16 output cast."""
|
| if a.dtype!=torch.bfloat16 or b.dtype!=torch.bfloat16:
|
| raise ValueError('BF16 matrix operands required')
|
| shape=(*a.shape[:-2],a.shape[-2],b.shape[-1])
|
| if a.is_cuda:
|
| left=a.reshape(-1,*a.shape[-2:]).contiguous()
|
| right=b.reshape(-1,*b.shape[-2:]).contiguous()
|
| if torch.cuda.get_device_capability(a.device)<(8,0):
|
| raise RuntimeError('Native BF16 MMA reference requires SM80+; older-SM emulation remains unvalidated')
|
| initial=None if accumulator is None else accumulator.reshape(-1,*shape[-2:]).contiguous()
|
| canonical=torch.cuda.get_device_capability(a.device)!=(10,0)
|
| return extension().mma(left,right,initial,canonical).reshape(shape)
|
| value=a.float()@b.float()
|
| return value if accumulator is None else accumulator+value
|
|
|
|
|
| def cuda_fma(a,b,c):
|
| a,b,c=torch.broadcast_tensors(a.float(),b.float(),c.float())
|
| return extension().fma(a.contiguous(),b.contiguous(),c.contiguous())
|
|
|
|
|
| def f32_batched_product(a,b,*,streams,stride):
|
| if a.dtype!=torch.float32 or b.dtype!=torch.float32 or a.shape[:-2]!=b.shape[:-2]:
|
| raise ValueError('Matching FP32 batch dimensions required')
|
| shape=(*a.shape[:-2],a.shape[-2],b.shape[-1])
|
| left=a.reshape(-1,*a.shape[-2:]).contiguous();right=b.reshape(-1,*b.shape[-2:]).contiguous()
|
| return extension().bmm_reference(left,right,streams,stride).reshape(shape)
|
|
|
