Commit e0047df
committed
feat: lower pad, eye and triangular_solve in the native Expr compiler
Closes the top-level Expr op cluster on the Expr op-coverage checklist
(#188). Two new IR opcodes — `pad` (113) and `linalg_solve_triangular`
(114) — plus a dedicated `eye` materialise-as-constant clause.
- `pad` lowers to `mx::pad` ("constant" mode) — operands [input,
pad_value]; iattrs [axes, lows, highs]. Interior dilation (> 0)
raises matching Emily.Backend.pad/4, which has the same MLX
constraint.
- `eye` follows the same shape-static-creation pattern as iota:
Nx.eye on Nx.BinaryBackend at lower time, materialised as a
captured const. Nx.eye already handles rank > 2 by broadcasting the
identity onto the trailing two axes, matching
Emily.Backend.eye/2.
- `triangular_solve` decomposes Nx's four `transform_a` × `left_side`
combinations into transposes around the bare
`linalg_solve_triangular` opcode (operands [a, b]; iattrs [[upper]]),
mirroring Emily.Backend.triangular_solve/4. The C++ dispatcher
routes `mx::linalg::solve_triangular` to the CPU stream per call
(MLX's solver is CPU-only), matching c_src/ops/linalg.cpp.
Probe drops from 15 → 12 misses. Tests cover pad on 1-D / 2-D with
symmetric and asymmetric padding (plus the interior > 0 raise), eye
2-D and batched (rank-3 with leading batch broadcast), and
triangular_solve across all four `transform_a` × `left_side` cases
plus the upper-triangular variant.1 parent 981bac2 commit e0047df
4 files changed
Lines changed: 235 additions & 2 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
66 | 66 | | |
67 | 67 | | |
68 | 68 | | |
| 69 | + | |
| 70 | + | |
| 71 | + | |
| 72 | + | |
| 73 | + | |
| 74 | + | |
| 75 | + | |
| 76 | + | |
| 77 | + | |
| 78 | + | |
| 79 | + | |
69 | 80 | | |
70 | 81 | | |
71 | 82 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
207 | 207 | | |
208 | 208 | | |
209 | 209 | | |
| 210 | + | |
| 211 | + | |
| 212 | + | |
| 213 | + | |
| 214 | + | |
| 215 | + | |
| 216 | + | |
| 217 | + | |
| 218 | + | |
| 219 | + | |
| 220 | + | |
210 | 221 | | |
211 | 222 | | |
212 | | - | |
| 223 | + | |
213 | 224 | | |
214 | 225 | | |
215 | 226 | | |
| |||
765 | 776 | | |
766 | 777 | | |
767 | 778 | | |
| 779 | + | |
| 780 | + | |
| 781 | + | |
| 782 | + | |
| 783 | + | |
| 784 | + | |
| 785 | + | |
| 786 | + | |
| 787 | + | |
| 788 | + | |
| 789 | + | |
| 790 | + | |
| 791 | + | |
| 792 | + | |
| 793 | + | |
| 794 | + | |
| 795 | + | |
| 796 | + | |
| 797 | + | |
| 798 | + | |
| 799 | + | |
| 800 | + | |
| 801 | + | |
| 802 | + | |
| 803 | + | |
768 | 804 | | |
769 | 805 | | |
770 | 806 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
188 | 188 | | |
189 | 189 | | |
190 | 190 | | |
191 | | - | |
| 191 | + | |
| 192 | + | |
| 193 | + | |
| 194 | + | |
| 195 | + | |
| 196 | + | |
| 197 | + | |
| 198 | + | |
| 199 | + | |
192 | 200 | | |
193 | 201 | | |
194 | 202 | | |
| |||
829 | 837 | | |
830 | 838 | | |
831 | 839 | | |
| 840 | + | |
| 841 | + | |
| 842 | + | |
| 843 | + | |
| 844 | + | |
| 845 | + | |
| 846 | + | |
| 847 | + | |
| 848 | + | |
| 849 | + | |
832 | 850 | | |
833 | 851 | | |
834 | 852 | | |
| |||
945 | 963 | | |
946 | 964 | | |
947 | 965 | | |
| 966 | + | |
| 967 | + | |
| 968 | + | |
| 969 | + | |
| 970 | + | |
| 971 | + | |
| 972 | + | |
| 973 | + | |
| 974 | + | |
| 975 | + | |
| 976 | + | |
| 977 | + | |
| 978 | + | |
| 979 | + | |
| 980 | + | |
| 981 | + | |
| 982 | + | |
| 983 | + | |
| 984 | + | |
| 985 | + | |
| 986 | + | |
| 987 | + | |
| 988 | + | |
| 989 | + | |
| 990 | + | |
| 991 | + | |
| 992 | + | |
| 993 | + | |
| 994 | + | |
| 995 | + | |
| 996 | + | |
| 997 | + | |
| 998 | + | |
| 999 | + | |
| 1000 | + | |
| 1001 | + | |
| 1002 | + | |
| 1003 | + | |
| 1004 | + | |
| 1005 | + | |
| 1006 | + | |
| 1007 | + | |
| 1008 | + | |
| 1009 | + | |
| 1010 | + | |
| 1011 | + | |
| 1012 | + | |
| 1013 | + | |
| 1014 | + | |
| 1015 | + | |
| 1016 | + | |
| 1017 | + | |
| 1018 | + | |
| 1019 | + | |
| 1020 | + | |
| 1021 | + | |
| 1022 | + | |
| 1023 | + | |
| 1024 | + | |
| 1025 | + | |
| 1026 | + | |
| 1027 | + | |
| 1028 | + | |
| 1029 | + | |
| 1030 | + | |
| 1031 | + | |
948 | 1032 | | |
949 | 1033 | | |
950 | 1034 | | |
| |||
1364 | 1448 | | |
1365 | 1449 | | |
1366 | 1450 | | |
| 1451 | + | |
| 1452 | + | |
| 1453 | + | |
| 1454 | + | |
| 1455 | + | |
| 1456 | + | |
| 1457 | + | |
1367 | 1458 | | |
1368 | 1459 | | |
1369 | 1460 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
350 | 350 | | |
351 | 351 | | |
352 | 352 | | |
| 353 | + | |
| 354 | + | |
| 355 | + | |
| 356 | + | |
| 357 | + | |
| 358 | + | |
| 359 | + | |
| 360 | + | |
| 361 | + | |
| 362 | + | |
| 363 | + | |
| 364 | + | |
| 365 | + | |
| 366 | + | |
| 367 | + | |
| 368 | + | |
| 369 | + | |
| 370 | + | |
| 371 | + | |
| 372 | + | |
| 373 | + | |
| 374 | + | |
| 375 | + | |
| 376 | + | |
| 377 | + | |
| 378 | + | |
| 379 | + | |
| 380 | + | |
| 381 | + | |
| 382 | + | |
| 383 | + | |
| 384 | + | |
| 385 | + | |
| 386 | + | |
| 387 | + | |
| 388 | + | |
| 389 | + | |
| 390 | + | |
| 391 | + | |
| 392 | + | |
| 393 | + | |
| 394 | + | |
| 395 | + | |
| 396 | + | |
| 397 | + | |
| 398 | + | |
| 399 | + | |
| 400 | + | |
| 401 | + | |
| 402 | + | |
| 403 | + | |
| 404 | + | |
| 405 | + | |
| 406 | + | |
| 407 | + | |
| 408 | + | |
| 409 | + | |
| 410 | + | |
| 411 | + | |
| 412 | + | |
| 413 | + | |
| 414 | + | |
| 415 | + | |
| 416 | + | |
| 417 | + | |
| 418 | + | |
| 419 | + | |
| 420 | + | |
| 421 | + | |
| 422 | + | |
| 423 | + | |
| 424 | + | |
| 425 | + | |
| 426 | + | |
| 427 | + | |
| 428 | + | |
| 429 | + | |
| 430 | + | |
| 431 | + | |
| 432 | + | |
| 433 | + | |
| 434 | + | |
| 435 | + | |
| 436 | + | |
| 437 | + | |
| 438 | + | |
| 439 | + | |
| 440 | + | |
| 441 | + | |
| 442 | + | |
| 443 | + | |
| 444 | + | |
| 445 | + | |
| 446 | + | |
| 447 | + | |
353 | 448 | | |
354 | 449 | | |
355 | 450 | | |
| |||
0 commit comments