Windows support

[astrelsky]

Describe the intent of your PR here. Added Windows support, pretty self explanatory and should have been included from the beginning. (I'm aware the llvm and mlir dependencies were the main limiting factor).

Added

• Added test skip validation for missing gated weights to stabilize local test executions.

Changed

• Switched to using time.perf_counter_ns to prevent divide by 0 errors.
• Hoisted object FIFO acquires out of conditional structures to satisfy cyclostatic compiler analysis constraints.
• Swapped shell/awk commands for platform-independent Python execution blocks to handle symbol mapping on Windows.

PR Merge Checklist

1. The PR is rebased on the latest devel commit and pointing to devel.
2. Your PR has been reviewed and approved.
3. All checks are passing.

================================================== warnings summary ===================================================
iron/operators/dequant/test.py::test_dequant[iter0-input_length_1024-num_aie_columns_1-num_channels_1-tile_size_1024-group_size_32]
  D:\IRON\iron\operators\dequant\reference.py:42: UserWarning: torch.quantize_per_tensor, torch.quantize_per_channel and other quantized tensor creation functions that produce tensors with dtype torch.quint8, torch.qint8, and torch.qint32 are deprecated and will be removed in a future PyTorch release. Please see https://github.com/pytorch/pytorch/issues/184982 for more information. (Triggered internally at D:/b/torch/aten/src/ATen/quantized/Quantizer.cpp:116.)
    A = torch.quantize_per_channel(

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
============================== 2830 passed, 25 skipped, 1 warning in 1032.96s (0:17:12) ===============================

The test run results shown above was faster because I left the build folder in place after realizing how to fix the divide by 0 error I was encountering with the latency_us calculations, it had previously taken 2029.72s (0:33:49).

I used the llvm-aie, mlir-aie-no-rtti and mlir-air that were installed when I installed triton-windows. The restrictions in requirements.txt are too strict for the versions and there was no windows wheel for those.

llvm-aie                   21.0.0.2026071801+ce8c0f8f
mlir-aie-no-rtti           1.3.4.dev19+g92e9475
mlir-air                   0.0.1.2026071405+c7ee42f.no.rtti

[Copilot]

Pull request overview

This PR aims to improve Windows compatibility across IRON by removing Unix-shell dependencies in the build/compilation flow, improving local test stability when gated model weights are missing, and adjusting MHA worker control flow to satisfy cyclostatic compiler constraints.

Changes:

• Updated MHA PV worker FIFO acquire/release structure around conditional execution.
• Made symbol-map generation platform-independent by replacing sh|awk with a Python block on Windows.
• Improved Llama application test behavior for missing weights and made invocation use the active Python interpreter.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File	Description
iron/operators/mha/design.py	Restructures FIFO acquire/release and conditional PV loop behavior in the MHA design.
iron/common/compilation/base.py	Replaces Unix `sh
iron/applications/llama_3.2_1b/test.py	Adds skipping for missing Llama weights and uses the current interpreter for invocation.

Comments suppressed due to low confidence (2)

iron/operators/mha/design.py:583

• The final-iteration P/V/scale acquires are now unconditional. When loop_idx_kv <= 1, this will still consume one extra element from each FIFO (and release it) even though the loop_idx_kv > 1 branch is not taken, which will desynchronize FIFO ordering vs the QK/softmax workers.

                elem_in_p = of_p.acquire(1)
                elem_in_v = of_v.acquire(1)
                elt_of_out_scale3 = of_scale.acquire(1)
                with if_(loop_idx_kv > 1) as if_op:
                    matmul_PV(

iron/applications/llama_3.2_1b/test.py:47

• The command is executed with shell=True and built as a single string. On Windows this is fragile (e.g., sys.executable and weight paths can include spaces). Passing an argv list with shell=False is more reliable and avoids quoting issues.

    command = f"{sys.executable} {test_dir}/llama_npu.py {weights_dir}/llama3.2-1b/model.safetensors {weights_dir}/llama3.2-1b/tokenizer.model --num-tokens {num_tokens} --prompt-len {prompt_len}"

    result = subprocess.run(
        command,
        cwd=test_dir,
        shell=True,
        capture_output=True,
        text=True,
    )

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[astrelsky]

While technically possible, I think it is extremely unlikely. Someone who understands the underlying math and operator will have to take a closer look. This was left unchanged and the only portion that was changed was the locking to fix the compilation errors.

[andrej]

What are the changes in mha/design.py?

[astrelsky]

What are the changes in mha/design.py?

Fixing the compile error. You cannot use acquire and release in a conditional statement within a loop, it must be unconditional.

I had to dig up the error.

D:\IRON\build\MHA_h1_s16384_d64_kv0_npu2.mlir:415:21: error: 'aie.objectfifo.acquire' op cannot statically analyze cyclostatic acquire pattern: acquire/release is inside a conditional %17 = aie.objectfifo.acquire @outP0(Consume, 1) : !aie.objectfifosubview<memref<64x64xbf16>> ^ D:\IRON\build\MHA_h1_s16384_d64_kv0_npu2.mlir:415:21: note: see current operation: %1586 = "aie.objectfifo.acquire"() <{objFifo_name = @outP0_cons, port = 1 : i32, size = 1 : i32}> : () -> !aie.objectfifosubview<memref<64x64xbf16>> D:\IRON\build\MHA_h1_s16384_d64_kv0_npu2.mlir:2:3: error: Failures have been detected while processing an MLIR pass pipeline aie.device(npu2) { ^ D:\IRON\build\MHA_h1_s16384_d64_kv0_npu2.mlir:2:3: note: Pipeline failed while executing [`AIEObjectFifoStatefulTransform` on 'aie.device' operation: @main]: reproducer generated at `MHA_h1_s16384_d64_kv0_npu2.mlir.prj\resource_alloc_crash.mlir` AIECC COMPILATION FAILED Intermediate MLIR saved to: MHA_h1_s16384_d64_kv0_npu2.mlir.prj\aiecc_failure_1784384099_31012.mlir Repeater script generated: MHA_h1_s16384_d64_kv0_npu2.mlir.prj\aiecc_repeater_1784384099_31012.sh Error: Resource allocation pipeline failed Compilation failed

[andrej]

Ah, yes, I think that is with a newer MLIR-AIE compiler.

Would you mind submitting an MLIR-AIE upgrade PR including this MHA fix, and this Windows fix separately, since this is unrelated to Windows support?

Edit: I might have it the wrong way around, maybe CI is using a newer MLIR-AIE compiler and your local install is older? Please verify, but I have confirmed CI is currently passing for this test.

[astrelsky]

Ah, yes, I think that is with a newer MLIR-AIE compiler.

Would you mind submitting an MLIR-AIE upgrade PR including this MHA fix, and this Windows fix separately, since this is unrelated to Windows support?

I can. I'm just about out of time for the night unfortunately but I can get to it tomorrow afternoon and I'll take care of the formatting problem as well (I forgot to run it through black). I did however just confirm that the mha changes work with mlir_aie==1.3.5.dev20+g167f34d, so it is backwards compatible.

🤔

mlir-aie                   1.3.5.dev20+g167f34d
mlir-aie-no-rtti           1.3.4.dev19+g92e9475
mlir-air                   0.0.1.2026071405+c7ee42f.no.rtti

That's problematic.

You appear to be correct and I may have not noticed this at the time. I'll do more testing to figure out what is what tomorrow afternoon and sort it out. Thank you for your patience.