def test_compile_and_run(self):
mlirdir = resource_loader.get_data_files_path()
for filename in os.listdir(mlirdir):
if not filename.endswith('.mlir'):
continue
with open(os.path.join(mlirdir, filename), mode='r') as f:
mlir_function = f.read()
arg_attrs = []
with ir.Context() as ctx:
ctx.allow_unregistered_dialects = True
module = ir.Module.parse(mlir_function)
func = module.body.operations[0]
function_name = ir.StringAttr(
func.attributes['sym_name']).value
if _ARG_ATTRIBUTES_NAME in func.attributes:
arg_attrs = ir.ArrayAttr(
func.attributes[_ARG_ATTRIBUTES_NAME])
logging.info(f'processing {filename}')
start = time.perf_counter()
compiled = cpurt.compile(mlir_function,
function_name,
tf_cpurt.Specialization.ENABLED,
vectorize=False)
end = time.perf_counter()
logging.info(
f'compiled {filename} in {end-start:0.4f} seconds')
if not arg_attrs:
continue
args = []
for arg_attr in arg_attrs:
attr_dict = ir.DictAttr(arg_attr)
if _SHAPE_VALUE_ATTRIBUTE_NAME in attr_dict:
shape_value_attr = ir.DenseIntElementsAttr(
attr_dict[_SHAPE_VALUE_ATTRIBUTE_NAME])
shape_value = np.array(list(shape_value_attr)).astype(
np.int32)
args.append(shape_value)
elif _STATIC_TYPE_ATTRIBUTE_NAME in attr_dict:
static_type = ir.TypeAttr(
attr_dict[_STATIC_TYPE_ATTRIBUTE_NAME]).value
shaped_type = ir.ShapedType(static_type)
dims = []
for i in range(shaped_type.rank):
dims.append(shaped_type.get_dim_size(i))
np_element_type = TfCompileAndRunTest.mlir_type_to_np_type(
shaped_type.element_type)
arg = np.random.uniform(
-10000.0, 10000.0,
size=dims).astype(np_element_type)
args.append(arg)
if len(args) != len(arg_attrs):
logging.error(
'expected valid python_test_attrs attributes for each argument'
)
continue
start = time.perf_counter()
cpurt.execute(compiled, args)
end = time.perf_counter()
logging.info(
f'executed {filename} in {end-start:0.4f} seconds')
def test_compile_and_run(self):
filename = FLAGS.test_file_name
if not os.path.isabs(filename):
filename = os.path.join(resource_loader.get_data_files_path(), filename)
with gfile.GFile(filename, mode='r') as f:
mlir_function = f.read()
arg_attrs = []
with ir.Context() as ctx:
ctx.allow_unregistered_dialects = True
module = ir.Module.parse(mlir_function)
func = module.body.operations[0]
function_name = ir.StringAttr(func.attributes['sym_name']).value
# If the function has arguments, we expect argument attributes.
if func.regions[0].blocks[0].arguments:
self.assertIn(_ARG_ATTRIBUTES_NAME, func.attributes)
arg_attrs = ir.ArrayAttr(func.attributes[_ARG_ATTRIBUTES_NAME])
logging.info(f'processing {filename}')
start = time.perf_counter()
compiled = jitrt.compile(
mlir_function,
function_name,
tf_jitrt.Specialization.ENABLED,
vectorize=FLAGS.vectorize)
end = time.perf_counter()
logging.info(f'compiled {filename} in {end-start:0.4f} seconds')
np.random.seed(FLAGS.input_data_seed)
args = []
for arg_attr in arg_attrs:
attr_dict = ir.DictAttr(arg_attr)
if _SHAPE_VALUE_ATTRIBUTE_NAME in attr_dict:
shape_value_attr = ir.DenseIntElementsAttr(
attr_dict[_SHAPE_VALUE_ATTRIBUTE_NAME])
shape_value = np.array(list(shape_value_attr)).astype(np.int32)
args.append(shape_value)
elif _STATIC_TYPE_ATTRIBUTE_NAME in attr_dict:
static_type = ir.TypeAttr(
attr_dict[_STATIC_TYPE_ATTRIBUTE_NAME]).value
shaped_type = ir.ShapedType(static_type)
np_element_type = CompileAndRunTest.mlir_type_to_np_type(
shaped_type.element_type)
arg = np.random.uniform(
-10000.0, 10000.0, size=shaped_type.shape).astype(np_element_type)
args.append(arg)
self.assertEqual(len(args), len(arg_attrs))
start = time.perf_counter()
result = jitrt.execute(compiled, args)
end = time.perf_counter()
logging.info(f'executed {filename} in {end-start:0.4f} seconds')
if FLAGS.compare_with_tensorflow:
start = time.perf_counter()
expected = tfrt_fallback.run_tfrt_fallback(mlir_function, function_name,
args)
end = time.perf_counter()
logging.info(
f'executed {filename} via tfrt fallback in {end-start:0.4f} seconds'
)
np.testing.assert_allclose(result, expected, rtol=1e-5, atol=1e-5)
def test_compile_and_run(self):
filename = FLAGS.test_file_name
if not os.path.isabs(filename):
filename = os.path.join(resource_loader.get_data_files_path(),
filename)
with gfile.GFile(filename, mode='r') as f:
mlir_function = f.read()
arg_attrs = []
with ir.Context() as ctx:
ctx.allow_unregistered_dialects = True
module = ir.Module.parse(mlir_function)
func = module.body.operations[0]
function_name = ir.StringAttr(
func.attributes['sym_name']).value
if _ARG_ATTRIBUTES_NAME in func.attributes:
arg_attrs = ir.ArrayAttr(
func.attributes[_ARG_ATTRIBUTES_NAME])
logging.info(f'processing {filename}')
start = time.perf_counter()
compiled = jitrt.compile(mlir_function,
function_name,
tf_jitrt.Specialization.ENABLED,
vectorize=FLAGS.vectorize)
end = time.perf_counter()
logging.info(f'compiled {filename} in {end-start:0.4f} seconds')
if not arg_attrs:
return
np.random.seed(FLAGS.input_data_seed)
args = []
for arg_attr in arg_attrs:
attr_dict = ir.DictAttr(arg_attr)
if _SHAPE_VALUE_ATTRIBUTE_NAME in attr_dict:
shape_value_attr = ir.DenseIntElementsAttr(
attr_dict[_SHAPE_VALUE_ATTRIBUTE_NAME])
shape_value = np.array(list(shape_value_attr)).astype(
np.int32)
args.append(shape_value)
elif _STATIC_TYPE_ATTRIBUTE_NAME in attr_dict:
static_type = ir.TypeAttr(
attr_dict[_STATIC_TYPE_ATTRIBUTE_NAME]).value
shaped_type = ir.ShapedType(static_type)
np_element_type = CompileAndRunTest.mlir_type_to_np_type(
shaped_type.element_type)
arg = np.random.uniform(
-10000.0, 10000.0,
size=shaped_type.shape).astype(np_element_type)
args.append(arg)
if len(args) != len(arg_attrs):
logging.error(
'expected valid python_test_attrs attributes for each argument'
)
return
start = time.perf_counter()
jitrt.execute(compiled, args)
end = time.perf_counter()
logging.info(f'executed {filename} in {end-start:0.4f} seconds')
def attribute_to_var(attr):
import mlir.ir as ir
if attr is None:
return None
if not isinstance(attr, ir.Attribute):
raise TypeError("attribute_to_var only accepts MLIR Attributes")
# If it's not the root type, assume it's already been downcasted and don't do
# the expensive probing below.
if attr.__class__ != ir.Attribute:
return attr
try:
return ir.BoolAttr(attr).value
except ValueError:
pass
try:
return ir.IntegerAttr(attr).value
except ValueError:
pass
try:
return ir.StringAttr(attr).value
except ValueError:
pass
try:
return ir.TypeAttr(attr).value
except ValueError:
pass
try:
arr = ir.ArrayAttr(attr)
return [attribute_to_var(x) for x in arr]
except ValueError:
pass
try:
dict = ir.DictAttr(attr)
return {i.name: attribute_to_var(i.attr) for i in dict}
except ValueError:
pass
raise TypeError(f"Cannot convert {repr(attr)} to python value")
def test_compile_and_run(self):
filename = _TEST_FILE_NAME.value
if not os.path.isabs(filename):
filename = os.path.join(resource_loader.get_data_files_path(),
filename)
with gfile.GFile(filename, mode='r') as f:
mlir_function = f.read()
arg_attrs = []
with ir.Context() as ctx:
ctx.allow_unregistered_dialects = True
module = ir.Module.parse(mlir_function)
func = module.body.operations[0]
function_type = ir.FunctionType(
ir.TypeAttr(func.attributes[_FUNCTION_TYPE_NAME]).value)
function_name = ir.StringAttr(
func.attributes['sym_name']).value
# If the function has arguments, we expect argument attributes.
entry_block = func.regions[0].blocks[0]
if entry_block.arguments:
self.assertIn(_ARG_ATTRIBUTES_NAME, func.attributes)
arg_attrs = ir.ArrayAttr(
func.attributes[_ARG_ATTRIBUTES_NAME])
logging.info(f'processing {filename}')
start = time.perf_counter()
compiled = jitrt.compile(
mlir_function,
function_name,
tf_jitrt.Specialization.ENABLED,
vectorize=_VECTORIZE.value,
one_shot_bufferize=_ONE_SHOT_BUFFERIZE.value)
end = time.perf_counter()
logging.info(f'compiled {filename} in {end-start:0.4f} seconds')
np.random.seed(_INPUT_DATA_SEED.value)
args = []
for arg_attr in arg_attrs:
attr_dict = ir.DictAttr(arg_attr)
if _SHAPE_VALUE_ATTRIBUTE_NAME in attr_dict:
shape_value_attr = ir.DenseIntElementsAttr(
attr_dict[_SHAPE_VALUE_ATTRIBUTE_NAME])
shape_value = np.array(list(shape_value_attr)).astype(
np.int32)
args.append(shape_value)
elif _STATIC_TYPE_ATTRIBUTE_NAME in attr_dict:
static_type = ir.TypeAttr(
attr_dict[_STATIC_TYPE_ATTRIBUTE_NAME]).value
shaped_type = ir.ShapedType(static_type)
np_element_type = CompileAndRunTest.mlir_type_to_np_type(
shaped_type.element_type)
arg = np.random.uniform(
-10000.0, 10000.0,
size=shaped_type.shape).astype(np_element_type)
args.append(arg)
self.assertEqual(len(args), len(arg_attrs))
start = time.perf_counter()
result = jitrt.execute(compiled, args)
end = time.perf_counter()
logging.info(f'executed {filename} in {end-start:0.4f} seconds')
if _COMPARE_WITH_TENSORFLOW.value:
start = time.perf_counter()
expected = tfrt_fallback.run_tfrt_fallback(
mlir_function, function_name, args)
end = time.perf_counter()
logging.info(
f'executed {filename} via tfrt fallback in {end-start:0.4f} seconds'
)
if len(function_type.results) > 1:
# If there is more than one result, we need to iterate manually,
# otherwise np.testing.assert_allclose will complain if not all
# results have equal size.
self.assertEqual(len(result), len(expected))
for res, expect in zip(result, expected):
np.testing.assert_allclose(res,
expect,
rtol=1e-5,
atol=1e-5)
else:
np.testing.assert_allclose(result,
expected,
rtol=1e-5,
atol=1e-5)