def test_bmm():
test_ms_bmm((8, 16, 32), (8, 64, 32), 'float32', poly_sch=True) # matmul with batch
test_ms_bmm((1, 2, 32), (1, 1, 32), 'float32', poly_sch=True) # matmul with some axis equals to 1
test_ms_bmm((1, 1024, 1024), (1, 1024, 1024), 'float32', poly_sch=True)
# test_ms_bmm((1, 1024, 1024), (1, 1024, 1024), 'float32', (1, 1024, 1024)) # cannot store type float32
test_ms_bmm((1, 1024, 512), (1, 256, 512), 'float32', poly_sch=True)
return True
def bmm(poly_sch, fuzz_shape=None):
test_ms_bmm((8, 16, 32), (8, 64, 32), 'float32',
poly_sch=poly_sch) # matmul with batch
test_ms_bmm((1, 2, 32), (1, 1, 32), 'float32',
poly_sch=poly_sch) # matmul with some axis equals to 1
test_ms_bmm((1, 1024, 1024), (1, 1024, 1024), 'float32', poly_sch=poly_sch)
# test_ms_bmm((1, 1024, 1024), (1, 1024, 1024), 'float32', (1, 1024, 1024)) # cannot store type float32
test_ms_bmm((1, 1024, 512), (1, 256, 512), 'float32', poly_sch=poly_sch)
def bmm(poly_sch, fuzz_shape=None):
test_ms_bmm((1, 1024, 1024), (1, 1024, 1024), 'float32', poly_sch=poly_sch)
# test_ms_bmm((1, 1024, 1024), (1, 1024, 1024), 'float32', (1, 1024, 1024)) # cannot store type float32
test_ms_bmm((1, 1024, 512), (1, 256, 512), 'float32', poly_sch=poly_sch)
def bmm(poly_sch, fuzz_shape=None):
test_ms_bmm((768, 768), (768, 768),
'float16',
'float16',
layout1='NHDT',
layout2='NHDT',
layout_out='NHDT',
shape_bias=(1, ),
add_bias=False,
tensor_core=True,
poly_sch=poly_sch,
dim="0 0 128 128 0 1 128 128 0 2 64 4",
bind_block="6 6",
bind_thread="32 4")
test_ms_bmm((768, 768), (768, 768),
'float16',
'float16',
layout1='NHTD',
layout2='NHTD',
layout_out='NHDT',
shape_bias=(1, ),
add_bias=False,
tensor_core=True,
poly_sch=poly_sch,
dim="0 0 128 128 0 1 128 128 0 2 64 4",
bind_block="6 6",
bind_thread="32 8")
test_ms_bmm((32, 12, 128, 128), (32, 12, 128, 64),
'float16',
'float16',
layout1='NHDT',
layout2='NHTD',
layout_out='NHDT',
shape_bias=(1, ),
add_bias=False,
tensor_core=True,
poly_sch=poly_sch,
dim="0 0 1 1 0 1 1 1 0 2 64 64 0 3 64 64 0 4 64 4",
bind_block="12 32",
bind_thread="32 4")
test_ms_bmm((32, 12, 128, 64), (32, 12, 128, 64),
'float16',
'float16',
layout1='NHDT',
layout2='NHDT',
layout_out='NHDT',
shape_bias=(1, ),
add_bias=False,
tensor_core=True,
poly_sch=poly_sch,
dim="0 0 1 1 0 1 1 1 0 2 64 64 0 3 64 64 0 4 64 4",
bind_block="12 32",
bind_thread="32 4")
""" Bert Batch64
test_ms_bmm((8192, 768), (768, 768), 'float16', 'float16', layout1='NHDT', layout2='NHDT', layout_out='NHDT',
shape_bias=(1, ), add_bias=False, tensor_core=True, poly_sch=poly_sch,
dim="0 0 128 128 0 1 128 128 0 2 64 4", bind_block="6 64", bind_thread="32 4")
test_ms_bmm((8192, 768), (8192, 768), 'float16', 'float16', layout1='NHTD', layout2='NHTD', layout_out='NHDT',
shape_bias=(1, ), add_bias=False, tensor_core=True, poly_sch=poly_sch,
dim="0 0 64 64 0 1 64 64 0 2 64 4", bind_block="12 12", bind_thread="32 4")
test_ms_bmm((64, 12, 128, 128), (64, 12, 128, 64), 'float16', 'float16', layout1='NHDT', layout2='NHTD', layout_out='NHDT',
shape_bias=(1, ), add_bias=False, tensor_core=True, poly_sch=poly_sch,
dim="0 0 1 1 0 1 1 1 0 2 64 64 0 3 64 64 0 4 64 4", bind_block="12 64", bind_thread="32 4")
test_ms_bmm((64, 12, 128, 64), (64, 12, 128, 64), 'float16', 'float16', layout1='NHDT', layout2='NHDT', layout_out='NHDT',
shape_bias=(1, ), add_bias=False, tensor_core=True, poly_sch=poly_sch,
dim="0 0 1 1 0 1 1 1 0 2 64 64 0 1 128 128 0 2 64 4", bind_block="2 12 64", bind_thread="32 4")
test_ms_bmm((8192, 768), (8192, 3072), 'float16', 'float16', layout1='NHTD', layout2='NHTD', layout_out='NHDT',
shape_bias=(1, ), add_bias=False, tensor_core=True, poly_sch=poly_sch,
dim="0 0 128 128 0 1 128 128 0 2 64 4", bind_block="24 6", bind_thread="32 8")
test_ms_bmm((8192, 3072), (8192, 768), 'float16', 'float16', layout1='NHTD', layout2='NHTD', layout_out='NHDT',
shape_bias=(1, ), add_bias=False, tensor_core=True, poly_sch=poly_sch,
dim="0 0 128 128 0 1 128 128 0 2 64 4", bind_block="6 24", bind_thread="32 8")
test_ms_bmm((8192, 768), (3072, 768), 'float16', 'float16', layout1='NHDT', layout2='NHDT', layout_out='NHDT',
shape_bias=(1, ), add_bias=False, tensor_core=True, poly_sch=poly_sch,
dim="0 0 128 128 0 1 128 128 0 2 64 4", bind_block="24 64", bind_thread="32 4")
test_ms_bmm((8192, 3072), (768, 3072), 'float16', 'float16', layout1='NHDT', layout2='NHDT', layout_out='NHDT',
shape_bias=(1, ), add_bias=False, tensor_core=True, poly_sch=poly_sch,
dim="0 0 128 128 0 1 128 128 0 2 64 4", bind_block="6 64", bind_thread="32 8")
test_ms_bmm((8192, 3072), (3072, 768), 'float16', 'float16', layout1='NHDT', layout2='NHTD', layout_out='NHDT',
shape_bias=(1, ), add_bias=False, tensor_core=True, poly_sch=poly_sch,
dim="0 0 128 128 0 1 128 128 0 2 64 4", bind_block="6 64", bind_thread="32 8")
test_ms_bmm((8192, 768), (768, 3072), 'float16', 'float16', layout1='NHDT', layout2='NHTD', layout_out='NHDT',
shape_bias=(1, ), add_bias=False, tensor_core=True, poly_sch=poly_sch,
dim="0 0 128 128 0 1 128 128 0 2 64 4", bind_block="24 64", bind_thread="32 4")
"""
""" Bert Batch32