PR #177922: [inductor] Add inline_asm_elementwise higher-order operator

• Repository: pytorch/pytorch
• Author: eellison
• State: closed | merged: False
• Link: https://github.com/pytorch/pytorch/pull/177922

Description (problem / solution / changelog)

Stack from ghstack (oldest at bottom):

• -> #177922

Adds a higher-order operator for inline PTX assembly that works in both eager and compiled modes:

• Eager: JIT compiles CUDA kernels via Jiterator with inline asm
• Compiled: Lowers to tl.inline_asm_elementwise in Triton via Inductor

This enables using PTX instructions not exposed through standard PyTorch ops (e.g., cvt.rn.satfinite.e2m1x2.f32 for NVFP4 quantization) while maintaining bitwise equivalence between eager and compiled execution.

Example usage:

from torch._higher_order_ops import inline_asm_elementwise

result = inline_asm_elementwise(
    x, y,
    asm_str="add.f32 $0, $1, $2;",
    constraints="=f,f,f",
    dtype=torch.float32,
)

Notes:

jitterator has limited support for the following:

• multiple outputs
• different output dtype than input
• pack != 1

In these cases, today, we error in jitterator and succeed in inductor.

We inherit the output striding behavior from jitterator in inductor/compilation (which follows eager pointwise ops).

Inductor details:

• Block size is not guaranteed to be a multiple of pack. Particularly, at the end of a persistent reduction, it's possible that xblock == 1. For this case, i've added a triton helper to pad the triton tensor to a multiple of pack, and then split to the actual output after. I suspect this is unlikely to occur but it's better to handle anyway.

• Inductor computes bf16/fp16 inside the kernel. For asm targeting these dtypes, we cast prior to invoking the asm. (even with emulate precision casts, we still compute in fp32, just add extra casts).

Otherwise it works more or less as the existing lowering.

reland of https://github.com/pytorch/pytorch/pull/175814

Written with Claude Code.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @chauhang @aakhundov @coconutruben @jataylo @Lucaskabela @mlazos

Changed files

• test/higher_order_ops/test_inline_asm_elementwise.py (added, +585/-0)
• torch/_dynamo/variables/higher_order_ops.py (modified, +26/-0)
• torch/_higher_order_ops/__init__.py (modified, +2/-0)
• torch/_higher_order_ops/inline_asm_elementwise.py (added, +299/-0)
• torch/_inductor/codegen/common.py (modified, +1/-0)
• torch/_inductor/codegen/triton.py (modified, +49/-4)
• torch/_inductor/dtype_propagation.py (modified, +7/-1)
• torch/_inductor/lowering.py (modified, +37/-0)
• torch/_inductor/ops_handler.py (modified, +1/-0)
• torch/_inductor/runtime/triton_helpers.py (modified, +24/-0)
• torch/testing/_internal/hop_db.py (modified, +28/-0)