[Mlir-commits] [mlir] [MLIR][DRAFT] Add loop annotation attribute interface (PR #131901)
Asher Mancinelli
llvmlistbot at llvm.org
Tue Mar 18 12:47:20 PDT 2025
https://github.com/ashermancinelli created https://github.com/llvm/llvm-project/pull/131901
This is a draft PR in case folks have comment on the RFC here: https://discourse.llvm.org/t/rfc-attribute-interface-for-loop-annotation-metadata/85311
---
Metadata that live on loops should be propagated through Affine, SCF, and all dialects with ops that conform to `LoopLikeInterface`, but today, conversions need to depend on the LLVM dialect directly in order to propagate them through conversion passes and are often missed.
This is what currently happens in SCF to CF conversion:
```
// in SCFToControlFlow.cpp
// Let the CondBranchOp carry the LLVM attributes from the ForOp, such as the
// llvm.loop_annotation attribute.
SmallVector<NamedAttribute> llvmAttrs;
llvm::copy_if(forOp->getAttrs(), std::back_inserter(llvmAttrs),
[](auto attr) {
return isa<LLVM::LLVMDialect>(attr.getValue().getDialect());
});
condBranchOp->setDiscardableAttrs(llvmAttrs);
```
We found gaps in other dialect’s propagation of this metadata, but making conversion passes between Affine, Linalg and other dialects with a LoopLikeInterface depend on the llvm dialect directly doesn’t seem right.
Instead, if we add an attribute interface (maybe `LoopAnnotationAttrInterface`) that `LoopAnnotationAttr` and any future loop attributes can add as a trait, conversion passes could look for and propagate this metadata on loop-like ops when creating ops in the target dialect.
The example above might look like this, with the dependency on the LLVM dialect removed:
```
// Let the CondBranchOp carry the LLVM attributes from the ForOp, such as the
// llvm.loop_annotation attribute.
SmallVector<NamedAttribute> loopAnnotationAttrs;
llvm::copy_if(forOp->getAttrs(), std::back_inserter(loopAnnotationAttrs),
[](auto attr) -> bool {
return mlir::isa<mlir::LoopAnnotationAttrInterface>(attr.getValue());
});
condBranchOp->setDiscardableAttrs(loopAnnotationAttrs);
```
Or if we’re okay adding an interface method like getLoopAnnotationAttributes to LoopLikeInterface:
```
// Let the CondBranchOp carry the LLVM attributes from the ForOp, such as the
// llvm.loop_annotation attribute.
condBranchOp->setDiscardableAttrs(forOp.getLoopAnnotationAttributes());
```
>From f5f69817fde681e5f3609480446a9d103a5df20c Mon Sep 17 00:00:00 2001
From: Asher Mancinelli <ashermancinelli at gmail.com>
Date: Tue, 18 Mar 2025 09:38:01 -0700
Subject: [PATCH] [mlir] Add loop annotation attribute interface
Interface used to propagate loop annotation attributes between passes
that shouldn't need to rely on the dialects they come from.
---
.../mlir/Dialect/LLVMIR/LLVMAttrDefs.td | 3 +-
mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h | 1 +
mlir/include/mlir/Interfaces/CMakeLists.txt | 6 ++++
.../Interfaces/LoopAnnotationAttrInterface.h | 15 ++++++++++
.../Interfaces/LoopAnnotationAttrInterface.td | 28 +++++++++++++++++++
.../mlir/Interfaces/LoopLikeInterface.h | 1 +
.../mlir/Interfaces/LoopLikeInterface.td | 17 +++++++++++
.../SCFToControlFlow/CMakeLists.txt | 1 -
.../SCFToControlFlow/SCFToControlFlow.cpp | 9 ++----
mlir/lib/Interfaces/CMakeLists.txt | 14 ++++++++++
.../LoopAnnotationAttrInterface.cpp | 3 ++
11 files changed, 89 insertions(+), 9 deletions(-)
create mode 100644 mlir/include/mlir/Interfaces/LoopAnnotationAttrInterface.h
create mode 100644 mlir/include/mlir/Interfaces/LoopAnnotationAttrInterface.td
create mode 100644 mlir/lib/Interfaces/LoopAnnotationAttrInterface.cpp
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
index ede9f2d365b20..a1a659ac40cdd 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
@@ -9,6 +9,7 @@
#ifndef LLVMIR_ATTRDEFS
#define LLVMIR_ATTRDEFS
+include "mlir/Interfaces/LoopAnnotationAttrInterface.td"
include "mlir/Dialect/LLVMIR/LLVMDialect.td"
include "mlir/Dialect/LLVMIR/LLVMInterfaces.td"
include "mlir/IR/AttrTypeBase.td"
@@ -204,7 +205,7 @@ def LoopUnswitchAttr : LLVM_Attr<"LoopUnswitch", "loop_unswitch"> {
let assemblyFormat = "`<` struct(params) `>`";
}
-def LoopAnnotationAttr : LLVM_Attr<"LoopAnnotation", "loop_annotation"> {
+def LoopAnnotationAttr : LLVM_Attr<"LoopAnnotation", "loop_annotation", [LoopAnnotationAttrInterface]> {
let description = [{
This attributes encapsulates "loop metadata". It is meant to decorate
branches that are "latches" (loop backedges) and maps to the `!llvm.loop`
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h
index 3ede857733242..77e0dd346f09c 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrs.h
@@ -16,6 +16,7 @@
#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
#include "mlir/IR/OpImplementation.h"
+#include "mlir/Interfaces/LoopAnnotationAttrInterface.h"
#include <optional>
#include "mlir/Dialect/LLVMIR/LLVMOpsEnums.h.inc"
diff --git a/mlir/include/mlir/Interfaces/CMakeLists.txt b/mlir/include/mlir/Interfaces/CMakeLists.txt
index d81298bb4daf0..0943c2af65d38 100644
--- a/mlir/include/mlir/Interfaces/CMakeLists.txt
+++ b/mlir/include/mlir/Interfaces/CMakeLists.txt
@@ -26,6 +26,12 @@ mlir_tablegen(MemorySlotTypeInterfaces.cpp.inc -gen-type-interface-defs)
add_public_tablegen_target(MLIRMemorySlotInterfacesIncGen)
add_dependencies(mlir-generic-headers MLIRMemorySlotInterfacesIncGen)
+set(LLVM_TARGET_DEFINITIONS LoopAnnotationAttrInterface.td)
+mlir_tablegen(LoopAnnotationAttrInterface.h.inc -gen-attr-interface-decls)
+mlir_tablegen(LoopAnnotationAttrInterface.cpp.inc -gen-attr-interface-defs)
+add_public_tablegen_target(MLIRLoopAnnotationAttrInterfaceIncGen)
+add_dependencies(mlir-generic-headers MLIRLoopAnnotationAttrInterfaceIncGen)
+
set(LLVM_TARGET_DEFINITIONS DataLayoutInterfaces.td)
mlir_tablegen(DataLayoutAttrInterface.h.inc -gen-attr-interface-decls)
mlir_tablegen(DataLayoutAttrInterface.cpp.inc -gen-attr-interface-defs)
diff --git a/mlir/include/mlir/Interfaces/LoopAnnotationAttrInterface.h b/mlir/include/mlir/Interfaces/LoopAnnotationAttrInterface.h
new file mode 100644
index 0000000000000..8584a39470192
--- /dev/null
+++ b/mlir/include/mlir/Interfaces/LoopAnnotationAttrInterface.h
@@ -0,0 +1,15 @@
+//===- LoopAnnotationAttrInterface.h ----------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_INTERFACES_LOOPANNOTATIONINTERFACE_H
+#define MLIR_INTERFACES_LOOPANNOTATIONINTERFACE_H
+
+#include "mlir/IR/Attributes.h"
+#include "mlir/Interfaces/LoopAnnotationAttrInterface.h.inc"
+
+#endif // MLIR_INTERFACES_LOOPANNOTATIONINTERFACE_H
diff --git a/mlir/include/mlir/Interfaces/LoopAnnotationAttrInterface.td b/mlir/include/mlir/Interfaces/LoopAnnotationAttrInterface.td
new file mode 100644
index 0000000000000..657cb490353f8
--- /dev/null
+++ b/mlir/include/mlir/Interfaces/LoopAnnotationAttrInterface.td
@@ -0,0 +1,28 @@
+//===- LoopAnnotationAttrInterface.td ----------------------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Defines the interface for attributes on loop-like operations.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_INTERFACES_LOOPANNOTATIONATTRINTERFACE
+#define MLIR_INTERFACES_LOOPANNOTATIONATTRINTERFACE
+
+include "mlir/IR/OpBase.td"
+
+def LoopAnnotationAttrInterface : AttrInterface<"LoopAnnotationAttrInterface"> {
+ let description = [{
+ Metadata that should live on loop-like operations.
+ }];
+ let cppNamespace = "::mlir";
+
+ let methods = [
+ ];
+}
+
+#endif // MLIR_INTERFACES_LOOPANNOTATIONATTRINTERFACE
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.h b/mlir/include/mlir/Interfaces/LoopLikeInterface.h
index 9925fc6ce6ca9..ed6579810386b 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.h
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.h
@@ -14,6 +14,7 @@
#define MLIR_INTERFACES_LOOPLIKEINTERFACE_H_
#include "mlir/IR/OpDefinition.h"
+#include "mlir/Interfaces/LoopAnnotationAttrInterface.h"
namespace mlir {
class RewriterBase;
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
index 6c95b4802837b..562a825799764 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -14,6 +14,7 @@
#define MLIR_INTERFACES_LOOPLIKEINTERFACE
include "mlir/IR/OpBase.td"
+include "mlir/Interfaces/LoopAnnotationAttrInterface.td"
//===----------------------------------------------------------------------===//
// Interfaces
@@ -232,6 +233,22 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
/*defaultImplementation=*/[{
return ::mlir::failure();
}]
+ >,
+ InterfaceMethod<[{
+ Returns the loop annotation attributes.
+ }],
+ /*retTy=*/"::mlir::SmallVector<::mlir::NamedAttribute>",
+ /*methodName=*/"getLoopAnnotationAttributes",
+ /*args=*/(ins),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ SmallVector<NamedAttribute> loopAnnotationAttrs;
+ llvm::copy_if($_op->getAttrs(), std::back_inserter(loopAnnotationAttrs),
+ [](auto attr) -> bool {
+ return mlir::isa<mlir::LoopAnnotationAttrInterface>(attr.getValue());
+ });
+ return loopAnnotationAttrs;
+ }]
>
];
diff --git a/mlir/lib/Conversion/SCFToControlFlow/CMakeLists.txt b/mlir/lib/Conversion/SCFToControlFlow/CMakeLists.txt
index 9c635ad7e664f..47366c38bb1d8 100644
--- a/mlir/lib/Conversion/SCFToControlFlow/CMakeLists.txt
+++ b/mlir/lib/Conversion/SCFToControlFlow/CMakeLists.txt
@@ -10,7 +10,6 @@ add_mlir_conversion_library(MLIRSCFToControlFlow
LINK_LIBS PUBLIC
MLIRArithDialect
MLIRControlFlowDialect
- MLIRLLVMDialect
MLIRSCFDialect
MLIRSCFTransforms
MLIRTransforms
diff --git a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
index 114d634629d77..405aae69bd05b 100644
--- a/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
+++ b/mlir/lib/Conversion/SCFToControlFlow/SCFToControlFlow.cpp
@@ -15,7 +15,6 @@
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
-#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Transforms/Transforms.h"
#include "mlir/IR/Builders.h"
@@ -377,12 +376,8 @@ LogicalResult ForLowering::matchAndRewrite(ForOp forOp,
// Let the CondBranchOp carry the LLVM attributes from the ForOp, such as the
// llvm.loop_annotation attribute.
- SmallVector<NamedAttribute> llvmAttrs;
- llvm::copy_if(forOp->getAttrs(), std::back_inserter(llvmAttrs),
- [](auto attr) {
- return isa<LLVM::LLVMDialect>(attr.getValue().getDialect());
- });
- condBranchOp->setDiscardableAttrs(llvmAttrs);
+ condBranchOp->setDiscardableAttrs(forOp.getLoopAnnotationAttributes());
+
// The result of the loop operation is the values of the condition block
// arguments except the induction variable on the last iteration.
rewriter.replaceOp(forOp, conditionBlock->getArguments().drop_front());
diff --git a/mlir/lib/Interfaces/CMakeLists.txt b/mlir/lib/Interfaces/CMakeLists.txt
index a25694cfff5f2..796385721ca77 100644
--- a/mlir/lib/Interfaces/CMakeLists.txt
+++ b/mlir/lib/Interfaces/CMakeLists.txt
@@ -10,6 +10,7 @@ set(LLVM_OPTIONAL_SOURCES
FunctionInterfaces.cpp
InferIntRangeInterface.cpp
InferTypeOpInterface.cpp
+ LoopAnnotationAttrInterface.cpp
LoopLikeInterface.cpp
MemorySlotInterfaces.cpp
ParallelCombiningOpInterface.cpp
@@ -65,6 +66,19 @@ add_mlir_library(MLIRFunctionInterfaces
add_mlir_interface_library(InferIntRangeInterface)
add_mlir_interface_library(InferTypeOpInterface)
+add_mlir_library(MLIRLoopAnnotationAttrInterface
+ LoopAnnotationAttrInterface.cpp
+
+ ADDITIONAL_HEADER_DIRS
+ ${MLIR_MAIN_INCLUDE_DIR}/mlir/Interfaces
+
+ DEPENDS
+ MLIRLoopAnnotationAttrInterfaceIncGen
+
+ LINK_LIBS PUBLIC
+ MLIRIR
+)
+
add_mlir_library(MLIRLoopLikeInterface
LoopLikeInterface.cpp
diff --git a/mlir/lib/Interfaces/LoopAnnotationAttrInterface.cpp b/mlir/lib/Interfaces/LoopAnnotationAttrInterface.cpp
new file mode 100644
index 0000000000000..e8d019072277d
--- /dev/null
+++ b/mlir/lib/Interfaces/LoopAnnotationAttrInterface.cpp
@@ -0,0 +1,3 @@
+#include "mlir/Interfaces/LoopAnnotationAttrInterface.h"
+using namespace mlir;
+#include "mlir/Interfaces/LoopAnnotationAttrInterface.cpp.inc"
More information about the Mlir-commits
mailing list