[mlir][memref][llvm] Infer llvm alias scopes attrs from memref.distinct_objects

mlir/include/mlir/Dialect/LLVMIR/Transforms/InferAliasScopeAttrs.h

@@ -0,0 +1,22 @@
+//==- InferAliasScopeAttrs.h - Infer LLVM alias scope attributes -*- 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_DIALECT_LLVMIR_TRANSFORMS_INFERALIASSCOPEATTRS_H
+#define MLIR_DIALECT_LLVMIR_TRANSFORMS_INFERALIASSCOPEATTRS_H
+
+#include <memory>
+
+namespace mlir {
+class Pass;
+namespace LLVM {
+#define GEN_PASS_DECL_LLVMINFERALIASSCOPEATTRIBUTES
+#include "mlir/Dialect/LLVMIR/Transforms/Passes.h.inc"
+} // namespace LLVM
+} // namespace mlir
+
+#endif // MLIR_DIALECT_LLVMIR_TRANSFORMS_INFERALIASSCOPEATTRS_H

mlir/include/mlir/Dialect/LLVMIR/Transforms/Passes.td

@@ -73,4 +73,20 @@ def DIScopeForLLVMFuncOpPass : Pass<"ensure-debug-info-scope-on-llvm-func", "::m
   ];
 }
 
+def LLVMInferAliasScopeAttributes
+    : InterfacePass<"llvm-infer-alias-scopes-attrs", "::mlir::FunctionOpInterface"> {
+  let summary = "Infer alias scopes for load/store ops";
+  let description = [{
+    This pass infers LLVM alias scope attributes for other dialect load and store
+    operations based on alias analysis information. It analyzes memory operations
+    within functions to determine which operations may or may not alias with each
+    other, and annotates them with appropriate LLVM alias scope metadata.
+
+    Alias scopes on other dialect ops are not robust enough against code
+    transformations (e.g. inlining) so this pass is intended to be run just
+    before *-to-llvm conversion.
+  }];
+  let dependentDialects = ["LLVM::LLVMDialect"];
+}
+
 #endif // MLIR_DIALECT_LLVMIR_TRANSFORMS_PASSES

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

@@ -153,7 +153,7 @@ def AssumeAlignmentOp : MemRef_Op<"assume_alignment", [
       The `assume_alignment` operation takes a memref and an integer alignment
       value. It returns a new SSA value of the same memref type, but associated
       with the assumption that the underlying buffer is aligned to the given
-      alignment. 
+      alignment.
 
       If the buffer isn't aligned to the given alignment, its result is poison.
       This operation doesn't affect the semantics of a program where the
@@ -168,14 +168,50 @@ def AssumeAlignmentOp : MemRef_Op<"assume_alignment", [
   let assemblyFormat = "$memref `,` $alignment attr-dict `:` type($memref)";
   let extraClassDeclaration = [{
     MemRefType getType() { return ::llvm::cast<MemRefType>(getResult().getType()); }
-    
+
     Value getViewSource() { return getMemref(); }
   }];
 
   let hasVerifier = 1;
   let hasFolder = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// DistinctObjectsOp
+//===----------------------------------------------------------------------===//
+
+def DistinctObjectsOp : MemRef_Op<"distinct_objects", [
+      Pure,
+      DistinctObjectsInterface,
+      DeclareOpInterfaceMethods<InferTypeOpInterface>
+      // ViewLikeOpInterface TODO: ViewLikeOpInterface only supports a single argument
+    ]> {
+  let summary = "assumption that acesses to specific memrefs will never alias";
+  let description = [{
+      The `distinct_objects` operation takes a list of memrefs and returns the same
+      memrefs, with the additional assumption that accesses to them will never
+      alias with each other. This means that loads and stores to different
+      memrefs in the list can be safely reordered.
+
+      If the memrefs do alias, the load/store behavior is undefined. This
+      operation doesn't affect the semantics of a valid program. It is
+      intended for optimization purposes, allowing the compiler to generate more
+      efficient code based on the non-aliasing assumption. The optimization is
+      best-effort.
+
+      Example:
+
+      ```mlir
+      %1, %2 = memref.distinct_objects %a, %b : memref<?xf32>, memref<?xf32>
+      ```
+  }];
+  let arguments = (ins Variadic<AnyMemRef>:$operands);
+  let results = (outs Variadic<AnyMemRef>:$results);
+
+  let assemblyFormat = "$operands attr-dict `:` type($operands)";
+  let hasVerifier = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // AllocOp
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Interfaces/ViewLikeInterface.td

@@ -414,4 +414,44 @@ def OffsetSizeAndStrideOpInterface : OpInterface<"OffsetSizeAndStrideOpInterface
   }];
 }
 
+def DistinctObjectsInterface : OpInterface<"DistinctObjectsInterface"> {
+  let description = [{
+    This intefaces indicates that pointer-like objects (such as memrefs) returned
+    from this operation will never alias with each other. This provides a
+    guarantee to optimization passes that accesses through different results
+    of this operation can be safely reordered, as they will never reference
+    overlapping memory locations.
+
+    Operations with this interface take multiple pointer-like operands
+    and return the same operands with additional non-aliasing guarantees.
+    If the access to the results of this operation aliases at runtime, the
+    behavior of such access is undefined.
+  }];
+  let cppNamespace = "::mlir";
+
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/[{ Return input pointer-like objects. }],
+      /*retTy=*/"::mlir::ValueRange",
+      /*methodName=*/"getDistinctOperands",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return $_op->getOperands();
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{ Return result pointer-like objects. }],
+      /*retTy=*/"::mlir::ValueRange",
+      /*methodName=*/"getDistinctResults",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return $_op->getResults();
+      }]
+    >
+  ];
+}
+
+
 #endif // MLIR_INTERFACES_VIEWLIKEINTERFACE

mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp

@@ -258,6 +258,33 @@ getAllocEffectFor(Value value,
   return success();
 }
 
+static Value getDistinctObjectsOperand(DistinctObjectsInterface op,
+                                       Value value) {
+  unsigned argNumber = cast<OpResult>(value).getResultNumber();
+  return op.getDistinctOperands()[argNumber];
+}
+
+static std::optional<AliasResult> checkDistinctObjects(Value lhs, Value rhs) {
+  // We should already checked that lhs and rhs are different.
+  assert(lhs != rhs && "lhs and rhs must be different");
+
+  // Result and corresponding operand must alias.
+  auto lhsOp = lhs.getDefiningOp<DistinctObjectsInterface>();
+  if (lhsOp && getDistinctObjectsOperand(lhsOp, lhs) == rhs)
+    return AliasResult::MustAlias;
+
+  auto rhsOp = rhs.getDefiningOp<DistinctObjectsInterface>();
+  if (rhsOp && getDistinctObjectsOperand(rhsOp, rhs) == lhs)
+    return AliasResult::MustAlias;
+
+  // If two different values come from the same `DistinctObjects` operation,
+  // they don't alias.
+  if (lhsOp && lhsOp == rhsOp)
+    return AliasResult::NoAlias;
+
+  return std::nullopt;
+}
+
 /// Given the two values, return their aliasing behavior.
 AliasResult LocalAliasAnalysis::aliasImpl(Value lhs, Value rhs) {
   if (lhs == rhs)
@@ -289,6 +316,9 @@ AliasResult LocalAliasAnalysis::aliasImpl(Value lhs, Value rhs) {
                : AliasResult::MayAlias;
   }
 
+  if (std::optional<AliasResult> result = checkDistinctObjects(lhs, rhs))
+    return *result;
+
   // Otherwise, neither of the values are constant so check to see if either has
   // an allocation effect.
   bool lhsHasAlloc = succeeded(getAllocEffectFor(lhs, lhsAlloc, lhsAllocScope));

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

@@ -465,6 +465,50 @@ struct AssumeAlignmentOpLowering
   }
 };
 
+struct DistinctObjectsOpLowering
+    : public ConvertOpToLLVMPattern<memref::DistinctObjectsOp> {
+  using ConvertOpToLLVMPattern<
+      memref::DistinctObjectsOp>::ConvertOpToLLVMPattern;
+  explicit DistinctObjectsOpLowering(const LLVMTypeConverter &converter)
+      : ConvertOpToLLVMPattern<memref::DistinctObjectsOp>(converter) {}
+
+  LogicalResult
+  matchAndRewrite(memref::DistinctObjectsOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    ValueRange operands = adaptor.getOperands();
+    if (operands.size() <= 1) {
+      // Fast path.
+      rewriter.replaceOp(op, operands);
+      return success();
+    }
+
+    Location loc = op.getLoc();
+    SmallVector<Value> ptrs;
+    for (auto [origOperand, newOperand] :
+         llvm::zip_equal(op.getOperands(), operands)) {
+      auto memrefType = cast<MemRefType>(origOperand.getType());
+      Value ptr = getStridedElementPtr(rewriter, loc, memrefType, newOperand,
+                                       /*indices=*/{});
+      ptrs.push_back(ptr);
+    }
+
+    auto cond =
+        LLVM::ConstantOp::create(rewriter, loc, rewriter.getI1Type(), 1);
+    // Generate separate_storage assumptions for each pair of pointers.
+    for (auto i : llvm::seq<size_t>(ptrs.size() - 1)) {
+      for (auto j : llvm::seq<size_t>(i + 1, ptrs.size())) {
+        Value ptr1 = ptrs[i];
+        Value ptr2 = ptrs[j];
+        LLVM::AssumeOp::create(rewriter, loc, cond,
+                               LLVM::AssumeSeparateStorageTag{}, ptr1, ptr2);
+      }
+    }
+
+    rewriter.replaceOp(op, operands);
+    return success();
+  }
+};
+
 // A `dealloc` is converted into a call to `free` on the underlying data buffer.
 // The memref descriptor being an SSA value, there is no need to clean it up
 // in any way.
@@ -878,6 +922,18 @@ struct GetGlobalMemrefOpLowering
   }
 };
 
+static SmallVector<std::pair<StringAttr, Attribute>>
+copyImportantAttrs(Operation *op) {
+  SmallVector<std::pair<StringAttr, Attribute>> attrs;
+  for (StringRef attrName : {LLVM::LLVMDialect::getAliasScopesAttrName(),
+                             LLVM::LLVMDialect::getNoAliasAttrName()}) {
+    auto nameAttr = StringAttr::get(op->getContext(), attrName);
+    if (auto attr = op->getAttr(nameAttr))
+      attrs.emplace_back(nameAttr, attr);
+  }
+  return attrs;
+}
+
 // Load operation is lowered to obtaining a pointer to the indexed element
 // and loading it.
 struct LoadOpLowering : public LoadStoreOpLowering<memref::LoadOp> {
@@ -888,15 +944,22 @@ struct LoadOpLowering : public LoadStoreOpLowering<memref::LoadOp> {
                   ConversionPatternRewriter &rewriter) const override {
     auto type = loadOp.getMemRefType();
 
+    SmallVector<std::pair<StringAttr, Attribute>> importantAttrs =
+        copyImportantAttrs(loadOp);
+
     // Per memref.load spec, the indices must be in-bounds:
     // 0 <= idx < dim_size, and additionally all offsets are non-negative,
     // hence inbounds and nuw are used when lowering to llvm.getelementptr.
     Value dataPtr = getStridedElementPtr(rewriter, loadOp.getLoc(), type,
                                          adaptor.getMemref(),
                                          adaptor.getIndices(), kNoWrapFlags);
-    rewriter.replaceOpWithNewOp<LLVM::LoadOp>(
+    auto newOp = rewriter.replaceOpWithNewOp<LLVM::LoadOp>(
         loadOp, typeConverter->convertType(type.getElementType()), dataPtr,
         loadOp.getAlignment().value_or(0), false, loadOp.getNontemporal());
+
+    for (auto [nameAttr, attr] : importantAttrs)
+      newOp->setAttr(nameAttr, attr);
+
     return success();
   }
 };
@@ -911,15 +974,22 @@ struct StoreOpLowering : public LoadStoreOpLowering<memref::StoreOp> {
                   ConversionPatternRewriter &rewriter) const override {
     auto type = op.getMemRefType();
 
+    SmallVector<std::pair<StringAttr, Attribute>> importantAttrs =
+        copyImportantAttrs(op);
+
     // Per memref.store spec, the indices must be in-bounds:
     // 0 <= idx < dim_size, and additionally all offsets are non-negative,
     // hence inbounds and nuw are used when lowering to llvm.getelementptr.
     Value dataPtr =
         getStridedElementPtr(rewriter, op.getLoc(), type, adaptor.getMemref(),
                              adaptor.getIndices(), kNoWrapFlags);
-    rewriter.replaceOpWithNewOp<LLVM::StoreOp>(op, adaptor.getValue(), dataPtr,
-                                               op.getAlignment().value_or(0),
-                                               false, op.getNontemporal());
+    auto newOp = rewriter.replaceOpWithNewOp<LLVM::StoreOp>(
+        op, adaptor.getValue(), dataPtr, op.getAlignment().value_or(0), false,
+        op.getNontemporal());
+
+    for (auto [nameAttr, attr] : importantAttrs)
+      newOp->setAttr(nameAttr, attr);
+
     return success();
   }
 };
@@ -1997,22 +2067,23 @@ void mlir::populateFinalizeMemRefToLLVMConversionPatterns(
   patterns.add<
       AllocaOpLowering,
       AllocaScopeOpLowering,
-      AtomicRMWOpLowering,
       AssumeAlignmentOpLowering,
+      AtomicRMWOpLowering,
       ConvertExtractAlignedPointerAsIndex,
       DimOpLowering,
+      DistinctObjectsOpLowering,
       ExtractStridedMetadataOpLowering,
       GenericAtomicRMWOpLowering,
       GetGlobalMemrefOpLowering,
       LoadOpLowering,
       MemRefCastOpLowering,
-      MemorySpaceCastOpLowering,
       MemRefReinterpretCastOpLowering,
       MemRefReshapeOpLowering,
+      MemorySpaceCastOpLowering,
       PrefetchOpLowering,
       RankOpLowering,
-      ReassociatingReshapeOpConversion<memref::ExpandShapeOp>,
       ReassociatingReshapeOpConversion<memref::CollapseShapeOp>,
+      ReassociatingReshapeOpConversion<memref::ExpandShapeOp>,
       StoreOpLowering,
       SubViewOpLowering,
       TransposeOpLowering,

mlir/lib/Dialect/LLVMIR/Transforms/CMakeLists.txt

@@ -3,6 +3,7 @@ add_mlir_dialect_library(MLIRLLVMIRTransforms
   DIExpressionLegalization.cpp
   DIExpressionRewriter.cpp
   DIScopeForLLVMFuncOp.cpp
+  InferAliasScopeAttrs.cpp
   InlinerInterfaceImpl.cpp
   LegalizeForExport.cpp
   OptimizeForNVVM.cpp