# Copyright (c) 2019 Graphcore Ltd. All rights reserved.
import sys
import popart

if len(sys.argv) != 2:
    raise RuntimeError("onnx model file name expected as argument")

model_file = sys.argv[1]

opts = popart.SessionOptions()
opts.logging = {"all": "TRACE"}
options = {"compileIPUCode": True, "numIPUs": 1}

# currently, with both Cpu and IpuModel, we have outstanding tasks
# T6384 and T6405, about the conv planner failing (after ~15 mins with Cpu)
device = popart.DeviceManager().createCpuDevice()
# createIpuModelDevice(options)

builder = popart.Builder(model_file)
graph_transformer = popart.GraphTransformer(builder.getModelProto())
graph_transformer.convertAllFixedPointInitializersToConstants()

# specific to the task, this output might need changing
output = builder.getOutputTensorIds()[0]
dataFlow = popart.DataFlow(1, {output: popart.AnchorReturnType("All")})

s = popart.Session(
    graph_transformer.getModelProto(), dataFlow=dataFlow, userOptions=opts
)
s.setDevice(device)
s.prepareDevice()