IREE exception while crating file line info

[prosenjitdhole]

We are facing one failing testcase on windows in shark-ai with following exception :

E       iree.compiler._mlir_libs._site_initialize.<locals>.MLIRError: Unable to parse attribute:
E       error: "#llvm.di_file<\22vanilla_attention.py\22 in \22C:\\Users\\prosenj\\Downloads\\TestShark\\shark-ai\\test_env\\Lib\\site-packages\\wave_lang\\kernel\\wave\\templates\22>":1:44: unknown escape in string literal

This exception is coming from following stack trace :

sharktank\tests\kernels\attention_wave_test.py:27:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
test_env\Lib\site-packages\iree\turbine\aot\exporter.py:348: in export
    cm = TransformedModule(context=context, import_to="import")
test_env\Lib\site-packages\iree\turbine\aot\compiled_module.py:750: in __new__
    CompiledModule.run_import(self, import_to)
test_env\Lib\site-packages\iree\turbine\aot\compiled_module.py:572: in run_import
    CompiledModule.expand_custom_ops(inst)
test_env\Lib\site-packages\iree\turbine\aot\compiled_module.py:585: in expand_custom_ops
    p.run()
test_env\Lib\site-packages\iree\turbine\transforms\general\custom_op_expansion.py:72: in run
    self.expand_func(mr.op)
test_env\Lib\site-packages\iree\turbine\transforms\general\custom_op_expansion.py:95: in expand_func
    self.expand_custom_op(custom_op_reg, op)
test_env\Lib\site-packages\iree\turbine\transforms\general\custom_op_expansion.py:120: in expand_custom_op
    op_reg.generate(ksel, kb)
sharktank\sharktank\kernels\mlir_kernel.py:350: in generate
    mlir_spec = func(*input_values, *([None] * len(result_args)))
sharktank\sharktank\kernels\wave\attention.py:112: in wave_bhsd_flash_attention
    asm = get_wave_flash_attention_asm(
sharktank\sharktank\kernels\wave\attention.py:63: in get_wave_flash_attention_asm
    base_attention = wave_compile(options, base_attention_func)
test_env\Lib\site-packages\wave_lang\kernel\wave\compile.py:345: in wave_compile
    ) = kernel._trace_and_get_kernel_signature(options)
test_env\Lib\site-packages\wave_lang\kernel\wave\wave.py:982: in _trace_and_get_kernel_signature
    *self.compile_to_mlir(trace, context, module_op, options=options),
test_env\Lib\site-packages\wave_lang\kernel\wave\wave.py:685: in compile_to_mlir
    dispatch_entrypoint = exe.define_entrypoint(
test_env\Lib\site-packages\wave_lang\kernel\compiler\dispatch_codegen.py:156: in define_entrypoint
    func_location_with_di = create_debug_info_for_kernel(

debug_attr.py is passing a proper escaped string as the directory name (it is directly retrieved from os.path.dirname()) but somehow while invoking the MLIR code to parse the attribute string is not understanding the escape sequence.

    full_path = primary_loc.filename
    filename = os.path.basename(full_path)
    directory = os.path.dirname(os.path.abspath(full_path))
   .....
    # Create debug attributes
    di_file_str = f'#llvm.di_file<"{filename}" in "{directory}">'
    
    di_file = Attribute.parse(di_file_str)

If we change the directory name as follows (inside debug_attr.py from wave_lang), exception goes away :

directory = re.sub(r'\\', r'\\\\', directory)

This is not a correct fix. We need a fix inside LLVM code for parsing attribute string with escaped directory name in windows.

Steps to reproduce :
Please follow https://github.com/nod-ai/shark-ai/actions/runs/19112717161/workflow?pr=2498 on windows. This is a windows only issue.
Command to run :

pytest -n 20 sharktank\tests\kernels\attention_wave_test.py::wave_attention::test_wave_attention_causal --cov=sharktank --cov-report xml:cov.xml --cov-config=.coveragerc --durations=10 --log-cli-level=info -v

Note :
Currently this failing test is marked as XFAIL, in order to reproduce need to comment the XFAIL on test sharktank\tests\kernels\attention_wave_test.py::wave_attention::test_wave_attention_causal

[rkayaith]

Are you able to reproduce this parsing error directly? This seems to work fine locally:

>>> from iree.compiler.ir import Attribute, Context
>>> ctx = Context()
>>> path = r"foo\\bar"
>>> Attribute.parse(f'#llvm.di_file<"foo" in "{path}">', context=ctx)
Attribute(#llvm.di_file<"foo" in "foo\\bar">)