| 72 | } |
| 73 | |
| 74 | void time_binary_ops() { |
| 75 | int M = 1000, N = 100, K = 10; |
| 76 | auto condition = mx::random::randint(0, 2, {M, N, K}); |
| 77 | auto a = mx::random::uniform({M, N, K}); |
| 78 | auto b = mx::random::uniform({M, N, K}); |
| 79 | auto device = mx::default_device(); |
| 80 | mx::eval(a, b); |
| 81 | |
| 82 | TIME(mx::add, a, b, device); |
| 83 | TIME(mx::subtract, a, b, device); |
| 84 | TIME(mx::multiply, a, b, device); |
| 85 | TIME(mx::divide, a, b, device); |
| 86 | TIME(mx::maximum, a, b, device); |
| 87 | TIME(mx::minimum, a, b, device); |
| 88 | TIME(mx::where, condition, a, b, device); |
| 89 | |
| 90 | condition = mx::array({true}); |
| 91 | b = mx::random::uniform({1}); |
| 92 | mx::eval(b); |
| 93 | TIMEM("scalar", mx::add, a, b, device); |
| 94 | TIMEM("vector-scalar", mx::subtract, a, b, device); |
| 95 | TIMEM("scalar-vector", mx::subtract, b, a, device); |
| 96 | TIMEM("scalar", mx::multiply, a, b, device); |
| 97 | TIMEM("vector-scalar", mx::divide, a, b, device); |
| 98 | TIMEM("scalar-vector", mx::divide, b, a, device); |
| 99 | TIMEM("scalar-vector", mx::where, condition, a, b, device); |
| 100 | |
| 101 | condition = mx::broadcast_to(mx::array({true}), {1000, 100}); |
| 102 | a = mx::broadcast_to(mx::random::uniform({1}), {1000, 100}); |
| 103 | b = mx::broadcast_to(mx::random::uniform({1}), {1000, 100}); |
| 104 | mx::eval(a, b); |
| 105 | TIMEM("scalar-scalar broadcast", mx::add, a, b, device); |
| 106 | TIMEM("scalar-scalar broadcast", mx::subtract, a, b, device); |
| 107 | TIMEM("scalar-scalar broadcast", mx::multiply, a, b, device); |
| 108 | TIMEM("scalar-scalar broadcast", mx::divide, a, b, device); |
| 109 | TIMEM("scalar-scalar broadcast", mx::where, condition, a, b, device); |
| 110 | } |
| 111 | |
| 112 | void time_strided_ops() { |
| 113 | int M = 50, N = 50, O = 50, P = 50; |