Some fixes to python temp file saving and tracing. (#7989)
* There was a bug with tfs and compiling to in memory output: it would return an empty byte array.
* Added tfs for core-input.mlir when compiling from a string.
* Added tfs for core-command-line.txt to capture the command line.
* On exception the trace will finalize the record it is working on (the call was being completely ommitted, which is not what you want, considering this is often for failure repros).
diff --git a/bindings/python/iree/runtime/function.py b/bindings/python/iree/runtime/function.py
index 826fac9..a7e3874 100644
--- a/bindings/python/iree/runtime/function.py
+++ b/bindings/python/iree/runtime/function.py
@@ -98,54 +98,55 @@
call_trace = None # type: Optional[tracing.CallTrace]
if self._tracer:
call_trace = self._tracer.start_call(self._vm_function)
+ try:
+ # Initialize the capacity to our total number of args, since we should
+ # be below that when doing a flat invocation. May want to be more
+ # conservative here when considering nesting.
+ inv = Invocation(self._device)
+ ret_descs = self._ret_descs
- # Initialize the capacity to our total number of args, since we should
- # be below that when doing a flat invocation. May want to be more
- # conservative here when considering nesting.
- inv = Invocation(self._device)
- ret_descs = self._ret_descs
+ # Merge keyword args in by name->position mapping.
+ if kwargs:
+ args = list(args)
+ len_delta = self._max_named_arg_index - len(args) + 1
+ if len_delta > 0:
+ # Fill in MissingArgument placeholders before arranging kwarg input.
+ # Any remaining placeholders will fail arity checks later on.
+ args.extend([MissingArgument] * len_delta)
- # Merge keyword args in by name->position mapping.
- if kwargs:
- args = list(args)
- len_delta = self._max_named_arg_index - len(args) + 1
- if len_delta > 0:
- # Fill in MissingArgument placeholders before arranging kwarg input.
- # Any remaining placeholders will fail arity checks later on.
- args.extend([MissingArgument] * len_delta)
+ for kwarg_key, kwarg_value in kwargs.items():
+ try:
+ kwarg_index = self._named_arg_indices[kwarg_key]
+ except KeyError:
+ raise ArgumentError(f"specified kwarg '{kwarg_key}' is unknown")
+ args[kwarg_index] = kwarg_value
- for kwarg_key, kwarg_value in kwargs.items():
- try:
- kwarg_index = self._named_arg_indices[kwarg_key]
- except KeyError:
- raise ArgumentError(f"specified kwarg '{kwarg_key}' is unknown")
- args[kwarg_index] = kwarg_value
+ arg_list = VmVariantList(len(args))
+ ret_list = VmVariantList(len(ret_descs) if ret_descs is not None else 1)
+ _merge_python_sequence_to_vm(inv, arg_list, args, self._arg_descs)
+ if call_trace:
+ call_trace.add_vm_list(arg_list, "args")
+ self._vm_context.invoke(self._vm_function, arg_list, ret_list)
+ if call_trace:
+ call_trace.add_vm_list(ret_list, "results")
- arg_list = VmVariantList(len(args))
- ret_list = VmVariantList(len(ret_descs) if ret_descs is not None else 1)
- _merge_python_sequence_to_vm(inv, arg_list, args, self._arg_descs)
- if call_trace:
- call_trace.add_vm_list(arg_list, "args")
- self._vm_context.invoke(self._vm_function, arg_list, ret_list)
- if call_trace:
- call_trace.add_vm_list(ret_list, "results")
-
- # Un-inline the results to align with reflection, as needed.
- reflection_aligned_ret_list = ret_list
- if self._has_inlined_results:
- reflection_aligned_ret_list = VmVariantList(1)
- reflection_aligned_ret_list.push_list(ret_list)
- returns = _extract_vm_sequence_to_python(inv, reflection_aligned_ret_list,
- ret_descs)
- if call_trace:
- call_trace.end_call()
- return_arity = len(returns)
- if return_arity == 1:
- return returns[0]
- elif return_arity == 0:
- return None
- else:
- return tuple(returns)
+ # Un-inline the results to align with reflection, as needed.
+ reflection_aligned_ret_list = ret_list
+ if self._has_inlined_results:
+ reflection_aligned_ret_list = VmVariantList(1)
+ reflection_aligned_ret_list.push_list(ret_list)
+ returns = _extract_vm_sequence_to_python(inv, reflection_aligned_ret_list,
+ ret_descs)
+ return_arity = len(returns)
+ if return_arity == 1:
+ return returns[0]
+ elif return_arity == 0:
+ return None
+ else:
+ return tuple(returns)
+ finally:
+ if call_trace:
+ call_trace.end_call()
def _parse_abi_dict(self, vm_function: VmFunction):
reflection = vm_function.reflection
diff --git a/llvm-external-projects/iree-compiler-api/python/iree/compiler/tools/core.py b/llvm-external-projects/iree-compiler-api/python/iree/compiler/tools/core.py
index 2281ad7..b675b07 100644
--- a/llvm-external-projects/iree-compiler-api/python/iree/compiler/tools/core.py
+++ b/llvm-external-projects/iree-compiler-api/python/iree/compiler/tools/core.py
@@ -184,10 +184,8 @@
cl.append(f"--iree-hal-target-backends={target_backend}")
# Output file.
- output_file = tfs.alloc_optional("core-output.bin",
- export_as=options.output_file)
- if output_file:
- cl.append(f"-o={output_file}")
+ if options.output_file:
+ cl.append(f"-o={options.output_file}")
# Translation to perform.
cl.append("--iree-mlir-to-vm-bytecode-module")
@@ -237,10 +235,25 @@
"""
with TempFileSaver.implicit() as tfs:
options = CompilerOptions(**kwargs)
+ retained_output_file = tfs.alloc_optional("core-output.bin",
+ export_as=options.output_file)
+ if options.output_file:
+ options.output_file = retained_output_file
cl = build_compile_command_line(input_file, tfs, options)
+
+ # Save a temp file with the command line.
+ retained_cl = tfs.alloc_optional("core-command-line.txt")
+ if retained_cl:
+ with open(retained_cl, "wt") as f:
+ f.write(" ".join(cl))
+
result = invoke_immediate(cl)
if options.output_file:
return None
+ # Output as string needs to write to the retained output file itself.
+ if retained_output_file:
+ with open(retained_output_file, "wb") as f:
+ f.write(result)
return result
@@ -255,11 +268,32 @@
was specified in the options.
"""
with TempFileSaver.implicit() as tfs:
+ retained_input_file = tfs.alloc_optional("core-input.mlir")
+ if retained_input_file:
+ with open(retained_input_file,
+ "wt" if isinstance(input_str, str) else "wb") as f:
+ f.write(input_str)
options = CompilerOptions(**kwargs)
+ retained_output_file = tfs.alloc_optional("core-output.bin",
+ export_as=options.output_file)
+ if options.output_file:
+ options.output_file = retained_output_file
cl = build_compile_command_line("-", tfs, options)
input_bytes = input_str.encode("utf-8") if isinstance(input_str,
str) else input_str
+
+ # Save a temp file with the command line.
+ retained_cl = tfs.alloc_optional("core-command-line.txt")
+ if retained_cl:
+ with open(retained_cl, "wt") as f:
+ f.write(" ".join(cl))
+
result = invoke_immediate(cl, immediate_input=input_bytes)
if options.output_file:
return None
+
+ # Output as string needs to write to the retained output file itself.
+ if retained_output_file:
+ with open(retained_output_file, "wb") as f:
+ f.write(result)
return result
diff --git a/llvm-external-projects/iree-compiler-api/unittests/tools/compiler_core_test.py b/llvm-external-projects/iree-compiler-api/unittests/tools/compiler_core_test.py
index 996fc29..f0accd6 100644
--- a/llvm-external-projects/iree-compiler-api/unittests/tools/compiler_core_test.py
+++ b/llvm-external-projects/iree-compiler-api/unittests/tools/compiler_core_test.py
@@ -199,6 +199,54 @@
self.assertEqual(temp_contents, output_contents)
temp_dir.cleanup()
+ def testExplicitTempFileSaverCompileToStrTextInput(self):
+ temp_dir = tempfile.TemporaryDirectory()
+ with iree.compiler.tools.TempFileSaver(temp_dir.name):
+ output = iree.compiler.tools.compile_str(
+ SIMPLE_MUL_ASM,
+ input_type="mhlo",
+ target_backends=iree.compiler.tools.DEFAULT_TESTING_BACKENDS)
+ self.assertIsNotNone(output)
+ self.assertGreater(len(output), 0)
+
+ # There should be a core-input.mlir and core-output.bin in the temp dir.
+ expected_temp_file = os.path.join(temp_dir.name, "core-output.bin")
+ self.assertTrue(os.path.exists(expected_temp_file))
+ with open(expected_temp_file, "rb") as f:
+ temp_output = f.read()
+ self.assertEqual(output, temp_output)
+
+ expected_temp_file = os.path.join(temp_dir.name, "core-input.mlir")
+ self.assertTrue(os.path.exists(expected_temp_file))
+ with open(expected_temp_file, "rt") as f:
+ input_contents = f.read()
+ self.assertEqual(SIMPLE_MUL_ASM, input_contents)
+ temp_dir.cleanup()
+
+ def testExplicitTempFileSaverBinaryInput(self):
+ temp_dir = tempfile.TemporaryDirectory()
+ with iree.compiler.tools.TempFileSaver(temp_dir.name):
+ output = iree.compiler.tools.compile_str(
+ SIMPLE_MUL_ASM,
+ input_type="mhlo",
+ target_backends=iree.compiler.tools.DEFAULT_TESTING_BACKENDS)
+ self.assertIsNotNone(output)
+ self.assertGreater(len(output), 0)
+
+ # There should be a core-input.mlir and core-output.bin in the temp dir.
+ expected_temp_file = os.path.join(temp_dir.name, "core-output.bin")
+ self.assertTrue(os.path.exists(expected_temp_file))
+ with open(expected_temp_file, "rb") as f:
+ temp_output = f.read()
+ self.assertEqual(output, temp_output)
+
+ expected_temp_file = os.path.join(temp_dir.name, "core-input.mlir")
+ self.assertTrue(os.path.exists(expected_temp_file))
+ with open(expected_temp_file, "rt") as f:
+ input_contents = f.read()
+ self.assertEqual(SIMPLE_MUL_ASM, input_contents)
+ temp_dir.cleanup()
+
def testEnvTempFileSaver(self):
temp_dir = tempfile.TemporaryDirectory()
os.environ["IREE_SAVE_TEMPS"] = temp_dir.name
