def build_model(weights_data: np.array, input_shape: Tuple[int]
) -> Tuple[_ir.Ir, HostToDeviceStream, DeviceToHostStream]:
"""Build the model using popart.ir API.
Args:
weights_data (np.array): The (non-streamed) data of the weights
input_shape (tuple): The shape of the streamed input tensor
Returns:
(tuple): tuple containing:
ir._pb_ir(_ir.Ir): The underlying IR
t_in_h2d(HostToDeviceStream): The input stream of t_in
t_out_d2h (DeviceToHostStream): The output stream of t_out
"""
ir = pir.Ir()
main = ir.main_graph()
with main:
weights = pir.variable(weights_data, name="weights")
# Load t_in from host
t_in_h2d = pir.h2d_stream(input_shape, pir.float32, name="t_in_stream")
# Operations on IPU 0
with pir.virtual_graph(0):
t_in = ops.host_load(t_in_h2d, "t_in")
t_1 = ops.matmul(t_in, weights)
# Copy to IPU 1
t_1_c = ops.ipu_copy(t_1, 1)
# Operations on IPU 1
with pir.virtual_graph(1):
t_2 = ops.gelu(t_1_c)
# Copy to IPU 2
t_2_c = ops.ipu_copy(t_2, 2)
# Operations on IPU 2
with pir.virtual_graph(2):
t_out = ops.softmax(t_2_c, axis=1)
t_out_d2h = pir.d2h_stream(t_out.shape,
pir.float32,
name="t_out_stream")
ops.host_store(t_out_d2h, t_out)
return ir._pb_ir, t_in_h2d, t_out_d2h
def test_fn_with_no_producer(self):
ir = pir.Ir()
g = ir.main_graph()
with g:
a = pir.variable(1)
a_1 = ops.ipu_copy(a, 1, 0)
assert len(g.get_variables()) == 1
assert len(g.get_tensors()) == 2
assert contains_op_of_type("IpuCopy", _ir.op.IpuCopyOp, g)
def test_fn_with_no_producer_error(self):
ir = pir.Ir()
g = ir.main_graph()
with g:
a = pir.variable(1)
with pytest.raises(ValueError) as excinfo:
a_1 = ops.ipu_copy(a, 1)
msg = str(excinfo.value)
assert "Could not infer virtual graph" in msg
assert "Please specify `source`" in msg
def test_fn_with_producer(self):
ir = pir.Ir()
g = ir.main_graph()
with g:
a = pir.variable(1)
with pir.virtual_graph(0):
a_0 = a + 1
ir._pb_ir.logIr()
a_1 = ops.ipu_copy(a_0, 1)
assert len(g.get_variables()) == 1
assert len(g.get_tensors()) == 4
assert contains_op_of_type("IpuCopy", _ir.op.IpuCopyOp, g)
def test_fn_producer_without_virtual_graph(self):
ir = pir.Ir()
g = ir.main_graph()
with g:
a = pir.variable(1)
# None will disable annotations of virtual_graph. Not recommended
# but ipu_copy should throw an error if it happens.
with pir.virtual_graph(None):
a_0 = a + 1
with pytest.raises(ValueError) as excinfo:
a_1 = ops.ipu_copy(a_0, 1)
msg = str(excinfo.value)
assert "Could not infer virtual graph" in msg
def copy_to_ipu(self,
destination: int,
source: Optional[int] = None) -> 'Tensor':
"""
Copies a Tensor to a virtual graph.
Args:
destination (int):
Ipu for the tensor to be copied to.
source (Optional[int]):
Ipu for the tensor to be copied from.
By default, the source will be taken from the producer of the tensor.
If the tensor does not have a producer a source MUST be provided.
"""
import popart.ir.ops as ops
return ops.ipu_copy(self, destination, source)
def copy_to_ipu(self, dst: int, src: int) -> 'Tensor':
"""Returns ops.ipu_copy(self, dst, src).
Must provide a src value."""
import popart.ir.ops as ops
return ops.ipu_copy(self, dst, src)