Method _get_kernel_code

test/amd/test_integration.py:78–93

(self, op_fn)

  """Test that we can parse tinygrad kernel disassembly."""

  def _get_kernel_code(self, op_fn) -> bytes:
    from tinygrad import Tensor
    from tinygrad.helpers import Target
    from tinygrad.codegen import to_program
    from tinygrad.renderer.llvmir import AMDLLVMRenderer
    from tinygrad.runtime.support.elf import elf_loader
    from tinygrad.uop.ops import Ops

    result = op_fn(Tensor)
    linear = result.schedule_linear()
    sink_items = [call for call in linear.src if call.src[0].op == Ops.SINK]
    assert len(sink_items) > 0, "No SINK in schedule"
    renderer = AMDLLVMRenderer(Target("AMD", arch='gfx1100'))
    prg = to_program(sink_items[0].src[0], renderer)
    lib = renderer.compiler.compile(prg.src[3].arg)
    return next(s.content for s in elf_loader(lib)[1] if s.name == ".text")

  def test_simple_add_kernel(self):
    """Generate a simple add kernel from tinygrad and verify disassembly."""