[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.

Author: ashermancinelli

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`

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"

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)

mlir/include/mlir/Interfaces/LoopAnnotationAttrInterface.h

@@ -0,0 +1,7 @@
+#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

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

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;

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;
+      }]
     >
   ];
 

mlir/lib/Conversion/SCFToControlFlow/CMakeLists.txt

@@ -10,7 +10,6 @@ add_mlir_conversion_library(MLIRSCFToControlFlow
   LINK_LIBS PUBLIC
   MLIRArithDialect
   MLIRControlFlowDialect
-  MLIRLLVMDialect
   MLIRSCFDialect
   MLIRSCFTransforms
   MLIRTransforms

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());

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