def test_repeat_simple_addition(self, repeat_count: int): """Test that a simple x = x + 1 repeated `repeat_count` times will produce x = `repeat_count` Args: repeat_count (int): Number of times to repeat. """ ir = pir.Ir() main = ir.main_graph() with main: one = pir.constant(0, pir.dtypes.int32) add_one = AddOne() add_one_graph = ir.create_graph(add_one, one) y = ops.repeat(add_one_graph, repeat_count, one, subgraph_in_to_parent_in={}) d2h = pir.d2h_stream(y.shape, pir.dtypes.int32, name="y_stream") ops.host_store(d2h, y) r_y = run_ir(ir, 1, d2h.tensor_id(), {}) assert r_y == repeat_count
def test_scaled_add_t_t(self, inplace): ir = pir.Ir() g = ir.main_graph() with g: X = pir.variable(np.ones((2, 2)), dtype=pir.float32) Y = pir.variable(np.ones((2, 2)), dtype=pir.float32) if inplace: ops.scaled_add_(X, Y, a=pir.constant(0.9), b=pir.variable(0.1)) else: ops.scaled_add(X, Y, a=pir.constant(0.9), b=pir.variable(0.1)) assert len(g.get_tensors()) == 5 assert len(g.get_variables()) == 3 if inplace: assert contains_op_of_type("ScaledAddLhsInplace", _ir.op.ScaledAddLhsInplaceOp, g) else: assert contains_op_of_type("ScaledAdd", _ir.op.ScaledAddOp, g)
def test_tensor_id_conflict(): ir = pir.Ir() main = ir.main_graph() with main: name0 = pir.variable(1, name="tensor").id name1 = pir.variable(1, name="tensor").id name2 = pir.constant(1, name="tensor").id assert name0 == "tensor" ids = [name0, name1, name2] assert len(ids) == len(set(ids))
def test_dampened_add_square(self): ir = pir.Ir() g = ir.main_graph() with g: a = pir.variable(1) b = pir.constant(2) c = ops.var_updates.accumulate_square_(a, b, 0.999) assert contains_op_of_type("Accumulate", _ir.op.AccumulateOp, g) op = g._pb_graph.getOps()[0] op.getAccumulationType() == _ir.AccumulationType.DampenedAddSquare
def test_modified(): ir = pir.Ir() g = ir.main_graph() with g, pir.in_sequence(): x = pir.variable(1) sg = ir.create_graph( lambda x: ops.var_updates.accumulate_(x, pir.constant(1)), x) ops.call(sg, x) # type: ignore # Store x x_non_modify_stream = pir.d2h_stream(x.shape, x.dtype) ops.host_store(x_non_modify_stream, x) info = ops.call_with_info(sg, x) info.set_op_input_modified(x) x_modifiy_stream = pir.d2h_stream(x.shape, x.dtype) ops.host_store(x_modifiy_stream, x) ir = ir._pb_ir dataFlow = popart.DataFlow(batchesPerStep=1, anchorTensors={ x_non_modify_stream.tensor_id(): popart.AnchorReturnType("All"), x_modifiy_stream.tensor_id(): popart.AnchorReturnType("All"), }) ir.setDataFlow(dataFlow) opts = ir.getSessionOptions() opts.useHostCopyOps = True opts.enableExplicitMainLoops = True opts.aliasZeroCopy = True opts.explicitRecomputation = True ir.updateVertices() session = popart.InferenceSession.fromIr( ir=ir, deviceInfo=tu.create_test_device()) session.prepareDevice() # Create buffers for anchors anchors = session.initAnchorArrays() # Run the model stepio = popart.PyStepIO(inputs={}, outputs=anchors) session.weightsFromHost() session.run(stepio) assert anchors[x_non_modify_stream.tensor_id()] == 1 assert anchors[x_modifiy_stream.tensor_id()] == 2
def test_add(self): ir = pir.Ir() g = ir.main_graph() with g: a = pir.variable(1) b = pir.constant(2) c = ops.var_updates.copy_var_update_(a, b) assert len(g.get_tensors()) == 3 assert len(g.get_variables()) == 1 assert contains_op_of_type("CopyVarUpdate", _ir.op.CopyVarUpdateOp, g)
def test_tensor_fns(self): ir = pir.Ir() g = ir.main_graph() with g: a = pir.variable(1) a_1 = a.copy_to_ipu(1, 0) a_0 = a_1.copy_to_ipu(0) c = pir.constant(2) c_1 = c.copy_to_ipu(1, 0) c_0 = c_1.copy_to_ipu(0)
def test_mean(self): ir = pir.Ir() g = ir.main_graph() with g: a = pir.variable(1) b = pir.constant(2) step = pir.variable(0) c = ops.var_updates.accumulate_mean_(a, b, step) assert contains_op_of_type("Accumulate", _ir.op.AccumulateOp, g) op = g._pb_graph.getOps()[0] op.getAccumulationType() == _ir.AccumulationType.Mean
def test_add(self): ir = pir.Ir() g = ir.main_graph() with g: a = pir.variable(1) b = pir.constant(2) c = ops.var_updates.accumulate_(a, b) assert len(g.get_tensors()) == 3 assert len(g.get_variables()) == 1 assert contains_op_of_type("Accumulate", _ir.op.AccumulateOp, g) op = g._pb_graph.getOps()[0] op.getAccumulationType() == _ir.AccumulationType.Add
def test_remote_load_graph(self, use_offset: bool, tensor_shape: Tuple[int, ...], repeats: int, tensor_dtype: dtype, inplace: bool) -> None: """Test that the graph is correct when using the remote load op Args: use_offset (bool): Whether or not to use offset tensor_shape (Tuple[int, ...]): The shape of the tensor to be loaded repeats (int): The number of tensors potentially stored in the buffer tensor_dtype (dtype): The type of the tensors to be loaded inplace (bool): Whether or not to use the inplace version of the op """ ir = pir.Ir() g = ir.main_graph() with g: t = pir.variable( np.random.rand(*tensor_shape).astype(tensor_dtype.as_numpy())) if use_offset: offset = pir.constant([1], name='offset') # With this option the graph should contain # 1. t # 2. offset # 3. out n_tensors = 3 else: offset = None # With this option the graph should contain # 1. t # 2. out n_tensors = 2 rbh = RemoteBufferHandle(remote_buffer_id=1, tensor_shape=tensor_shape, tensor_dtype=tensor_dtype, repeats=repeats) op = ops.remote_load if not inplace else ops.remote_load_ op(t, offset, rbh) assert len(g.get_tensors()) == n_tensors # Only t is a variable assert len(g.get_variables()) == 1 type_string = "RemoteLoad" if not inplace else "RemoteLoadInplace" pb_type = _ir.op.exchange.RemoteLoadOp if not inplace else _ir.op.exchange.RemoteLoadInplaceOp assert contains_op_of_type(type_string, pb_type, g) # Clean-up so that the RemoteBufferHandle gets reset RemoteBufferHandle._buffers = {}
def test_from_pb_type(self): """Test the from_pb_tensor returns the correct python type""" ir = pir.Ir() main = ir.main_graph() with main: a = pir.variable(1) c = pir.constant(2) assert isinstance(a, Variable) new_a = pir.Tensor._from_pb_tensor(a._pb_tensor) assert isinstance(new_a, Variable) assert isinstance(c, Constant) new_c = pir.Tensor._from_pb_tensor(c._pb_tensor) assert isinstance(new_c, Constant)
def test_adam_wd_updater(self): ir = pir.Ir() g = ir.main_graph() with g: w = pir.variable(1, name='w') m = pir.variable(1, name='m') v = pir.variable(2, name='v') wd = pir.constant(0.2, name='wd') updater = ops.var_updates.adam_updater(m, v, weight=w, weight_decay=wd) assert len(g.get_tensors()) == 5 assert contains_op_of_type("AdamUpdater", _ir.op.AdamUpdaterOp, g)
def test_adam_wd_updater_invalid(self): ir = pir.Ir() g = ir.main_graph() with g: m = pir.variable(1, name='m') v = pir.variable(2, name='v') t = pir.variable(1, name='t') wd = pir.constant(0.2, name='wd') with pytest.raises(ValueError) as excinfo: updater = ops.var_updates.adam_updater(m, v, time_step=t, weight_decay=wd) message = str(excinfo.value) assert "Weight decay requires weight to be not None." in message
def test_adam_bias_wd_updater(self): ir = pir.Ir() g = ir.main_graph() with g: m = pir.variable(1, name='m') v = pir.variable(2, name='v') w = pir.variable(1, name='w') t = pir.variable(2, name='t') wd = pir.constant(0.2, name='wd') b1 = 0.9 b2 = 0.99 updater = ops.var_updates.adam_updater(m, v, w, t, wd, b1, b2) assert len(g.get_tensors()) == 6 assert contains_op_of_type("AdamUpdater", _ir.op.AdamUpdaterOp, g)
def build_model( data: Dict[str, np.array]) -> Tuple[_ir.Ir, Dict[str, DeviceToHostStream]]: """Build a model for storing and loading tensors from the remote buffer. Args: data(Dict[str, np.array]) : Dict of the data to be stored and loaded from the remote buffer Returns: (tuple): tuple containing: ir._pb_ir (_ir.Ir): The underlying IR d2h_streams (Dict[str, DeviceToHostStream]): The output streams """ ir = pir.Ir() main = ir.main_graph() with main: # Placeholder for tensor ids tensors = {} # Create variable tensors from the data for name in data.keys(): tensors[name] = pir.variable(data[name], name=name) # Placeholder for device to host streams d2h_streams = {} # Store and load the first tensor without specifying the remote buffer handle or offset ops.remote_store(t=tensors["store_in_1"]) tensors["load_out_1"] = ops.remote_load(t=tensors["load_in_1"]) tensors["load_out_1_inplace"] = ops.remote_load_( t=tensors["load_in_1_inplace"]) # Anchor the input tensors to the load operator d2h_streams = make_anchor(d2h_streams, tensors, "load_in_1") d2h_streams = make_anchor(d2h_streams, tensors, "load_in_1_inplace") # Anchor the output tensors of the load operator d2h_streams = make_anchor(d2h_streams, tensors, "load_out_1") d2h_streams = make_anchor(d2h_streams, tensors, "load_out_1_inplace") # Store and load the second and third tensor using the same buffer id # Buffer 1 should already be assigned implicitly, so we chose a different id rbh = RemoteBufferHandle( remote_buffer_id=42, tensor_shape=tensors["store_in_2"]._pb_tensor.info.shape(), tensor_dtype=dtype.as_dtype( tensors["store_in_2"]._pb_tensor.info.data_type_lcase()), repeats=2) # Index starts at 0 offset_tensor_2 = pir.constant(0, name="offset_2") offset_tensor_3 = pir.constant(1, name="offset_3") ops.remote_store(t=tensors["store_in_2"], offset=offset_tensor_2, remote_buffer_handle=rbh) ops.remote_store(t=tensors["store_in_3"], offset=offset_tensor_3, remote_buffer_handle=rbh) tensors["load_out_2"] = ops.remote_load(t=tensors["load_in_2"], offset=offset_tensor_2, remote_buffer_handle=rbh) tensors["load_out_3_inplace"] = ops.remote_load_( t=tensors["load_in_3_inplace"], offset=offset_tensor_3, remote_buffer_handle=rbh) # Anchor the input tensors to the load operator d2h_streams = make_anchor(d2h_streams, tensors, "load_in_2") d2h_streams = make_anchor(d2h_streams, tensors, "load_in_3_inplace") # Anchor the output tensors of the load operator d2h_streams = make_anchor(d2h_streams, tensors, "load_out_2") d2h_streams = make_anchor(d2h_streams, tensors, "load_out_3_inplace") return ir._pb_ir, d2h_streams
def build_model( ) -> Tuple[_ir.Ir, pir.HostToDeviceStream, pir.DeviceToHostStream, pir. DeviceToHostStream, pir.DeviceToHostStream, pir. DeviceToHostStream, np.ndarray, np.ndarray]: ir = pir.Ir() main = ir.main_graph() with main: x_h2d = pir.h2d_stream(_IN_SHAPE, pir.float32, name="x_stream") x = ops.host_load(x_h2d, "x") W_data = np.random.normal(0, 0.1, _WEIGHT_SHAPE).astype(np.float32) b_data = np.zeros(_BIAS_SHAPE, dtype=np.float32) W = pir.variable(W_data, name="W") b = pir.variable(b_data, name="b") lin = Linear() lin_graph = ir.create_graph(lin, x, out_features=_OUT_FEATURES) lin_call_info = ops.call_with_info(lin_graph, x, subgraph_in_to_parent_in={ lin.W: W, lin.b: b }) y = lin_call_info.get_output_tensors()[0] assert y.shape == _OUT_SHAPE y_d2h = pir.d2h_stream(y.shape, y.dtype, name="x_stream") ops.host_store(y_d2h, y) lin_bwd_info = pir.transforms.autodiff.autodiff(lin_graph) lin_bwd_graph = lin_bwd_info.graph with main: grad_seed = pir.constant(np.ones(_OUT_SHAPE, np.float32)) tensors_required_for_bwd = pir.transforms.autodiff.get_expected_forward_inputs_from_call( lin_call_info, lin_bwd_info) lin_bwd_call_info = ops.call_with_info( lin_bwd_graph, grad_seed, subgraph_in_to_parent_in=tensors_required_for_bwd) ##### Extract parent graph x_grad, W_grad, b_grad expected_outputs = lin_bwd_info.expected_outputs x_grad, W_grad, b_grad = None, None, None sg_x = lin_call_info.op_in_to_subgraph_in_tensor(x) sg_W = lin_call_info.op_in_to_subgraph_in_tensor(W) sg_b = lin_call_info.op_in_to_subgraph_in_tensor(b) def get_grad_tensor_in_main_graph_from_fwdgrad_expected_connection( ec: pir.transforms.autodiff.ExpectedConnection) -> pir.Tensor: # If (t, FwdGrad) appears at index i in expected_outputs, it is # guaranteed that t’ (the grad of t) appears at output index i in the # grad graph. sg_out_idx = expected_outputs.index(ec) op_out_idx = lin_bwd_call_info.subgraph_in_to_op_in_index(sg_out_idx) parent_grad = lin_bwd_call_info.get_op_output_tensor(op_out_idx) return parent_grad for ec in expected_outputs: # Should always be the case for expected_outputs assert ec.connection_type == pir.transforms.autodiff.ExpectedConnectionType.FwdGrad sg_fwd_tensor = ec.fwd_tensor if sg_fwd_tensor == sg_x: x_grad = get_grad_tensor_in_main_graph_from_fwdgrad_expected_connection( ec) elif sg_fwd_tensor == sg_W: W_grad = get_grad_tensor_in_main_graph_from_fwdgrad_expected_connection( ec) elif sg_fwd_tensor == sg_b: b_grad = get_grad_tensor_in_main_graph_from_fwdgrad_expected_connection( ec) assert x_grad is not None assert W_grad is not None assert b_grad is not None # HostStore grads and collect d2h streams def host_store_and_return_d2h_stream( grad: pir.Tensor) -> pir.DeviceToHostStream: with main: d2h = pir.d2h_stream(grad.shape, grad.dtype, name=grad.name + "_stream") ops.host_store(d2h, grad) return d2h x_grad_d2h = host_store_and_return_d2h_stream(x_grad) W_grad_d2h = host_store_and_return_d2h_stream(W_grad) b_grad_d2h = host_store_and_return_d2h_stream(b_grad) assert x_grad_d2h is not None assert W_grad_d2h is not None assert b_grad_d2h is not None return ir._pb_ir, x_h2d, y_d2h, x_grad_d2h, W_grad_d2h, b_grad_d2h, W_data, b_data
def test_remote_store_graph(self, use_offset: bool, use_remote_buffer_id: bool, use_rbh: bool, tensor_shape: Tuple[int, ...], repeats: int, tensor_dtype: dtype) -> None: """Test that the graph is correct when using the remote store op. Args: use_offset (bool): Whether or not to use offset use_remote_buffer_id (bool): Whether or not to set the remote buffer_id use_rbh (bool): Whether or not to specify the remote buffer handle tensor_shape (Tuple[int, ...]): The shape of the tensor to be stored repeats (int): The number of tensors to potentially store in the buffer tensor_dtype (dtype): The type of the tensors to be stored """ ir = pir.Ir() g = ir.main_graph() with g: t = pir.variable( np.random.rand(*tensor_shape).astype(tensor_dtype.as_numpy())) if use_offset: offset = pir.constant([1], name='offset') # With this option the graph should contain # 1. t # 2. offset n_tensors = 2 else: offset = None # With this option the graph should contain # 1. t n_tensors = 1 remote_buffer_id = 1 if use_remote_buffer_id else -1 if remote_buffer_id == -1: with pytest.raises(NotImplementedError): _ = RemoteBufferHandle(remote_buffer_id=remote_buffer_id, tensor_shape=tensor_shape, tensor_dtype=tensor_dtype, repeats=repeats) # Clean-up so that the RemoteBufferHandle gets reset RemoteBufferHandle._buffers = {} return if use_rbh: rbh = RemoteBufferHandle(remote_buffer_id=remote_buffer_id, tensor_shape=tensor_shape, tensor_dtype=tensor_dtype, repeats=repeats) else: rbh = None ops.remote_store(t, offset, rbh) assert len(g.get_tensors()) == n_tensors # Only t is a variable assert len(g.get_variables()) == 1 assert contains_op_of_type("RemoteStore", _ir.op.exchange.RemoteStoreOp, g) # Clean-up so that the RemoteBufferHandle gets reset RemoteBufferHandle._buffers = {}
# Copyright (c) 2021 Graphcore Ltd. All rights reserved. import numpy as np import popart.ir as pir import popart.ir.ops as ops import popart # Creating a model with popart.ir ir = pir.Ir() main = ir.main_graph() with main: a = pir.variable(3, dtype=pir.int8, name="variable_a") b = pir.constant(1, dtype=pir.int8, name="constant_b") # addition o = a + b # host store o_d2h = pir.d2h_stream(o.shape, o.dtype, name="output_stream") ops.host_store(o_d2h, o) dataFlow = popart.DataFlow( batchesPerStep=1, anchorTensors={o_d2h.tensor_id(): popart.AnchorReturnType("All")}) ir = ir._pb_ir ir.setDataFlow(dataFlow) opts = ir.getSessionOptions() opts.useHostCopyOps = True opts.enableExplicitMainLoops = True ir.updateVertices() ir.setIsPrepared()
def test_subgraph(): class ScaleNShift(pir.Module): def __init__(self): self.W: pir.Tensor = None self.b: pir.Tensor = None def build(self, x: pir.Tensor, out_features: int, bias: bool = True) -> pir.Tensor: self.W = pir.subgraph_input((x.shape[-1], out_features), pir.float32, "W") y = ops.mul(x, self.W) if bias: self.b = pir.subgraph_input((out_features, ), pir.float32, "b") y = y + self.b return y ir = pir.Ir() main = ir.main_graph() with main: h2d = pir.h2d_stream((16, 16), pir.dtypes.float32) x = ops.host_load(h2d, "x") W = pir.variable(np.random.normal(0, 0.1, (16, 16)), name="W") b = pir.variable(np.zeros(16), name="b", dtype=pir.dtypes.float32) ss = ScaleNShift() ss_graph = ir.create_graph(ss, x, out_features=16) call_info = ops.call_with_info(ss_graph, x, subgraph_in_to_parent_in={ ss.W: W, ss.b: b }) y = call_info.get_output_tensors()[0] d2h = pir.d2h_stream(y.shape, y.dtype) ops.host_store(d2h, y) assert len(ss_graph.get_input_tensors()) == 3 assert len(ss_graph.get_output_tensors()) == 1 ss_bwd_info = pir.transforms.autodiff.autodiff(ss_graph) # Check an additional output has been added to the fwd graph. assert len(ss_graph.get_output_tensors()) == 2 bwd_graph = ss_bwd_info.graph assert isinstance(bwd_graph, pir.Graph) assert len(ss_bwd_info.expected_inputs) == len( bwd_graph.get_input_tensors()) assert len(ss_bwd_info.expected_outputs) == len( bwd_graph.get_output_tensors()) for op in bwd_graph._pb_graph.getOps(): grad_ops = (_ir.op.SumOp, _ir.op.MulArg0GradOp, _ir.op.MulArg1GradOp, _ir.op.AddArg0GradOp, _ir.op.AddArg1GradOp) assert isinstance(op, grad_ops) with main: grad_seed = pir.constant(np.ones((16, 16), np.float32)) activations = pir.transforms.autodiff.get_expected_forward_inputs_from_call( call_info, ss_bwd_info) grads = ops.call(bwd_graph, grad_seed, subgraph_in_to_parent_in=activations) assert len(grads) == len(ss_bwd_info.expected_outputs)