(f:Callable)
| 40 | |
| 41 | def nir_instr(nc=1, bs=lambda: None, intrins=None, srcs=None, has_def=True, df=None, also=lambda: None, **contents): |
| 42 | def dec(f:Callable): |
| 43 | @functools.wraps(f) |
| 44 | def wrapper(*args, **kwargs) -> mesa.nir_def: |
| 45 | (ba:=inspect.signature(f).bind(*args, **kwargs)).apply_defaults() |
| 46 | def go(g): return g(**{nm: ba.arguments[nm] for nm in inspect.signature(g).parameters}) if callable(g) else g |
| 47 | |
| 48 | instr = f(*args, **kwargs) |
| 49 | if has_def: mesa.nir_def_init(instr.contents.instr, instr.contents._def, go(nc), go(bs)) |
| 50 | for k, v in go(intrins or {}).items(): |
| 51 | idx = mesa.nir_intrinsic_infos[instr.contents.intrinsic].index_map[g(f"NIR_INTRINSIC_{k}")] |
| 52 | assert idx > 0, "invalid intrinsic. mesa version mismatch?" |
| 53 | instr.contents.const_index[idx - 1] = go(v) |
| 54 | for i, src in enumerate(go(srcs or [])): ctypes.cast(instr.contents.src, ctypes.POINTER(mesa.nir_src))[i] = go(src) |
| 55 | for k,v in {k:vcomp for k,v in contents.items() if (vcomp:=go(v)) is not None}.items(): setattr(instr.contents, k, go(v)) |
| 56 | mesa.nir_builder_instr_insert(ba.arguments['b'], instr.contents.instr) |
| 57 | go(also) |
| 58 | return instr.contents._def if has_def else (mesa.nir_def() if df is None else go(df)) |
| 59 | return wrapper |
| 60 | return dec |
| 61 | |
| 62 | @nir_instr(nc=1, bs=lambda src: src.bit_size, exact=lambda b:b.exact, fp_fast_math=lambda b:b.fp_fast_math) |
no outgoing calls