MLIR SROA

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

@@ -6,7 +6,7 @@ include "mlir/Dialect/LLVMIR/LLVMAttrDefs.td"
 include "mlir/Dialect/LLVMIR/LLVMEnums.td"
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
-include "mlir/Interfaces/Mem2RegInterfaces.td"
+include "mlir/Interfaces/MemorySlotInterfaces.td"
 
 // Operations that correspond to LLVM intrinsics. With MLIR operation set being
 // extendable, there is no reason to introduce a hard boundary between "core"

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

@@ -22,7 +22,7 @@ include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
-include "mlir/Interfaces/Mem2RegInterfaces.td"
+include "mlir/Interfaces/MemorySlotInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 class LLVM_Builder<string builder> {
@@ -173,7 +173,8 @@ def LLVM_FNegOp : LLVM_UnaryFloatArithmeticOp<
 
 // Memory-related operations.
 def LLVM_AllocaOp : LLVM_Op<"alloca", 
-    [DeclareOpInterfaceMethods<PromotableAllocationOpInterface>]>,
+    [DeclareOpInterfaceMethods<PromotableAllocationOpInterface>,
+     DeclareOpInterfaceMethods<DestructurableAllocationOpInterface>]>,
   LLVM_MemOpPatterns {
   let arguments = (ins AnyInteger:$arraySize,
                    OptionalAttr<I64Attr>:$alignment,
@@ -232,7 +233,9 @@ def LLVM_AllocaOp : LLVM_Op<"alloca",
 }
 
 def LLVM_GEPOp : LLVM_Op<"getelementptr", [Pure,
-    DeclareOpInterfaceMethods<PromotableOpInterface>]> {
+    DeclareOpInterfaceMethods<PromotableOpInterface>,
+    DeclareOpInterfaceMethods<TypeSafeOpInterface>,
+    DeclareOpInterfaceMethods<DestructurableAccessorOpInterface>]> {
   let arguments = (ins LLVM_ScalarOrVectorOf<LLVM_AnyPointer>:$base,
                    Variadic<LLVM_ScalarOrVectorOf<AnyInteger>>:$dynamicIndices,
                    DenseI32ArrayAttr:$rawConstantIndices,
@@ -316,7 +319,8 @@ def LLVM_GEPOp : LLVM_Op<"getelementptr", [Pure,
 }
 
 def LLVM_LoadOp : LLVM_MemAccessOpBase<"load",
-    [DeclareOpInterfaceMethods<PromotableMemOpInterface>]> {
+    [DeclareOpInterfaceMethods<PromotableMemOpInterface>,
+     DeclareOpInterfaceMethods<TypeSafeOpInterface>]> {
   dag args = (ins Arg<LLVM_PointerTo<LLVM_LoadableType>, "", [MemRead]>:$addr,
               OptionalAttr<I64Attr>:$alignment,
               UnitAttr:$volatile_,
@@ -388,7 +392,8 @@ def LLVM_LoadOp : LLVM_MemAccessOpBase<"load",
 }
 
 def LLVM_StoreOp : LLVM_MemAccessOpBase<"store",
-    [DeclareOpInterfaceMethods<PromotableMemOpInterface>]> {
+    [DeclareOpInterfaceMethods<PromotableMemOpInterface>,
+     DeclareOpInterfaceMethods<TypeSafeOpInterface>]> {
   dag args = (ins LLVM_LoadableType:$value,
               Arg<LLVM_PointerTo<LLVM_LoadableType>,"",[MemWrite]>:$addr,
               OptionalAttr<I64Attr>:$alignment,

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h

@@ -16,6 +16,7 @@
 
 #include "mlir/IR/Types.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
+#include "mlir/Interfaces/MemorySlotInterfaces.h"
 #include <optional>
 
 namespace llvm {
@@ -103,6 +104,7 @@ DEFINE_TRIVIAL_LLVM_TYPE(LLVMMetadataType);
 class LLVMStructType
     : public Type::TypeBase<LLVMStructType, Type, detail::LLVMStructTypeStorage,
                             DataLayoutTypeInterface::Trait,
+                            DestructurableTypeInterface::Trait,
                             TypeTrait::IsMutable> {
 public:
   /// Inherit base constructors.
@@ -198,6 +200,12 @@ class LLVMStructType
 
   LogicalResult verifyEntries(DataLayoutEntryListRef entries,
                               Location loc) const;
+
+  /// Destructs the struct into its indexed field types.
+  Optional<DenseMap<Attribute, Type>> getSubelementIndexMap();
+
+  /// Returns which type is stored at a given integer index within the struct.
+  Type getTypeAtIndex(Attribute index);
 };
 
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.td

@@ -12,6 +12,7 @@
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/Interfaces/DataLayoutInterfaces.td"
+include "mlir/Interfaces/MemorySlotInterfaces.td"
 
 /// Base class for all LLVM dialect types.
 class LLVMType<string typeName, string typeMnemonic, list<Trait> traits = []>
@@ -24,7 +25,8 @@ class LLVMType<string typeName, string typeMnemonic, list<Trait> traits = []>
 //===----------------------------------------------------------------------===//
 
 def LLVMArrayType : LLVMType<"LLVMArray", "array", [
-    DeclareTypeInterfaceMethods<DataLayoutTypeInterface, ["getTypeSize"]>]> {
+    DeclareTypeInterfaceMethods<DataLayoutTypeInterface, ["getTypeSize"]>,
+    DeclareTypeInterfaceMethods<DestructurableTypeInterface>]> {
   let summary = "LLVM array type";
   let description = [{
     The `!llvm.array` type represents a fixed-size array of element types.

mlir/include/mlir/Interfaces/CMakeLists.txt

@@ -7,7 +7,6 @@ add_mlir_interface(DestinationStyleOpInterface)
 add_mlir_interface(InferIntRangeInterface)
 add_mlir_interface(InferTypeOpInterface)
 add_mlir_interface(LoopLikeInterface)
-add_mlir_interface(Mem2RegInterfaces)
 add_mlir_interface(ParallelCombiningOpInterface)
 add_mlir_interface(RuntimeVerifiableOpInterface)
 add_mlir_interface(ShapedOpInterfaces)
@@ -17,6 +16,14 @@ add_mlir_interface(ValueBoundsOpInterface)
 add_mlir_interface(VectorInterfaces)
 add_mlir_interface(ViewLikeInterface)
 
+set(LLVM_TARGET_DEFINITIONS MemorySlotInterfaces.td)
+mlir_tablegen(MemorySlotOpInterfaces.h.inc -gen-op-interface-decls)
+mlir_tablegen(MemorySlotOpInterfaces.cpp.inc -gen-op-interface-defs)
+mlir_tablegen(MemorySlotTypeInterfaces.h.inc -gen-type-interface-decls)
+mlir_tablegen(MemorySlotTypeInterfaces.cpp.inc -gen-type-interface-defs)
+add_public_tablegen_target(MLIRMemorySlotInterfacesIncGen)
+add_dependencies(mlir-generic-headers MLIRMemorySlotInterfacesIncGen)
+
 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/Mem2RegInterfaces.h → mlir/include/mlir/Interfaces/MemorySlotInterfaces.h

@@ -6,8 +6,8 @@
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MLIR_INTERFACES_MEM2REGINTERFACES_H
-#define MLIR_INTERFACES_MEM2REGINTERFACES_H
+#ifndef MLIR_INTERFACES_MEMORYSLOTINTERFACES_H
+#define MLIR_INTERFACES_MEMORYSLOTINTERFACES_H
 
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/OpDefinition.h"
@@ -23,6 +23,13 @@ struct MemorySlot {
   Type elemType;
 };
 
+/// Memory slot attached with information about its destructuring procedure.
+struct DestructurableMemorySlot : MemorySlot {
+  /// Maps an index within the memory slot to the type of the pointer that
+  /// will be generated to access the element directly.
+  DenseMap<Attribute, Type> elementPtrs;
+};
+
 /// Returned by operation promotion logic requesting the deletion of an
 /// operation.
 enum class DeletionKind {
@@ -34,6 +41,7 @@ enum class DeletionKind {
 
 } // namespace mlir
 
-#include "mlir/Interfaces/Mem2RegInterfaces.h.inc"
+#include "mlir/Interfaces/MemorySlotOpInterfaces.h.inc"
+#include "mlir/Interfaces/MemorySlotTypeInterfaces.h.inc"
 
-#endif // MLIR_INTERFACES_MEM2REGINTERFACES_H
+#endif // MLIR_INTERFACES_MEMORYSLOTINTERFACES_H

mlir/include/mlir/Interfaces/Mem2RegInterfaces.td → mlir/include/mlir/Interfaces/MemorySlotInterfaces.td

@@ -1,13 +1,13 @@
-//===-- Mem2RegInterfaces.td - Mem2Reg interfaces ----------*- tablegen -*-===//
+//===-- MemorySlotInterfaces.td - MemorySlot interfaces ----*- 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
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MLIR_INTERFACES_MEM2REGINTERFACES
-#define MLIR_INTERFACES_MEM2REGINTERFACES
+#ifndef MLIR_INTERFACES_MEMORYSLOTINTERFACES
+#define MLIR_INTERFACES_MEMORYSLOTINTERFACES
 
 include "mlir/IR/OpBase.td"
 
@@ -76,6 +76,9 @@ def PromotableMemOpInterface : OpInterface<"PromotableMemOpInterface"> {
     to memory slots. Loads and stores must be of whole values of the same
     type as the slot itself.
 
+    For a memory operation on a slot to be valid, it must operate on the slot
+    pointer *only as a pointer to an element of the type of the slot*.
+
     If the same operation does both loads and stores on the same slot, the
     load must semantically happen first.
   }];
@@ -152,21 +155,21 @@ def PromotableOpInterface : OpInterface<"PromotableOpInterface"> {
   let methods = [
     InterfaceMethod<[{
         Checks that this operation can be promoted to no longer use the provided
-        blocking uses, in the context of promoting `slot`.
+        blocking uses, in order to allow optimization.
 
         If the removal procedure of the use will require that other uses get
         removed, that dependency should be added to the `newBlockingUses`
         argument. Dependent uses must only be uses of results of this operation.
       }], "bool", "canUsesBeRemoved",
-      (ins "const ::mlir::MemorySlot &":$slot,
-           "const ::llvm::SmallPtrSetImpl<::mlir::OpOperand *> &":$blockingUses,
+      (ins "const ::llvm::SmallPtrSetImpl<::mlir::OpOperand *> &":$blockingUses,
            "::llvm::SmallVectorImpl<::mlir::OpOperand *> &":$newBlockingUses)
     >,
     InterfaceMethod<[{
         Transforms IR to ensure that the current operation does not use the
-        provided memory slot anymore. In contrast to `PromotableMemOpInterface`,
-        operations implementing this interface must not need access to the
-        reaching definition of the content of the slot.
+        provided blocking uses anymore. In contrast to
+        `PromotableMemOpInterface`, operations implementing this interface
+        must not need access to the reaching definition of the content of the
+        slot.
 
         During the transformation, *no operation should be deleted*.
         The operation can only schedule its own deletion by returning the
@@ -186,11 +189,148 @@ def PromotableOpInterface : OpInterface<"PromotableOpInterface"> {
       }],
       "::mlir::DeletionKind",
       "removeBlockingUses",
-      (ins "const ::mlir::MemorySlot &":$slot,
-           "const ::llvm::SmallPtrSetImpl<mlir::OpOperand *> &":$blockingUses,
+      (ins "const ::llvm::SmallPtrSetImpl<mlir::OpOperand *> &":$blockingUses,
            "::mlir::OpBuilder &":$builder)
     >,
   ];
 }
 
-#endif // MLIR_INTERFACES_MEM2REGINTERFACES
+def DestructurableAllocationOpInterface
+  : OpInterface<"DestructurableAllocationOpInterface"> {
+  let description = [{
+    Describes operations allocating memory slots of aggregates that can be
+    destructured into multiple smaller allocations.
+  }];
+  let cppNamespace = "::mlir";
+
+  let methods = [
+    InterfaceMethod<[{
+        Returns the list of slots for which destructuring should be attempted,
+        specifying in which way the slot should be destructured into subslots.
+        The subslots are indexed by attributes. This computes the type of the
+        pointers of each subslots to be generated. The type of the memory slot
+        must implement `DestructurableTypeInterface`.
+      }],
+      "::llvm::SmallVector<::mlir::DestructurableMemorySlot>",
+      "getDestructurableSlots",
+      (ins)
+    >,
+    InterfaceMethod<[{
+        Destructures this slot into multiple subslots. The newly generated slots
+        may belong to a different allocator. The original slot must still exist
+        at the end of this call.
+
+        The builder is located at the beginning of the block where the slot
+        pointer is defined.
+      }],
+      "::llvm::DenseMap<::mlir::Attribute, ::mlir::MemorySlot>",
+      "destructure",
+      (ins "const ::mlir::DestructurableMemorySlot &":$slot,
+           "::llvm::SmallPtrSetImpl<::mlir::Attribute> &":$usedIndices,
+           "::mlir::OpBuilder &":$builder)
+    >,
+    InterfaceMethod<[{
+        Hook triggered once the destructuring of a slot is complete, meaning the
+        original slot is no longer being refered to and could be deleted.
+        This will only be called for slots declared by this operation.
+      }],
+      "void", "handleDestructuringComplete",
+      (ins "const ::mlir::DestructurableMemorySlot &":$slot)
+    >,
+  ];
+}
+
+def TypeSafeOpInterface : OpInterface<"TypeSafeOpInterface"> {
+  let description = [{
+    Describes operations using memory slots in a type-safe manner.
+  }];
+  let cppNamespace = "::mlir";
+
+  let methods = [
+    InterfaceMethod<[{
+        Returns whether all accesses in this operation to the provided slot are
+        done in a type-safe manner. To be type-safe, the access must only load
+        the value in this type as the type of the slot, and without assuming any
+        context around the slot. For example, a type-safe load must not load
+        outside the bounds of the slot.
+
+        If the type-safety of the accesses depends on the type-safety of the
+        accesses to further memory slots, the result of this method will be
+        conditioned to the type-safety of the accesses to the slots added by
+        this method to `mustBeSafelyUsed`.
+      }],
+      "::mlir::LogicalResult",
+      "ensureOnlyTypeSafeAccesses",
+      (ins "const ::mlir::MemorySlot &":$slot,
+           "::mlir::SmallVectorImpl<::mlir::MemorySlot> &":$mustBeSafelyUsed)
+    >
+  ];
+}
+
+def DestructurableAccessorOpInterface
+  : OpInterface<"DestructurableAccessorOpInterface"> {
+  let description = [{
+    Describes operations that can access a sub-element of a destructurable slot.
+  }];
+  let cppNamespace = "::mlir";
+
+  let methods = [
+    InterfaceMethod<[{
+        For a given destructurable memory slot, returns whether this operation can
+        rewire its uses of the slot to use the slots generated after
+        destructuring. This may involve creating new operations, and usually
+        amounts to checking if the pointer types match.
+
+        This method must also register the indices it will access within the
+        `usedIndices` set. If the accessor generates new slots mapping to
+        subelements, they must be registered in `mustBeSafelyUsed` to ensure
+        they are used in a locally type-safe manner.
+      }],
+      "bool",
+      "canRewire",
+      (ins "const ::mlir::DestructurableMemorySlot &":$slot,
+           "::llvm::SmallPtrSetImpl<::mlir::Attribute> &":$usedIndices,
+           "::mlir::SmallVectorImpl<::mlir::MemorySlot> &":$mustBeSafelyUsed)
+    >,
+    InterfaceMethod<[{
+        Rewires the use of a slot to the generated subslots, without deleting
+        any operation. Returns whether the accessor should be deleted.
+      }],
+      "::mlir::DeletionKind",
+      "rewire",
+      (ins "const ::mlir::DestructurableMemorySlot &":$slot,
+           "::llvm::DenseMap<::mlir::Attribute, ::mlir::MemorySlot> &":$subslots)
+    >
+  ];
+}
+
+def DestructurableTypeInterface
+  : TypeInterface<"DestructurableTypeInterface"> {
+  let description = [{
+    Describes a type that can be broken down into indexable sub-element types.
+  }];
+  let cppNamespace = "::mlir";
+
+  let methods = [
+    InterfaceMethod<[{
+        Destructures the type into subelements into a map of index attributes to
+        types of subelements. Returns nothing if the type cannot be destructured.
+      }],
+      "::llvm::Optional<::llvm::DenseMap<::mlir::Attribute, ::mlir::Type>>",
+      "getSubelementIndexMap",
+      (ins)
+    >,
+    InterfaceMethod<[{
+        Indicates which type is held at the provided index, returning a null
+        Type if no type could be computed. While this can return information
+        even when the type cannot be completely destructured, it must be coherent
+        with the types returned by `getSubelementIndexMap` when they exist.
+      }],
+      "::mlir::Type",
+      "getTypeAtIndex",
+      (ins "::mlir::Attribute":$index)
+    >
+  ];
+}
+
+#endif // MLIR_INTERFACES_MEMORYSLOTINTERFACES