def test_th_code():
src = '''
@kernel
def main(in_ptr, out_ptr, th):
y = get_local_id(0)
x = get_local_id(1)
index = 16*y + x
in_v = in_ptr[index]
out_v = 255 if in_v > th else 0
out_ptr[index] = out_v
'''
comp = CompilerDriver(8)
main_object = comp.run(src)
code = [InstrAdapter(x) for x in main_object.code]
pattern = '''
imm r2,0
lid r4,r2
imm r2,1
lid r5,r2
imm r2,16
mul r6,r2,r4
add r2,r6,r5
add r4,r3,r2
memr r3,r4
imm r4,255
imm r5,0
cmp r3,r0
mov r0,r5
mov {GT} r0,r4
add r3,r1,r2
memw r3,r0
'''.strip()
assert match_code(code, pattern)
def test_logic_variable2():
''' Test if a delayed condition captures the correct version of the tested variable. '''
src = '''
@kernel
def main(b):
b_gt = b > 0
d = 4 if b_gt else b
b += 1
c = 3 if b_gt else 1
'''
comp = CompilerDriver(8)
main_object = comp.run(src)
code = [InstrAdapter(x) for x in main_object.code]
pattern = '''
imm r1,0
imm r2,4
cmp r0,r1
mov r3,r0
mov {GT} r3,r2
imm r2,1
add r4,r0,r2
imm r2,3
imm r5,1
cmp r0,r1
mov r0,r5
mov {GT} r0,r2
'''.strip()
assert match_code(code, pattern)
def compile_code(src_code, args, block_size, nr_reg = 8, no_sequencer=True): ''' Compile a kernel and return a codegen function. ''' comp = CompilerDriver(nr_reg, no_sequencer=True) main_object = comp.run(src_code) patched_object = Compiler.patch_arguments_before_run(main_object, [args[argname] for argname in main_object.arguments]) def codegen_func(code, block_size, args): for x in patched_object.code: yield InstrAdapter(x, use_reg_wrapper=True) code = Code() code.set_generator(codegen_func, block_size, args) return code
def test_2D_const_list(): ''' Test code generation with 2D constants list. Note that this looks easy to optimise, but it is not so simple, because q is only know right before execution. ''' from blip.code.BlipCompiler import CompilerDriver, Compiler src = ''' @kernel def main(q): a = [[1, 2, 3], [4, 5, 6]] acc = 0 for i in range(2): for j in range(3): acc += q*a[i][j] return acc ''' comp = CompilerDriver(16, no_sequencer=True) main_object = comp.run(src) patched_object = Compiler.patch_arguments_before_run(main_object, [41]) code = [InstrAdapter(x) for x in patched_object.code] pattern = ''' imm r0,41 imm r1,0 imm r2,0 imm r3,0 imm r4,1 mul r5,r0,r4 add r4,r1,r5 imm r1,1 imm r5,2 mul r6,r0,r5 add r5,r4,r6 imm r4,2 imm r6,3 mul r7,r0,r6 add r6,r5,r7 imm r5,1 imm r7,0 imm r8,4 mul r9,r0,r8 add r8,r6,r9 imm r6,1 imm r9,5 mul r10,r0,r9 add r9,r8,r10 imm r8,2 imm r10,6 mul r11,r0,r10 add r0,r9,r11 mov r9,r0 '''.strip() assert match_code(code, pattern)
def test_const_array_access(): src = ''' @kernel def main(): r = [1, 2, 3] b = r[1] ''' comp = CompilerDriver(8) main_object = comp.run(src) code = [InstrAdapter(x) for x in main_object.code] pattern = ''' imm r0,1 imm r1,2 '''.strip() assert match_code(code, pattern)
def test_patch_arguments(): from blip.code.BlipCompiler import CompilerDriver, Compiler src = ''' @kernel def main(q): return q + 1 ''' comp = CompilerDriver(8) main_object = comp.run(src) patched_object = Compiler.patch_arguments_before_run(main_object, [41]) code = [InstrAdapter(x) for x in patched_object.code] pattern = ''' imm r0,41 imm r1,1 add r2,r0,r1 mov _, r2 '''.strip() assert match_code(code, pattern)
def test_logic_variable():
src = '''
@kernel
def main(b):
a = b > 0
c = 3 if a else 1
'''
comp = CompilerDriver(8)
main_object = comp.run(src)
code = [InstrAdapter(x) for x in main_object.code]
pattern = '''
imm r1,0
imm r2,3
imm r3,1
cmp r0,r1
mov r0,r3
mov {GT} r0,r2
'''.strip()
assert match_code(code, pattern)
def test_delayed_memderef(): src = ''' @kernel def main(q): b = 3 a = q[b] b = 5 c = a ''' comp = CompilerDriver(8) main_object = comp.run(src) code = [InstrAdapter(x) for x in main_object.code] pattern = ''' imm r0,3 add r2,r1,r0 memr r0,r2 imm r1,5 '''.strip() assert match_code(code, pattern)
def test_memderef_inside_condassign():
src = '''
@kernel
def main():
ind = 3
z = 1
b = z[4] if z > 0 else 0
'''
comp = CompilerDriver(8)
main_object = comp.run(src)
code = [InstrAdapter(x) for x in main_object.code]
pattern = '''
imm r0,3
imm r1,1
imm r2,0
imm r3,4
memr_imm r4,5
imm r5,0
cmp r1,r2
mov r1,r5
mov {GT} r1,r4
'''.strip()
assert match_code(code, pattern)
def test_conv_code_noseq():
''' Test a simple 3x1 convolution. '''
src = '''
@kernel
def main(in_ptr, out_ptr, th):
coeff = [-1, 0, 1]
y = get_local_id(0)
x = get_local_id(1)
index = 16*y
acc = 0
current_x = x - 1
for i in range(3):
in_v = in_ptr[index+current_x]
acc += (coeff[i] * in_v if current_x < 16 else 0) if current_x >= 0 else 0
current_x += 1
out_ptr[index+x] = acc
'''
comp = CompilerDriver(16, no_sequencer=True)
main_object = comp.run(src)
code = [InstrAdapter(x) for x in main_object.code]
pattern = '''
imm r0,0
lid r3,r0
imm r0,1
lid r4,r0
imm r0,16
mul r5,r0,r3
imm r0,0
imm r3,1
sub r6,r4,r3
imm r3,0
add r7,r5,r6
add r8,r2,r7
memr r7,r8
imm r8,0
imm r9,16
imm r10,-1
mul r11,r10,r7
imm r7,0
cmp r6,r9
mov r9,r7
mov {LT} r9,r11
imm r7,0
cmp r6,r8
mov r8,r7
mov {GE} r8,r9
add r7,r0,r8
imm r0,1
add r8,r6,r0
imm r0,1
add r6,r5,r8
add r9,r2,r6
memr r6,r9
imm r9,0
imm r10,16
imm r11,0
mul r12,r11,r6
imm r6,0
cmp r8,r10
mov r10,r6
mov {LT} r10,r12
imm r6,0
cmp r8,r9
mov r9,r6
mov {GE} r9,r10
add r6,r7,r9
imm r7,1
add r9,r8,r7
imm r7,2
add r8,r5,r9
add r10,r2,r8
memr r2,r10
imm r8,0
imm r10,16
imm r11,1
mul r12,r11,r2
imm r2,0
cmp r9,r10
mov r10,r2
mov {LT} r10,r12
imm r2,0
cmp r9,r8
mov r8,r2
mov {GE} r8,r10
add r2,r6,r8
imm r6,1
add r8,r9,r6
add r6,r5,r4
add r4,r1,r6
memw r4,r2
'''.strip()
assert match_code(code, pattern)
