from cmath import e
from bangpy import tcp
from bangpy.tcp.runtime import TaskType
from bangpy.platform.bang_config import TARGET
import bangpy as bp
import numpy as np
target = "mlu290"
tcps = tcp.TCP(target)
#设置核信息
task_type = TaskType.UNION16
tcps.launch_cluster(task_type.value)
task_num = task_type.value * 4
tcps.launch_task(task_num, 1, 1)
#设置数据类型和数据个数
dtype = bp.float32
dtype_sz = dtype.bytes
length = 2**20
shape = (length, )
#设置硬件相关信息
compute_row = 32
nram_size = TARGET(target).nram_size
single_buffer_size = 64 * 128 // dtype.bytes * compute_row
task_id = tcps.taskId
#每次处理的数据量
data_calculated_each_time = single_buffer_size // dtype_sz
data_calculated_each_task = length // task_num
with tcps.if_scope(task_id == task_num - 1):
data_calculated_each_task = length // task_num + length % task_num
#每个task需要循环多少次
loop_num = tcps.Scalar(bp.int32, name="loop_num")
loop_num.assign(data_calculated_each_task // data_calculated_each_time)
# global的数据buffer
input = tcps.Buffer(shape=shape, name="input", dtype=dtype, scope="global")
target = tcps.Buffer(shape=shape, name="target", dtype=dtype, scope="global")
out = tcps.Buffer(shape=shape, name="output", dtype=dtype, scope="global")
reduction = tcps.SizeVar(name="reduction")
# 分块处理相关变量
start = tcps.Scalar(bp.int32, name="start")
start.assign(task_id * (length // task_num))
stop = tcps.Scalar(bp.int32, name="stop")
buffer_input = tcps.Buffer(shape=(2*data_calculated_each_time,), name="input_N", dtype=dtype, scope="nram")
buffer_target = tcps.Buffer(shape=(2*data_calculated_each_time,), name="target_N", dtype=dtype, scope="nram")
buffer_out = tcps.Buffer(shape=(2*data_calculated_each_time,), name="OUTPUT_N", dtype=dtype, scope="nram")
def compute(ou, inp, tar, i) :
tcps.log(inp[(i%2)*data_calculated_each_time : (i%2+1)*data_calculated_each_time], inp[(i%2)*data_calculated_each_time : (i%2+1)*data_calculated_each_time], high_precision = False)
tcps.log(ou[(i%2)*data_calculated_each_time : (i%2+1)*data_calculated_each_time], tar[(i%2)*data_calculated_each_time : (i%2+1)*data_calculated_each_time], high_precision = False)
tcps.subtract(ou[(i%2)*data_calculated_each_time : (i%2+1)*data_calculated_each_time], ou[(i%2)*data_calculated_each_time : (i%2+1)*data_calculated_each_time], inp[(i%2)*data_calculated_each_time : (i%2+1)*data_calculated_each_time])
tcps.multiply(ou[(i%2)*data_calculated_each_time : (i%2+1)*data_calculated_each_time], tar[(i%2)*data_calculated_each_time : (i%2+1)*data_calculated_each_time], ou[(i%2)*data_calculated_each_time : (i%2+1)*data_calculated_each_time])
## 第 0 次 gdram -> nram 传输
tcps.memcpy_async(buffer_input[:data_calculated_each_time], input[start : start + data_calculated_each_time])
tcps.memcpy_async(buffer_target[:data_calculated_each_time], target[start : start + data_calculated_each_time])
tcps.sync_ipu()
## 第 0 次 comopute
compute(buffer_out, buffer_input, buffer_target, 0)
## 第 1 次 gdram -> nram 传输
tcps.memcpy_async(buffer_input[data_calculated_each_time : data_calculated_each_time*2], input[start+ data_calculated_each_time : start + data_calculated_each_time*2])
tcps.memcpy_async(buffer_target[data_calculated_each_time : data_calculated_each_time*2], target[start+ data_calculated_each_time : start + data_calculated_each_time*2])
#对于每个task里面的每次buffer计算进行流水线处理
with tcps.for_range(begin = 0, end = loop_num-2, stage = 0) as i :
tcps.sync_ipu()
## 第 i 次 nram -> gdram 传输
tcps.memcpy_async(out[start+data_calculated_each_time*i : start+data_calculated_each_time*(i+1)], buffer_out[(i%2) * data_calculated_each_time: (i%2+1) * data_calculated_each_time])
## 第 i+1 次 compute
compute(buffer_out, buffer_input, buffer_target, i+1)
## 第 i+2 次 gdram -> nram 传输
tcps.memcpy_async(buffer_input[((i+2)%2)*data_calculated_each_time : ((i+2)%2 + 1)*data_calculated_each_time], input[start + (i+2)*data_calculated_each_time : start + (i+3)*data_calculated_each_time])
tcps.memcpy_async(buffer_target[((i+2)%2)*data_calculated_each_time : ((i+2)%2 + 1)*data_calculated_each_time], target[start + (i+2)*data_calculated_each_time : start + (i+3)*data_calculated_each_time])
tcps.sync_ipu()
## 第 n-1 次 nram -> gdram 传输
tcps.memcpy_async(out[start+data_calculated_each_time*(loop_num-2) : start+data_calculated_each_time * (loop_num-1)], buffer_out[((loop_num-2)%2) * data_calculated_each_time: ((loop_num-2)%2+1) * data_calculated_each_time])
## 第 n 次 compute
compute(buffer_out, buffer_input, buffer_target, loop_num-1)
## 第 n 次 nram -> gdram 传输
tcps.memcpy_async(out[start+data_calculated_each_time*(loop_num-1) : start+data_calculated_each_time * (loop_num)], buffer_out[((loop_num-1)%2) * data_calculated_each_time: ((loop_num-1)%2+1) * data_calculated_each_time])
#处理每个task的尾部数据
with tcps.if_scope(data_calculated_each_task % data_calculated_each_time != 0):
start.assign(start + data_calculated_each_time * loop_num)
stop.assign(start + data_calculated_each_task % data_calculated_each_time)
#data copy
tcps.memcpy(buffer_input[:stop-start], input[start:stop])
tcps.memcpy(buffer_target[:stop-start], target[start:stop])
tcps.log(buffer_input, buffer_input, high_precision = False)
tcps.log(buffer_out, buffer_target, high_precision = False)
tcps.subtract(buffer_out, buffer_out, buffer_input)
tcps.multiply(buffer_out, buffer_target, buffer_out)
tcps.memcpy(out[start:stop], buffer_out[:stop-start])
f = tcps.BuildBANG(inputs=[input, target, reduction], outputs=[out], kernel_name="kidivloss")
print(f.get_source())
######################################################################################
data_input = np.random.uniform(low=0, high=1, size=(length,))
data_target = np.random.uniform(low=0, high=1, size=(length,)).astype(dtype.as_numpy_dtype)
# data_input = (np.array([0.2]*length)).reshape(shape)
# data_target = (np.array([0.25]*length)).reshape(shape)
# print('data_input : ', data_input)
# print('data_target : ', data_target)
data_out = np.multiply(data_target,np.subtract(np.log(data_target), np.log(data_input)))
print(data_out)
print("-------------result----------------------------------------")
dev = bp.device(0)
data_input_dev = bp.Array(data_input.astype(dtype.as_numpy_dtype), dev)
data_target_dev = bp.Array(data_target.astype(dtype.as_numpy_dtype), dev)
data_out_dev = bp.Array(np.zeros(data_out.shape, dtype.as_numpy_dtype), dev)
reduction = 0
f(data_input_dev, data_target_dev, reduction, data_out_dev)
print(data_out_dev)
evaluator = f.time_evaluator(number=100, repeat=2, min_repeat_ms=0)
print(evaluator)
t = evaluator(data_input_dev, data_target_dev,reduction, data_out_dev).mean * 1e3
print(
"tutorial : %f ms"
% t
)
io_efficiency = (length * 3 * dtype_sz) / (2**30) / (t / 1000)
print("io_efficiency : ", io_efficiency)