This is an archive of the discontinued LLVM Phabricator instance.

[MLIR][LLVMDialect] Added branch weights attribute to CondBrOp
ClosedPublic

Authored by georgemitenkov on Jul 13 2020, 12:21 AM.

Details

Summary

This patch introduces branch weights metadata to llvm.cond_br op in LLVM Dialect. It is modelled as optional ElementsAttr, for example:

llvm.cond_br %cond weights(dense<[1, 3]> : vector<2xi32>), ^bb1, ^bb2

When exporting to proper LLVM, this attribute is transformed into metadata node. The test for metadata creation is added to ../Target/llvmir.mlir.

Diff Detail

Event Timeline

Herald added a project: Restricted Project. · View Herald TranscriptJul 13 2020, 12:21 AM
rriddle added inline comments.Jul 13 2020, 12:51 PM
mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
180

Can you please avoid writing C++ parsers if you don't have to? Or file bugs for things that are missing? Seems like the only thing missing here is parsing optional attributes, as it looks like you are rewriting the parser for ArrayAttr. I would expect https://reviews.llvm.org/D83712 to be all you need to make this declarative.

georgemitenkov marked an inline comment as done.Jul 14 2020, 2:29 AM
georgemitenkov added inline comments.
mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
180

Thanks for the comment! I have followed a similar approach taken for some ops in the codebase where optional parsing is done manually, but I guess the codebase can be not very consistent then.

Using assembly format is much nicer, I agree. I will make changes to use the patch you mentioned.

mehdi_amini added inline comments.Jul 14 2020, 4:18 PM
mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
180

Many ops in the codebase also predates the declarative assembly feature, any migration of existing ops is welcome (either River or I would be happy to review patches to migrate to it)

rriddle added inline comments.Jul 14 2020, 4:28 PM
mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
180

Submitted the patch I referred to previously.

No worries George, it's hard to keep track of what can and can't be done sometimes in the codebase especially when not everything is using the desired thing. As Mehdi mentions, a lot of operations still haven't been updated to use the declarative form as new feature have been added. I try to do sweeps occasionally to port new ones over, but MLIR is getting larger and it can be hard to keep track of everything. If you see some in the future could you file bugs, or send a patch to fix if you can(that would be awesome). (Same applies to any other area of the codebase that is a bit inconsistent with the current recommended practices)

Thanks

Removed custom parser and printer for cond_br, and used assembly format for parsing optional attributes instead.

georgemitenkov marked 4 inline comments as done.Jul 15 2020, 1:38 AM
georgemitenkov added inline comments.
mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
180

Great, thanks!

ftynse requested changes to this revision.Jul 16 2020, 2:55 AM

Thanks! A couple of small comments.

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
474

Can this use DenseElementsAttr instead of ArrayAttr? It is more efficient and likely easier to work with in this use case.

483

Please format

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
447

Can we declare this above the if, initialize to nullptr and remove the duplication of calls to builder.CreateCondBr?

This revision now requires changes to proceed.Jul 16 2020, 2:55 AM
georgemitenkov marked 2 inline comments as done.Jul 16 2020, 4:18 AM
georgemitenkov added inline comments.
mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
474

If using DenseElementsAttr, then the op would be printed as

llvm.cond_br %cond dense<[5, 10]> : vector<2xi32>, ^bb1, ^bb2

which is not as concise and clear as if it was using ArrayAttr?

Addressed comments on branch weights attributes:

  • improved format
  • removed the duplication of calls to builder.CreateCondBr()

TODO: decide on whether to use DenseElementsAttr or ArrayAttr

ftynse added inline comments.Jul 16 2020, 4:31 AM
mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
474

I wouldn't motivate modeling choices by pretty-printing. It should be possible to wrap it with an extra keyword for better readability:

llvm.cond_br %cond weights(dense<[5, 10]> : vector<2xi32>), ^bb1, ^bb2
georgemitenkov marked 2 inline comments as done.

Better formatting in builder.

georgemitenkov marked an inline comment as done.Jul 16 2020, 4:35 AM
georgemitenkov added inline comments.
mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
474

I see, thanks!

georgemitenkov marked 2 inline comments as done.

Changed branch_weights from ArrayAttr to ElementsAttr.

georgemitenkov edited the summary of this revision. (Show Details)Jul 16 2020, 4:55 AM
ftynse accepted this revision.Jul 16 2020, 5:53 AM

Thanks!

This revision is now accepted and ready to land.Jul 16 2020, 5:53 AM
rriddle added inline comments.Jul 16 2020, 10:01 AM
mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
474

If you could constrain the type of the attribute in ODS to one with a builder, it should be elided from the printed format as well.

mehdi_amini added inline comments.Jul 16 2020, 2:55 PM
mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
474

The pretty-printing would be nice if we could customize it even more, e.g. something like:

llvm.cond_br %cond ^bb1, ^bb2  weights[5, 10]
georgemitenkov marked an inline comment as done.Jul 20 2020, 5:37 AM
georgemitenkov added inline comments.
mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
474

@rriddle @mehdi_amini I am not sure we can do this at the moment with the current infrastructure unless using a custom parser? Especially if using ElementsAttr.

branch_weights are not heavily used explicitly so I guess this printing will suffice for now?

This revision was automatically updated to reflect the committed changes.