From 7a513df00552310c078f38d2d1a8aa20e9cf03e6 Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Thu, 4 Sep 2025 17:25:26 +0200 Subject: [PATCH 1/6] [mlir][memref] Introduce `memref.distinct_objects` op --- .../mlir/Dialect/MemRef/IR/MemRefOps.td | 39 ++++++++++++++- .../Conversion/MemRefToLLVM/MemRefToLLVM.cpp | 49 +++++++++++++++++-- mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp | 19 +++++++ .../MemRefToLLVM/memref-to-llvm.mlir | 19 +++++++ mlir/test/Dialect/MemRef/ops.mlir | 9 ++++ 5 files changed, 130 insertions(+), 5 deletions(-) diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td index 671cc05e963b4..933fb87de30ab 100644 --- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td +++ b/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,7 +168,7 @@ def AssumeAlignmentOp : MemRef_Op<"assume_alignment", [ let assemblyFormat = "$memref `,` $alignment attr-dict `:` type($memref)"; let extraClassDeclaration = [{ MemRefType getType() { return ::llvm::cast(getResult().getType()); } - + Value getViewSource() { return getMemref(); } }]; @@ -176,6 +176,41 @@ def AssumeAlignmentOp : MemRef_Op<"assume_alignment", [ let hasFolder = 1; } +//===----------------------------------------------------------------------===// +// DistinctObjectsOp +//===----------------------------------------------------------------------===// + +def DistinctObjectsOp : MemRef_Op<"distinct_objects", [ + Pure, + DeclareOpInterfaceMethods + // 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 a list of + memrefs of the same types, with the additional assumption that accesses to + these memrefs 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 behavior is undefined. This operation doesn't + affect the semantics of a program where the non-aliasing assumption holds + true. 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, memref + ``` + }]; + let arguments = (ins Variadic:$operands); + let results = (outs Variadic:$results); + + let assemblyFormat = "$operands attr-dict `:` type($operands)"; + let hasVerifier = 1; +} + //===----------------------------------------------------------------------===// // AllocOp //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp index 262e0e7a30c63..571e5000b3f51 100644 --- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp +++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp @@ -465,6 +465,48 @@ struct AssumeAlignmentOpLowering } }; +struct DistinctObjectsOpLowering + : public ConvertOpToLLVMPattern { + using ConvertOpToLLVMPattern< + memref::DistinctObjectsOp>::ConvertOpToLLVMPattern; + explicit DistinctObjectsOpLowering(const LLVMTypeConverter &converter) + : ConvertOpToLLVMPattern(converter) {} + + LogicalResult + matchAndRewrite(memref::DistinctObjectsOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + ValueRange operands = adaptor.getOperands(); + if (operands.empty()) { + rewriter.eraseOp(op); + return success(); + } + Location loc = op.getLoc(); + SmallVector ptrs; + for (auto [origOperand, newOperand] : + llvm::zip_equal(op.getOperands(), operands)) { + auto memrefType = cast(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(ptrs.size() - 1)) { + for (auto j : llvm::seq(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. @@ -1997,22 +2039,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, ReassociatingReshapeOpConversion, + ReassociatingReshapeOpConversion, StoreOpLowering, SubViewOpLowering, TransposeOpLowering, diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp index 5d15d5f6e3de4..62fb57f24a870 100644 --- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp +++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp @@ -542,6 +542,25 @@ OpFoldResult AssumeAlignmentOp::fold(FoldAdaptor adaptor) { return getMemref(); } +//===----------------------------------------------------------------------===// +// DistinctObjectsOp +//===----------------------------------------------------------------------===// + +LogicalResult DistinctObjectsOp::verify() { + if (getOperandTypes() != getResultTypes()) + return emitOpError("operand types and result types must match"); + return success(); +} + +LogicalResult DistinctObjectsOp::inferReturnTypes( + MLIRContext * /*context*/, std::optional /*location*/, + ValueRange operands, DictionaryAttr /*attributes*/, + OpaqueProperties /*properties*/, RegionRange /*regions*/, + SmallVectorImpl &inferredReturnTypes) { + llvm::copy(operands.getTypes(), std::back_inserter(inferredReturnTypes)); + return success(); +} + //===----------------------------------------------------------------------===// // CastOp //===----------------------------------------------------------------------===// diff --git a/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir b/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir index 45b1a1f1ca40c..3eb8df093af10 100644 --- a/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir +++ b/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir @@ -195,6 +195,25 @@ func.func @assume_alignment(%0 : memref<4x4xf16>) { // ----- +// ALL-LABEL: func @distinct_objects +// ALL-SAME: (%[[ARG0:.*]]: memref, %[[ARG1:.*]]: memref, %[[ARG2:.*]]: memref) +func.func @distinct_objects(%arg0: memref, %arg1: memref, %arg2: memref) -> (memref, memref, memref) { +// ALL-DAG: %[[CAST_0:.*]] = builtin.unrealized_conversion_cast %[[ARG0]] : memref to !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> +// ALL-DAG: %[[CAST_1:.*]] = builtin.unrealized_conversion_cast %[[ARG1]] : memref to !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> +// ALL-DAG: %[[CAST_2:.*]] = builtin.unrealized_conversion_cast %[[ARG2]] : memref to !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> +// ALL: %[[PTR_0:.*]] = llvm.extractvalue %[[CAST_0]][1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> +// ALL: %[[PTR_1:.*]] = llvm.extractvalue %[[CAST_1]][1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> +// ALL: %[[PTR_2:.*]] = llvm.extractvalue %[[CAST_2]][1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> +// ALL: %[[TRUE:.*]] = llvm.mlir.constant(true) : i1 +// ALL: llvm.intr.assume %[[TRUE]] ["separate_storage"(%[[PTR_0]], %[[PTR_1]] : !llvm.ptr, !llvm.ptr)] : i1 +// ALL: llvm.intr.assume %[[TRUE]] ["separate_storage"(%[[PTR_0]], %[[PTR_2]] : !llvm.ptr, !llvm.ptr)] : i1 +// ALL: llvm.intr.assume %[[TRUE]] ["separate_storage"(%[[PTR_1]], %[[PTR_2]] : !llvm.ptr, !llvm.ptr)] : i1 + %1, %2, %3 = memref.distinct_objects %arg0, %arg1, %arg2 : memref, memref, memref + return %1, %2, %3 : memref, memref, memref +} + +// ----- + // CHECK-LABEL: func @assume_alignment_w_offset // CHECK-INTERFACE-LABEL: func @assume_alignment_w_offset func.func @assume_alignment_w_offset(%0 : memref<4x4xf16, strided<[?, ?], offset: ?>>) { diff --git a/mlir/test/Dialect/MemRef/ops.mlir b/mlir/test/Dialect/MemRef/ops.mlir index 6c2298a3f8acb..a90c9505a8405 100644 --- a/mlir/test/Dialect/MemRef/ops.mlir +++ b/mlir/test/Dialect/MemRef/ops.mlir @@ -302,6 +302,15 @@ func.func @assume_alignment(%0: memref<4x4xf16>) { return } +// CHECK-LABEL: func @distinct_objects +// CHECK-SAME: (%[[ARG0:.*]]: memref, %[[ARG1:.*]]: memref, %[[ARG2:.*]]: memref) +func.func @distinct_objects(%arg0: memref, %arg1: memref, %arg2: memref) -> (memref, memref, memref) { + // CHECK: %[[RES:.*]]:3 = memref.distinct_objects %[[ARG0]], %[[ARG1]], %[[ARG2]] : memref, memref, memref + %1, %2, %3 = memref.distinct_objects %arg0, %arg1, %arg2 : memref, memref, memref + // CHECK: return %[[RES]]#0, %[[RES]]#1, %[[RES]]#2 : memref, memref, memref + return %1, %2, %3 : memref, memref, memref +} + // CHECK-LABEL: func @expand_collapse_shape_static func.func @expand_collapse_shape_static( %arg0: memref<3x4x5xf32>, From 1e24829dbc19f9d3afbb8d11cd1d9a334c9c2c80 Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Wed, 10 Sep 2025 17:25:38 +0200 Subject: [PATCH 2/6] verifier test --- mlir/test/Dialect/MemRef/invalid.mlir | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/mlir/test/Dialect/MemRef/invalid.mlir b/mlir/test/Dialect/MemRef/invalid.mlir index 3f96d907632b7..67951f8ef0765 100644 --- a/mlir/test/Dialect/MemRef/invalid.mlir +++ b/mlir/test/Dialect/MemRef/invalid.mlir @@ -1169,3 +1169,11 @@ func.func @expand_shape_invalid_output_shape( into memref<2x15x20xf32, strided<[60000, 4000, 2], offset: 100>> return } + +// ----- + +func.func @Invalid_distinct_objects(%arg0: memref, %arg1: memref) -> (memref, memref) { + // expected-error @+1 {{operand types and result types must match}} + %0, %1 = "memref.distinct_objects"(%arg0, %arg1) : (memref, memref) -> (memref, memref) + return %0, %1 : memref, memref +} From 4b943e82d1727e00512c35fb0034beeab7346e1e Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Tue, 23 Sep 2025 23:13:13 +0200 Subject: [PATCH 3/6] comments, reject 0-poerand version --- .../mlir/Dialect/MemRef/IR/MemRefOps.td | 20 +++++++++---------- .../Conversion/MemRefToLLVM/MemRefToLLVM.cpp | 6 ++++-- mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp | 4 ++++ .../MemRefToLLVM/memref-to-llvm.mlir | 11 ++++++++++ mlir/test/Dialect/MemRef/invalid.mlir | 10 +++++++++- 5 files changed, 38 insertions(+), 13 deletions(-) diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td index 933fb87de30ab..f75e311645426 100644 --- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td +++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td @@ -187,16 +187,16 @@ def DistinctObjectsOp : MemRef_Op<"distinct_objects", [ ]> { let summary = "assumption that acesses to specific memrefs will never alias"; let description = [{ - The `distinct_objects` operation takes a list of memrefs and returns a list of - memrefs of the same types, with the additional assumption that accesses to - these memrefs 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 behavior is undefined. This operation doesn't - affect the semantics of a program where the non-aliasing assumption holds - true. 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. + 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: diff --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp index 571e5000b3f51..64270726f4a01 100644 --- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp +++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp @@ -476,10 +476,12 @@ struct DistinctObjectsOpLowering matchAndRewrite(memref::DistinctObjectsOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { ValueRange operands = adaptor.getOperands(); - if (operands.empty()) { - rewriter.eraseOp(op); + if (operands.size() <= 1) { + // Fast path. + rewriter.replaceOp(op, operands); return success(); } + Location loc = op.getLoc(); SmallVector ptrs; for (auto [origOperand, newOperand] : diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp index 62fb57f24a870..0bca922b0c804 100644 --- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp +++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp @@ -549,6 +549,10 @@ OpFoldResult AssumeAlignmentOp::fold(FoldAdaptor adaptor) { LogicalResult DistinctObjectsOp::verify() { if (getOperandTypes() != getResultTypes()) return emitOpError("operand types and result types must match"); + + if (getOperandTypes().empty()) + return emitOpError("expected at least one operand"); + return success(); } diff --git a/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir b/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir index 3eb8df093af10..0cbe064572911 100644 --- a/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir +++ b/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir @@ -214,6 +214,17 @@ func.func @distinct_objects(%arg0: memref, %arg1: memref, %arg2: m // ----- +// ALL-LABEL: func @distinct_objects_noop +// ALL-SAME: (%[[ARG0:.*]]: memref) +func.func @distinct_objects_noop(%arg0: memref) -> memref { +// 1-operand version is noop +// ALL-NEXT: return %[[ARG0]] + %1 = memref.distinct_objects %arg0 : memref + return %1 : memref +} + +// ----- + // CHECK-LABEL: func @assume_alignment_w_offset // CHECK-INTERFACE-LABEL: func @assume_alignment_w_offset func.func @assume_alignment_w_offset(%0 : memref<4x4xf16, strided<[?, ?], offset: ?>>) { diff --git a/mlir/test/Dialect/MemRef/invalid.mlir b/mlir/test/Dialect/MemRef/invalid.mlir index 67951f8ef0765..5ff292058ccc1 100644 --- a/mlir/test/Dialect/MemRef/invalid.mlir +++ b/mlir/test/Dialect/MemRef/invalid.mlir @@ -1172,8 +1172,16 @@ func.func @expand_shape_invalid_output_shape( // ----- -func.func @Invalid_distinct_objects(%arg0: memref, %arg1: memref) -> (memref, memref) { +func.func @distinct_objects_types_mismatch(%arg0: memref, %arg1: memref) -> (memref, memref) { // expected-error @+1 {{operand types and result types must match}} %0, %1 = "memref.distinct_objects"(%arg0, %arg1) : (memref, memref) -> (memref, memref) return %0, %1 : memref, memref } + +// ----- + +func.func @distinct_objects_0_operands() { + // expected-error @+1 {{expected at least one operand}} + "memref.distinct_objects"() : () -> () + return +} From 98af92170063138588f98758bf8e8a75b24ae74d Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Wed, 24 Sep 2025 12:34:15 +0200 Subject: [PATCH 4/6] LocalAliasAnalysis --- .../mlir/Dialect/MemRef/IR/MemRefOps.td | 1 + .../mlir/Interfaces/ViewLikeInterface.td | 40 +++++++++++++++++++ .../AliasAnalysis/LocalAliasAnalysis.cpp | 30 ++++++++++++++ mlir/test/Analysis/test-alias-analysis.mlir | 16 ++++++++ 4 files changed, 87 insertions(+) diff --git a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td index f75e311645426..b4900720c6f74 100644 --- a/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td +++ b/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td @@ -182,6 +182,7 @@ def AssumeAlignmentOp : MemRef_Op<"assume_alignment", [ def DistinctObjectsOp : MemRef_Op<"distinct_objects", [ Pure, + DistinctObjectsInterface, DeclareOpInterfaceMethods // ViewLikeOpInterface TODO: ViewLikeOpInterface only supports a single argument ]> { diff --git a/mlir/include/mlir/Interfaces/ViewLikeInterface.td b/mlir/include/mlir/Interfaces/ViewLikeInterface.td index ed213bfdae337..2c3adb4a7fd6a 100644 --- a/mlir/include/mlir/Interfaces/ViewLikeInterface.td +++ b/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 diff --git a/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp b/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp index 8062b474539fd..15686fc6d4031 100644 --- a/mlir/lib/Analysis/AliasAnalysis/LocalAliasAnalysis.cpp +++ b/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(value).getResultNumber(); + return op.getDistinctOperands()[argNumber]; +} + +static std::optional 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(); + if (lhsOp && getDistinctObjectsOperand(lhsOp, lhs) == rhs) + return AliasResult::MustAlias; + + auto rhsOp = rhs.getDefiningOp(); + 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 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)); diff --git a/mlir/test/Analysis/test-alias-analysis.mlir b/mlir/test/Analysis/test-alias-analysis.mlir index 8cbee61c78b45..d71adee05c7a3 100644 --- a/mlir/test/Analysis/test-alias-analysis.mlir +++ b/mlir/test/Analysis/test-alias-analysis.mlir @@ -256,3 +256,19 @@ func.func @constants(%arg: memref<2xf32>) attributes {test.ptr = "func"} { return } + +// ----- + +// CHECK-LABEL: Testing : "distinct_objects" +// CHECK-DAG: func.region0#0 <-> func.region0#1: MayAlias + +// CHECK-DAG: distinct#0 <-> distinct#1: NoAlias +// CHECK-DAG: distinct#0 <-> func.region0#0: MustAlias +// CHECK-DAG: distinct#1 <-> func.region0#0: MayAlias +// CHECK-DAG: distinct#0 <-> func.region0#1: MayAlias +// CHECK-DAG: distinct#1 <-> func.region0#1: MustAlias + +func.func @distinct_objects(%arg: memref, %arg1: memref) attributes {test.ptr = "func"} { + %0, %1 = memref.distinct_objects %arg, %arg1 {test.ptr = "distinct"} : memref, memref + return +} From f492b57daf45a9ddbd9cb12366a3e3840a8e29d1 Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Wed, 24 Sep 2025 13:53:58 +0200 Subject: [PATCH 5/6] infer alias scopes --- .../LLVMIR/Transforms/InferAliasScopeAttrs.h | 22 ++++ .../mlir/Dialect/LLVMIR/Transforms/Passes.td | 16 +++ .../Dialect/LLVMIR/Transforms/CMakeLists.txt | 1 + .../Transforms/InferAliasScopeAttrs.cpp | 112 ++++++++++++++++++ .../Dialect/LLVMIR/infer-alias-attrs.mlir | 14 +++ 5 files changed, 165 insertions(+) create mode 100644 mlir/include/mlir/Dialect/LLVMIR/Transforms/InferAliasScopeAttrs.h create mode 100644 mlir/lib/Dialect/LLVMIR/Transforms/InferAliasScopeAttrs.cpp create mode 100644 mlir/test/Dialect/LLVMIR/infer-alias-attrs.mlir diff --git a/mlir/include/mlir/Dialect/LLVMIR/Transforms/InferAliasScopeAttrs.h b/mlir/include/mlir/Dialect/LLVMIR/Transforms/InferAliasScopeAttrs.h new file mode 100644 index 0000000000000..e0ec027d95b4d --- /dev/null +++ b/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 + +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 diff --git a/mlir/include/mlir/Dialect/LLVMIR/Transforms/Passes.td b/mlir/include/mlir/Dialect/LLVMIR/Transforms/Passes.td index 961909d5c8d27..62cdca9062c2a 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/Transforms/Passes.td +++ b/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 diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/CMakeLists.txt b/mlir/lib/Dialect/LLVMIR/Transforms/CMakeLists.txt index d4ff0955c5d0e..c263f4292c338 100644 --- a/mlir/lib/Dialect/LLVMIR/Transforms/CMakeLists.txt +++ b/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 diff --git a/mlir/lib/Dialect/LLVMIR/Transforms/InferAliasScopeAttrs.cpp b/mlir/lib/Dialect/LLVMIR/Transforms/InferAliasScopeAttrs.cpp new file mode 100644 index 0000000000000..06bf380eea386 --- /dev/null +++ b/mlir/lib/Dialect/LLVMIR/Transforms/InferAliasScopeAttrs.cpp @@ -0,0 +1,112 @@ +//===- InferAliasScopeAttrs.cpp - Infer LLVM alias scope attributes -------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/LLVMIR/Transforms/InferAliasScopeAttrs.h" + +#include "mlir/Analysis/AliasAnalysis.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace LLVM { +#define GEN_PASS_DEF_LLVMINFERALIASSCOPEATTRIBUTES +#include "mlir/Dialect/LLVMIR/Transforms/Passes.h.inc" +} // namespace LLVM +} // namespace mlir + +using namespace mlir; + +static Value getBasePtr(Operation *op) { + // TODO: we need a common interface to get the base ptr. + if (auto loadOp = dyn_cast(op)) + return loadOp.getMemRef(); + + if (auto storeOp = dyn_cast(op)) + return storeOp.getMemRef(); + + return nullptr; +} + +namespace { + +struct LLVMInferAliasScopeAttrs + : public LLVM::impl::LLVMInferAliasScopeAttributesBase< + LLVMInferAliasScopeAttrs> { + void runOnOperation() override { + SmallVector memOps; + getOperation().walk([&](MemoryEffectOpInterface op) { + if ((op.hasEffect() || + op.hasEffect()) && + getBasePtr(op)) + memOps.emplace_back(op); + }); + + if (memOps.empty()) + return markAllAnalysesPreserved(); + + auto &aliasAnalysis = getAnalysis(); + MLIRContext *ctx = &getContext(); + + LLVM::AliasScopeDomainAttr domain; + llvm::SmallDenseMap aliasScopes; + auto getScope = [&](Operation *op) -> LLVM::AliasScopeAttr { + if (!domain) + domain = LLVM::AliasScopeDomainAttr::get(ctx); + + auto scope = + cast_if_present(aliasScopes.lookup(op)); + if (scope) + return scope; + + scope = LLVM::AliasScopeAttr::get(domain); + aliasScopes[op] = scope; + return scope; + }; + + DenseMap> noaliasScopes; + + // TODO: This is quadratic in the number of memOps, can we do better? + for (Operation *op : memOps) { + for (Operation *otherOp : memOps) { + if (op == otherOp) + continue; + + Value basePtr = getBasePtr(op); + assert(basePtr && "Expected base ptr"); + Value otherBasePtr = getBasePtr(otherOp); + assert(otherBasePtr && "Expected base ptr"); + if (!aliasAnalysis.alias(basePtr, otherBasePtr).isNo()) + continue; + + noaliasScopes[op].insert(getScope(otherOp)); + } + } + + if (noaliasScopes.empty()) + return markAllAnalysesPreserved(); + + auto aliasScopesName = + StringAttr::get(ctx, LLVM::LLVMDialect::getAliasScopesAttrName()); + auto noaliasName = + StringAttr::get(ctx, LLVM::LLVMDialect::getNoAliasAttrName()); + + // We are intentionally using discardable attributes here because those are + // generally not robust against codegen transformations (e.g. inlining) and + // this pass is intended to be run just before *-to-llvm conversion. + for (Operation *op : memOps) { + if (auto aliasScope = aliasScopes.lookup(op)) + op->setAttr(aliasScopesName, ArrayAttr::get(ctx, {aliasScope})); + + auto it = noaliasScopes.find(op); + if (it != noaliasScopes.end()) + op->setAttr(noaliasName, ArrayAttr::get(ctx, it->second.getArrayRef())); + } + } +}; +} // namespace diff --git a/mlir/test/Dialect/LLVMIR/infer-alias-attrs.mlir b/mlir/test/Dialect/LLVMIR/infer-alias-attrs.mlir new file mode 100644 index 0000000000000..91a83a88989bf --- /dev/null +++ b/mlir/test/Dialect/LLVMIR/infer-alias-attrs.mlir @@ -0,0 +1,14 @@ +// RUN: mlir-opt -pass-pipeline='builtin.module(func.func(llvm-infer-alias-scopes-attrs))' %s | FileCheck %s + +// CHECK-LABEL: distinct_objects +func.func @distinct_objects(%arg0: memref, %arg1: memref, %arg2: memref) { + %c0 = arith.constant 0 : index + %1, %2, %3 = memref.distinct_objects %arg0, %arg1, %arg2 : memref, memref, memref +// CHECK: memref.load {{.*}} {alias_scopes = [#[[SCOPE0:.*]]], llvm.noalias = [#[[SCOPE1:.*]], #[[SCOPE2:.*]]]} +// CHECK: memref.store {{.*}} {alias_scopes = [#[[SCOPE1]]], llvm.noalias = [#[[SCOPE0]], #[[SCOPE2]]]} : memref +// CHECK: memref.store {{.*}} {alias_scopes = [#[[SCOPE2]]], llvm.noalias = [#[[SCOPE0]], #[[SCOPE1]]]} : memref + %4 = memref.load %1[%c0] : memref + memref.store %4, %2[%c0] : memref + memref.store %4, %3[%c0] : memref + return +} From aceff3733a3a851be37d92d328857cb24517b321 Mon Sep 17 00:00:00 2001 From: Ivan Butygin Date: Wed, 24 Sep 2025 14:11:51 +0200 Subject: [PATCH 6/6] copy attrs --- .../Conversion/MemRefToLLVM/MemRefToLLVM.cpp | 34 ++++++++++++++++--- .../MemRefToLLVM/memref-to-llvm.mlir | 16 +++++++++ 2 files changed, 46 insertions(+), 4 deletions(-) diff --git a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp index 64270726f4a01..a4ac07ace4172 100644 --- a/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp +++ b/mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp @@ -922,6 +922,18 @@ struct GetGlobalMemrefOpLowering } }; +static SmallVector> +copyImportantAttrs(Operation *op) { + SmallVector> 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 { @@ -932,15 +944,22 @@ struct LoadOpLowering : public LoadStoreOpLowering { ConversionPatternRewriter &rewriter) const override { auto type = loadOp.getMemRefType(); + SmallVector> 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( + auto newOp = rewriter.replaceOpWithNewOp( 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(); } }; @@ -955,15 +974,22 @@ struct StoreOpLowering : public LoadStoreOpLowering { ConversionPatternRewriter &rewriter) const override { auto type = op.getMemRefType(); + SmallVector> 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(op, adaptor.getValue(), dataPtr, - op.getAlignment().value_or(0), - false, op.getNontemporal()); + auto newOp = rewriter.replaceOpWithNewOp( + op, adaptor.getValue(), dataPtr, op.getAlignment().value_or(0), false, + op.getNontemporal()); + + for (auto [nameAttr, attr] : importantAttrs) + newOp->setAttr(nameAttr, attr); + return success(); } }; diff --git a/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir b/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir index 0cbe064572911..102943ee7ffb1 100644 --- a/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir +++ b/mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir @@ -820,6 +820,22 @@ func.func @store_with_alignment(%arg0 : memref<32xf32>, %arg1 : f32, %arg2 : ind // ----- +#alias_scope_domain = #llvm.alias_scope_domain> +#alias_scope0 = #llvm.alias_scope, domain = #alias_scope_domain> +#alias_scope1 = #llvm.alias_scope, domain = #alias_scope_domain> + +// ALL-LABEL: func @load_store_alias_attrs +func.func @load_store_alias_attrs(%arg0: memref, %arg1: memref) { +// ALL: llvm.load {{.*}} {alias_scopes = [#[[SCOPE0:.*]]], llvm.noalias = [#[[SCOPE1:.*]]]} +// ALL: llvm.store {{.*}} {alias_scopes = [#[[SCOPE1]]], llvm.noalias = [#[[SCOPE0]]]} : f32, !llvm.ptr + %c0 = arith.constant 0 : index + %0 = memref.load %arg0[%c0] {alias_scopes = [#alias_scope0], llvm.noalias = [#alias_scope1]} : memref + memref.store %0, %arg1[%c0] {alias_scopes = [#alias_scope1], llvm.noalias = [#alias_scope0]} : memref + func.return +} + +// ----- + // Ensure unconvertable memory space not cause a crash // CHECK-LABEL: @alloca_unconvertable_memory_space