search

iree-org / iree

A retargetable MLIR-based machine learning compiler and runtime toolkit.
http://iree.dev/
Apache License 2.0
2.59k stars 580 forks source link

canonicalize dropping the `lowering_config` attached to an operation. #18697

Open pashu123 opened 3 days ago

pashu123 commented 3 days ago

What happened?

module_prefill_bs4$async_dispatch_20.mlir:9:7: error: not able to find the lowering config

  func.func @prefill_bs4$async_dispatch_20_pack_f16() {
  ^

module_prefill_bs4$async_dispatch_20.mlir:9:7: note: see current operation:

Steps to reproduce your issue

Input IR: https://gist.github.com/pashu123/17e3f882b8ddee143c367a6743d918d5

Command: iree-opt --canonicalize inp.mlir

The output IR won't have the lowering_config attribute attached to the tensor.pack operation.

What component(s) does this issue relate to?

No response

Version information

No response

Additional context

No response

hanhanW commented 3 days ago

The fix is to add the attributes when we create a new tensor.pack op: https://github.com/llvm/llvm-project/blob/78089d5845d14bf61bdc06209c32f0fb34927c68/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp#L4326-L4346

Or you can use in-place update like:

    rewriter.startOpModification(packOp);
    packOp.getSourceMutable().assign(source); // or .set(); I don't remember. :p
    packOp.getDestMutable().assign(dest);
    rewriter.finalizeOpModification(packOp);
pashu123 commented 2 days ago

The fix is to add the attributes when we create a new tensor.pack op: https://github.com/llvm/llvm-project/blob/78089d5845d14bf61bdc06209c32f0fb34927c68/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp#L4326-L4346

Or you can use in-place update like:

    rewriter.startOpModification(packOp);
    packOp.getSourceMutable().assign(source); // or .set(); I don't remember. :p
    packOp.getDestMutable().assign(dest);
    rewriter.finalizeOpModification(packOp);

Added the Patch here: https://github.com/llvm/llvm-project/pull/111261