jit_primitives.py 571 B

1234567891011121314151617
  1. """Primitive functions that interface with the Modelverse JIT."""
  2. import modelverse_kernel.primitives as primitive_functions
  3. # pylint: disable=I0011,C0103
  4. def get_jit_enabled(**kwargs):
  5. """Checks if the JIT is enabled."""
  6. result, = yield [("CNV", [kwargs['mvk'].jit.jit_enabled])]
  7. raise primitive_functions.PrimitiveFinished(result)
  8. def set_jit_enabled(a, **kwargs):
  9. """Enables or disables the JIT."""
  10. val, result = yield [("RV", [a]), ("CN", [])]
  11. kwargs['mvk'].jit.set_jit_enabled(val)
  12. raise primitive_functions.PrimitiveFinished(result)