[TORCH] Added flex_attention hop function

[keshavvinayak01]

Description

• Added support for PyTorch's flex_attention Higher-Order Operator in torch-mlir.
• Implemented Torch_AtenFlexAttentionOp with 6 operands (query, key, value, scale, enable_gqa, return_lse) and 2 optional attributes (score_mod_fn, mask_mod_fn) for function references.
• The FX importer (_import_hop_flex_attention) correctly extracts score/mask modification functions from get_attr nodes using module IDs, following the while_loop HOP pattern.
• Includes TODO markers for kernel_options performance tuning parameters.
• Imports flex_attention from PyTorch FX graphs into valid MLIR.

[zjgarvey]

This does enable importing to mlir.

However, the changes don't actually provide "support" for this op, since the torch op can neither be decomposed nor lowered to any other dialects.

Although we could review/merge this and subsequently add a lowering path for the op in MLIR, I would personally prefer the e2e support is added in the same PR as the import support.

This is a rather unique operator, so having passing e2e tests would give me a lot more confidence in the choices made here. Otherwise I'm basically just hoping that what you did generally makes sense (or doing a significant amount of work myself to check it out), because there really isn't much precedent for these kinds of choices in our codebase.

[Groverkss]

This does enable importing to mlir.

However, the changes don't actually provide "support" for this op, since the torch op can neither be decomposed nor lowered to any other dialects.

Although we could review/merge this and subsequently add a lowering path for the op in MLIR, I would personally prefer the e2e support is added in the same PR as the import support.

This is a rather unique operator, so having passing e2e tests would give me a lot more confidence in the choices made here. Otherwise I'm basically just hoping that what you did generally makes sense (or doing a significant amount of work myself to check it out), because there really isn't much precedent for these kinds of choices in our codebase.

The only thing needed to have this passing e2e tests is implementing TilingInterface for this operation:

torch-mlir/lib/Dialect/TMTensor/IR/TMTensorOps.cpp

Line 199 in 8d563af

LogicalResult AttentionOp::generateScalarImplementation(OpBuilder &b,

With that said, it's an unreasonable bar to set that every operation must compile e2e through torch-mlir. Torch-MLIR is not a compiler, even though it has tests for e2e paths. The project docs explicitly call out this:

Torch-MLIR is primarily a project that is integrated into compilers to bridge them to PyTorch and ONNX. If contemplating a new integration, it may be helpful to refer to existing downstreams:

IREE
Blade
While most of the project is exercised via testing paths, there are some ways that an end user can directly use the APIs without further integration:

It should be okay to land support for ops through the importer without it running e2e tests in torch-mlir. I've looked at the implementation of e2e tests for more complex ops like attention, and they are not good implementations, they don't add much value.

We should as a project allow landing PRs that add support to the importer seperately from e2e tests (Atleast for HOPs). I don't think having a dummy implementation for an op should be the bar to land an operation.