|
| 1 | +import sys |
| 2 | +import os |
| 3 | + |
| 4 | +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) |
| 5 | + |
| 6 | +import torch |
| 7 | +import infinicore |
| 8 | +from infinicore.ops.gemm import gemm as ic_gemm |
| 9 | +from framework.base import BaseOperatorTest, TensorSpec, TestCase |
| 10 | +from framework.tensor import TensorInitializer |
| 11 | +from framework.runner import GenericTestRunner |
| 12 | + |
| 13 | +# ============================================================================== |
| 14 | +# Operator-specific configuration |
| 15 | +# ============================================================================== |
| 16 | + |
| 17 | +# Test cases format: (operation_mode, nbatch, m, n, k, a_strides, b_strides, c_strides) |
| 18 | +# If nbatch is None: a_shape=(m, k), b_shape=(k, n), c_shape=(m, n) |
| 19 | +# If nbatch is provided: a_shape=(nbatch, m, k), b_shape=(nbatch, k, n), c_shape=(nbatch, m, n) |
| 20 | +# Aligned with test/infiniop/gemm.py shapes/strides and per-case alpha/beta |
| 21 | +# Each item: (alpha, beta, operation_mode, nbatch, m, n, k, a_strides, b_strides, c_strides) |
| 22 | +_TEST_CASES_DATA = [ |
| 23 | + # (1) alpha=1.0, beta=0.0, a=(1,2048), b=(2048,2048), c=(1,2048) |
| 24 | + (1.0, 0.0, TestCase.BOTH, None, 1, 2048, 2048, None, None, None), |
| 25 | + # (2) alpha=1.0, beta=0.0, a=(2,4,2048), b=(2,2048,2048), c=(2,4,2048) |
| 26 | + (1.0, 0.0, TestCase.BOTH, 2, 4, 2048, 2048, None, None, None), |
| 27 | + # (3) alpha=1.0, beta=0.0, strided (4096,1) |
| 28 | + (1.0, 0.0, TestCase.BOTH, None, 1, 2048, 2048, (4096, 1), (4096, 1), (4096, 1)), |
| 29 | + # (4) alpha=1.0, beta=1.0, only meaningful for IN_PLACE (needs existing C) |
| 30 | + (1.0, 1.0, TestCase.IN_PLACE, None, 6, 2560, 2048, (2048, 1), (1, 2048), (2560, 1)), |
| 31 | + # (5) alpha=1.0/8.0, beta=0.0, a=(4,48,64), b=(4,64,6), c=(4,48,6) |
| 32 | + (1.0 / 8.0, 0.0, TestCase.BOTH, 4, 48, 6, 64, None, None, None), |
| 33 | +] |
| 34 | + |
| 35 | + |
| 36 | +def parse_test_cases(data): |
| 37 | + """ |
| 38 | + Parse gemm test case data according to format: |
| 39 | + (operation_mode, nbatch, m, n, k, a_strides, b_strides, c_strides) |
| 40 | + """ |
| 41 | + alpha = data[0] |
| 42 | + beta = data[1] |
| 43 | + operation_mode = data[2] |
| 44 | + nbatch = data[3] |
| 45 | + m, n, k = data[4], data[5], data[6] |
| 46 | + a_strides = data[7] if len(data) > 7 else None |
| 47 | + b_strides = data[8] if len(data) > 8 else None |
| 48 | + c_strides = data[9] if len(data) > 9 else None |
| 49 | + |
| 50 | + # Determine shapes based on batch dimension |
| 51 | + if nbatch is None: |
| 52 | + a_shape = (m, k) |
| 53 | + b_shape = (k, n) |
| 54 | + c_shape = (m, n) |
| 55 | + else: |
| 56 | + a_shape = (nbatch, m, k) |
| 57 | + b_shape = (nbatch, k, n) |
| 58 | + c_shape = (nbatch, m, n) |
| 59 | + |
| 60 | + # Create input specifications |
| 61 | + inputs = [] |
| 62 | + |
| 63 | + # Tensor a |
| 64 | + if a_strides is not None: |
| 65 | + inputs.append(TensorSpec.from_strided_tensor(a_shape, a_strides)) |
| 66 | + else: |
| 67 | + inputs.append(TensorSpec.from_tensor(a_shape)) |
| 68 | + |
| 69 | + # Tensor b |
| 70 | + if b_strides is not None: |
| 71 | + inputs.append(TensorSpec.from_strided_tensor(b_shape, b_strides)) |
| 72 | + else: |
| 73 | + inputs.append(TensorSpec.from_tensor(b_shape)) |
| 74 | + |
| 75 | + # Output tensor |
| 76 | + if c_strides is not None: |
| 77 | + output = TensorSpec.from_strided_tensor( |
| 78 | + c_shape, |
| 79 | + c_strides, |
| 80 | + init_mode=TensorInitializer.ONES if beta != 0.0 else TensorInitializer.RANDOM, |
| 81 | + ) |
| 82 | + else: |
| 83 | + output = TensorSpec.from_tensor( |
| 84 | + c_shape, |
| 85 | + init_mode=TensorInitializer.ONES if beta != 0.0 else TensorInitializer.RANDOM, |
| 86 | + ) |
| 87 | + |
| 88 | + return TestCase(operation_mode, inputs, output, alpha=alpha, beta=beta) |
| 89 | + |
| 90 | + |
| 91 | +# Parse test cases |
| 92 | +_TEST_CASES = [parse_test_cases(data) for data in _TEST_CASES_DATA] |
| 93 | + |
| 94 | +# Data types |
| 95 | +_TENSOR_DTYPES = [infinicore.float16, infinicore.bfloat16, infinicore.float32] |
| 96 | + |
| 97 | +# Tolerance |
| 98 | +_TOLERANCE_MAP = { |
| 99 | + infinicore.float16: {"atol": 0, "rtol": 1e-2}, |
| 100 | + infinicore.float32: {"atol": 0, "rtol": 1e-3}, |
| 101 | + infinicore.bfloat16: {"atol": 0, "rtol": 5e-2}, |
| 102 | +} |
| 103 | + |
| 104 | + |
| 105 | +class OpTest(BaseOperatorTest): |
| 106 | + """GEMM test with simplified test case parsing |
| 107 | +
|
| 108 | + Note: We test default alpha=1.0 and beta=0.0 so it should match torch.matmul. |
| 109 | + """ |
| 110 | + |
| 111 | + def __init__(self): |
| 112 | + super().__init__("Gemm") |
| 113 | + |
| 114 | + def get_test_cases(self): |
| 115 | + return _TEST_CASES |
| 116 | + |
| 117 | + def get_tensor_dtypes(self): |
| 118 | + return _TENSOR_DTYPES |
| 119 | + |
| 120 | + def get_tolerance_map(self): |
| 121 | + return _TOLERANCE_MAP |
| 122 | + |
| 123 | + def torch_operator(self, a, b, out=None, **kwargs): |
| 124 | + alpha = kwargs.get("alpha", 1.0) |
| 125 | + beta = kwargs.get("beta", 0.0) |
| 126 | + mm = torch.matmul(a, b) |
| 127 | + if out is None: |
| 128 | + return mm.mul(alpha) |
| 129 | + out.mul_(beta) |
| 130 | + out.add_(mm, alpha=alpha) |
| 131 | + return out |
| 132 | + |
| 133 | + def infinicore_operator(self, a, b, out=None, **kwargs): |
| 134 | + alpha = kwargs.get("alpha", 1.0) |
| 135 | + beta = kwargs.get("beta", 0.0) |
| 136 | + if out is None: |
| 137 | + return ic_gemm(a, b, alpha=alpha, beta=beta) |
| 138 | + return ic_gemm(a, b, alpha=alpha, beta=beta, out=out) |
| 139 | + |
| 140 | + |
| 141 | +def main(): |
| 142 | + """Main entry point""" |
| 143 | + runner = GenericTestRunner(OpTest) |
| 144 | + runner.run_and_exit() |
| 145 | + |
| 146 | + |
| 147 | +if __name__ == "__main__": |
| 148 | + main() |
0 commit comments