Method test_copies_2_devs

test/backend/test_graph.py:189–209

(self)

  @needs_second_gpu
  def test_copies_2_devs(self):
    self.skip_if_not_multigraph()
    d0, d1 = Device.DEFAULT, f"{Device.DEFAULT}:1"
    b0 = [make_buffer(d0, fill=True) for _ in range(3)]
    b1 = [make_buffer(d1, fill=True)]
    c: dict[Buffer,UOp] = {}

    calls = [
      UOp(Ops.COPY).call(get_buf_uop(b1[0],c), get_buf_uop(b0[0],c), metadata=()),
      get_ast(d0, 2).call(get_buf_uop(b0[2],c), get_buf_uop(b0[0],c), get_buf_uop(b0[1],c), metadata=()),
    ]

    out = [b1[0], b0[2]]
    zero_bufs(out)
    run_schedule(calls)
    expected = {buf: np.frombuffer(buf.as_memoryview(), np.int32).copy() for buf in b0 + b1}

    for _ in range(RUN_CNT):
      zero_bufs(out)
      make_graph(Device[d0].graph, calls)([], {})
      for buf in b0 + b1: np.testing.assert_equal(expected[buf], np.frombuffer(buf.as_memoryview(), np.int32))

  def test_graph_offset_bufs(self):
    self.skip_if_not_multigraph()