This is an archive of the discontinued LLVM Phabricator instance.

Differential D142616

[mlir] Add nontemporal field to memref.load/store and convey to llvm.load/store
ClosedPublic

Authored by guraypp on Jan 26 2023, 5:06 AM.

Details

Reviewers
ftynse
nicolasvasilache
dcaballe
Commits
rG1cb91b421e7d: [mlir] Add nontemporal field to memref.load/store and convey to llvm.load/store
Summary

llvm.load op has nonTemporal field which is missing for memref.load and memref.store. This revision first adds nonTemporal field to memref's load/store op, then it lowers the field to llvm.load/store ops.

Diff Detail

Event Timeline

guraypp created this revision.Jan 26 2023, 5:06 AM
Herald added a reviewer: ftynse. · View Herald TranscriptJan 26 2023, 5:06 AM
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: Moerafaat, zero9178, bzcheeseman and 22 others. · View Herald Transcript
guraypp requested review of this revision.Jan 26 2023, 5:06 AM
Herald added a reviewer: nicolasvasilache. · View Herald TranscriptJan 26 2023, 5:06 AM
Herald added a reviewer: dcaballe. · View Herald Transcript
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: stephenneuendorffer, nicolasvasilache. · View Herald Transcript
guraypp updated this revision to Diff 492401.Jan 26 2023, 5:08 AM

cleanup

Harbormaster completed remote builds in B210091: Diff 492401.Jan 26 2023, 5:28 AM
guraypp updated this revision to Diff 493245.Jan 30 2023, 2:07 AM

rebase

nicolasvasilache added a comment.Jan 30 2023, 2:09 AM

Have you checked that foldings and canonicalizations don't just drop the new attribute?

Some of the folding / canonicalizations tests should be updated to cover this.

Harbormaster completed remote builds in B210727: Diff 493245.Jan 30 2023, 2:58 AM
guraypp updated this revision to Diff 494542.Feb 3 2023, 12:59 AM

Added flag propogation on the folding passes, and added new tests.

Harbormaster completed remote builds in B211662: Diff 494542.Feb 3 2023, 1:00 AM
nicolasvasilache accepted this revision.Feb 3 2023, 2:43 AM
nicolasvasilache added inline comments.
mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir
527

nit: nl

This revision is now accepted and ready to land.Feb 3 2023, 2:43 AM
guraypp updated this revision to Diff 494575.Feb 3 2023, 3:10 AM

rebase, address the comment

Harbormaster completed remote builds in B211684: Diff 494575.Feb 3 2023, 3:10 AM
guraypp updated this revision to Diff 494577.Feb 3 2023, 3:12 AM

rebase

Harbormaster completed remote builds in B211685: Diff 494577.Feb 3 2023, 3:55 AM
Closed by commit rG1cb91b421e7d: [mlir] Add nontemporal field to memref.load/store and convey to llvm.load/store (authored by guraypp). · Explain WhyFeb 3 2023, 5:03 AM
This revision was automatically updated to reflect the committed changes.
guraypp added a commit: rG1cb91b421e7d: [mlir] Add nontemporal field to memref.load/store and convey to llvm.load/store.

Revision Contents

PathSize
mlir/
lib/
Dialect/
MemRef/
Transforms/
6 lines
17 lines
test/
Conversion/
MemRefToLLVM/
11 lines
Dialect/
MemRef/
12 lines
16 lines
22 lines

Diff 494542

mlir/lib/Dialect/MemRef/Transforms/EmulateWideInt.cpp

Show First 20 Lines • Show All 60 Lines • ▼ Show 20 Lines matchAndRewrite(memref::LoadOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override { ConversionPatternRewriter &rewriter) const override {
Type newResTy = getTypeConverter()->convertType(op.getType()); Type newResTy = getTypeConverter()->convertType(op.getType());
if (!newResTy) if (!newResTy)
return rewriter.notifyMatchFailure( return rewriter.notifyMatchFailure(
op->getLoc(), llvm::formatv("failed to convert memref type: {0}", op->getLoc(), llvm::formatv("failed to convert memref type: {0}",
op.getMemRefType())); op.getMemRefType()));
rewriter.replaceOpWithNewOp<memref::LoadOp>( rewriter.replaceOpWithNewOp<memref::LoadOp>(
op, newResTy, adaptor.getMemref(), adaptor.getIndices()); op, newResTy, adaptor.getMemref(), adaptor.getIndices(),
op.getNontemporal());
return success(); return success();
} }
}; };
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// ConvertMemRefStore // ConvertMemRefStore
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
struct ConvertMemRefStore final : OpConversionPattern<memref::StoreOp> { struct ConvertMemRefStore final : OpConversionPattern<memref::StoreOp> {
using OpConversionPattern::OpConversionPattern; using OpConversionPattern::OpConversionPattern;
LogicalResult LogicalResult
matchAndRewrite(memref::StoreOp op, OpAdaptor adaptor, matchAndRewrite(memref::StoreOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override { ConversionPatternRewriter &rewriter) const override {
Type newTy = getTypeConverter()->convertType(op.getMemRefType()); Type newTy = getTypeConverter()->convertType(op.getMemRefType());
if (!newTy) if (!newTy)
return rewriter.notifyMatchFailure( return rewriter.notifyMatchFailure(
op->getLoc(), llvm::formatv("failed to convert memref type: {0}", op->getLoc(), llvm::formatv("failed to convert memref type: {0}",
op.getMemRefType())); op.getMemRefType()));
rewriter.replaceOpWithNewOp<memref::StoreOp>( rewriter.replaceOpWithNewOp<memref::StoreOp>(
op, adaptor.getValue(), adaptor.getMemref(), adaptor.getIndices()); op, adaptor.getValue(), adaptor.getMemref(), adaptor.getIndices(),
op.getNontemporal());
return success(); return success();
} }
}; };
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
// Pass Definition // Pass Definition
//===----------------------------------------------------------------------===// //===----------------------------------------------------------------------===//
▲ Show 20 Lines • Show All 64 Lines • Show Last 20 Lines

mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp

Show First 20 Lines • Show All 378 Lines • ▼ Show 20 Lines if (auto affineLoadOp = dyn_cast<AffineLoadOp>(loadOp.getOperation())) {
indices.assign(expandedIndices.begin(), expandedIndices.end()); indices.assign(expandedIndices.begin(), expandedIndices.end());
} }
SmallVector<Value, 4> sourceIndices; SmallVector<Value, 4> sourceIndices;
if (failed(resolveSourceIndicesSubView(loadOp.getLoc(), rewriter, subViewOp, if (failed(resolveSourceIndicesSubView(loadOp.getLoc(), rewriter, subViewOp,
indices, sourceIndices))) indices, sourceIndices)))
return failure(); return failure();
llvm::TypeSwitch<Operation *, void>(loadOp) llvm::TypeSwitch<Operation *, void>(loadOp)
.Case<AffineLoadOp, memref::LoadOp>([&](auto op) { .Case([&](AffineLoadOp op) {
rewriter.replaceOpWithNewOp<decltype(op)>(loadOp, subViewOp.getSource(), rewriter.replaceOpWithNewOp<AffineLoadOp>(loadOp, subViewOp.getSource(),
sourceIndices); sourceIndices);
}) })
.Case([&](memref::LoadOp op) {
rewriter.replaceOpWithNewOp<memref::LoadOp>(
loadOp, subViewOp.getSource(), sourceIndices, op.getNontemporal());
})
.Case([&](vector::TransferReadOp transferReadOp) { .Case([&](vector::TransferReadOp transferReadOp) {
rewriter.replaceOpWithNewOp<vector::TransferReadOp>( rewriter.replaceOpWithNewOp<vector::TransferReadOp>(
transferReadOp, transferReadOp.getVectorType(), transferReadOp, transferReadOp.getVectorType(),
subViewOp.getSource(), sourceIndices, subViewOp.getSource(), sourceIndices,
getPermutationMapAttr(rewriter.getContext(), subViewOp, getPermutationMapAttr(rewriter.getContext(), subViewOp,
transferReadOp.getPermutationMap()), transferReadOp.getPermutationMap()),
transferReadOp.getPadding(), transferReadOp.getPadding(),
/*mask=*/Value(), transferReadOp.getInBoundsAttr()); /*mask=*/Value(), transferReadOp.getInBoundsAttr());
▲ Show 20 Lines • Show All 86 Lines • ▼ Show 20 Lines if (auto affineStoreOp = dyn_cast<AffineStoreOp>(storeOp.getOperation())) {
indices.assign(expandedIndices.begin(), expandedIndices.end()); indices.assign(expandedIndices.begin(), expandedIndices.end());
} }
SmallVector<Value, 4> sourceIndices; SmallVector<Value, 4> sourceIndices;
if (failed(resolveSourceIndicesSubView(storeOp.getLoc(), rewriter, subViewOp, if (failed(resolveSourceIndicesSubView(storeOp.getLoc(), rewriter, subViewOp,
indices, sourceIndices))) indices, sourceIndices)))
return failure(); return failure();
llvm::TypeSwitch<Operation *, void>(storeOp) llvm::TypeSwitch<Operation *, void>(storeOp)
.Case<AffineStoreOp, memref::StoreOp>([&](auto op) { .Case([&](AffineStoreOp op) {
rewriter.replaceOpWithNewOp<decltype(op)>( rewriter.replaceOpWithNewOp<AffineStoreOp>(
storeOp, storeOp.getValue(), subViewOp.getSource(), sourceIndices); storeOp, storeOp.getValue(), subViewOp.getSource(), sourceIndices);
}) })
.Case([&](memref::StoreOp op) {
rewriter.replaceOpWithNewOp<memref::StoreOp>(
storeOp, storeOp.getValue(), subViewOp.getSource(), sourceIndices,
op.getNontemporal());
})
.Case([&](vector::TransferWriteOp op) { .Case([&](vector::TransferWriteOp op) {
rewriter.replaceOpWithNewOp<vector::TransferWriteOp>( rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
op, op.getValue(), subViewOp.getSource(), sourceIndices, op, op.getValue(), subViewOp.getSource(), sourceIndices,
getPermutationMapAttr(rewriter.getContext(), subViewOp, getPermutationMapAttr(rewriter.getContext(), subViewOp,
op.getPermutationMap()), op.getPermutationMap()),
op.getInBoundsAttr()); op.getInBoundsAttr());
}) })
.Default([](Operation *) { llvm_unreachable("unexpected operation."); }); .Default([](Operation *) { llvm_unreachable("unexpected operation."); });
▲ Show 20 Lines • Show All 109 Lines • Show Last 20 Lines

mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir

Show First 20 Lines • Show All 535 Lines • ▼ Show 20 Lines -> memref<f32>, index,
index, index index, index
return return
} }
// ----- // -----
// CHECK-LABEL: func @load_non_temporal( // CHECK-LABEL: func @load_non_temporal(
func.func @load_non_temporal(%arg0 : memref<32xf32, affine_map<(d0) -> (d0)>>) { func.func @load_non_temporal(%arg0 : memref<32xf32, affine_map<(d0) -> (d0)>>) {
%0 = memref.alloc() : memref<32xf32, affine_map<(d0) -> (d0)>>
%1 = arith.constant 7 : index %1 = arith.constant 7 : index
// CHECK: llvm.load %{{.*}} {nontemporal} : !llvm.ptr<f32> // CHECK: llvm.load %{{.*}} {nontemporal} : !llvm.ptr<f32>
%2 = memref.load %0[%1] {nontemporal = true} : memref<32xf32, affine_map<(d0) -> (d0)>> %2 = memref.load %arg0[%1] {nontemporal = true} : memref<32xf32, affine_map<(d0) -> (d0)>>
func.return func.return
} }
// ----- // -----
// CHECK-LABEL: func @store_non_temporal( // CHECK-LABEL: func @store_non_temporal(
func.func @store_non_temporal(%arg0 : memref<32xf32, affine_map<(d0) -> (d0)>>, %data : f32) { func.func @store_non_temporal(%input : memref<32xf32, affine_map<(d0) -> (d0)>>, %output : memref<32xf32, affine_map<(d0) -> (d0)>>) {
%0 = memref.alloc() : memref<32xf32, affine_map<(d0) -> (d0)>>
%1 = arith.constant 7 : index %1 = arith.constant 7 : index
%2 = memref.load %input[%1] {nontemporal = true} : memref<32xf32, affine_map<(d0) -> (d0)>>
// CHECK: llvm.store %{{.*}}, %{{.*}} {nontemporal} : !llvm.ptr<f32> // CHECK: llvm.store %{{.*}}, %{{.*}} {nontemporal} : !llvm.ptr<f32>
memref.store %data, %0[%1] {nontemporal = true} : memref<32xf32, affine_map<(d0) -> (d0)>> memref.store %2, %output[%1] {nontemporal = true} : memref<32xf32, affine_map<(d0) -> (d0)>>
func.return func.return
} }

mlir/test/Dialect/MemRef/canonicalize.mlir

Show First 20 Lines • Show All 888 Lines • ▼ Show 20 Linesfunc.func @fold_trivial_subviews(%m: memref<?xf32, strided<[?], offset: ?>>,
%0 = memref.subview %m[5] [%sz] [1] %0 = memref.subview %m[5] [%sz] [1]
: memref<?xf32, strided<[?], offset: ?>> : memref<?xf32, strided<[?], offset: ?>>
to memref<?xf32, strided<[?], offset: ?>> to memref<?xf32, strided<[?], offset: ?>>
%1 = memref.subview %0[0] [%sz] [1] %1 = memref.subview %0[0] [%sz] [1]
: memref<?xf32, strided<[?], offset: ?>> : memref<?xf32, strided<[?], offset: ?>>
to memref<?xf32, strided<[?], offset: ?>> to memref<?xf32, strided<[?], offset: ?>>
return %1 : memref<?xf32, strided<[?], offset: ?>> return %1 : memref<?xf32, strided<[?], offset: ?>>
} }
// -----
// CHECK-LABEL: func @load_store_nontemporal(
func.func @load_store_nontemporal(%input : memref<32xf32, affine_map<(d0) -> (d0)>>, %output : memref<32xf32, affine_map<(d0) -> (d0)>>) {
%1 = arith.constant 7 : index
// CHECK: memref.load %{{.*}}[%{{.*}}] {nontemporal = true} : memref<32xf32>
%2 = memref.load %input[%1] {nontemporal = true} : memref<32xf32, affine_map<(d0) -> (d0)>>
// CHECK: memref.store %{{.*}}, %{{.*}}[%{{.*}}] {nontemporal = true} : memref<32xf32>
memref.store %2, %output[%1] {nontemporal = true} : memref<32xf32, affine_map<(d0) -> (d0)>>
func.return
}

mlir/test/Dialect/MemRef/emulate-wide-int.mlir

Show All 38 Lines
func.func @alloc_load_store_i64() { func.func @alloc_load_store_i64() {
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%c1 = arith.constant 1 : i64 %c1 = arith.constant 1 : i64
%m = memref.alloc() : memref<4xi64, 1> %m = memref.alloc() : memref<4xi64, 1>
%v = memref.load %m[%c0] : memref<4xi64, 1> %v = memref.load %m[%c0] : memref<4xi64, 1>
memref.store %c1, %m[%c0] : memref<4xi64, 1> memref.store %c1, %m[%c0] : memref<4xi64, 1>
return return
} }
// CHECK-LABEL: func @alloc_load_store_i64_nontemporal
// CHECK: [[C1:%.+]] = arith.constant dense<[1, 0]> : vector<2xi32>
// CHECK-NEXT: [[M:%.+]] = memref.alloc() : memref<4xvector<2xi32>, 1>
// CHECK-NEXT: [[V:%.+]] = memref.load [[M]][{{%.+}}] {nontemporal = true} : memref<4xvector<2xi32>, 1>
// CHECK-NEXT: memref.store [[C1]], [[M]][{{%.+}}] {nontemporal = true} : memref<4xvector<2xi32>, 1>
// CHECK-NEXT: return
func.func @alloc_load_store_i64_nontemporal() {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : i64
%m = memref.alloc() : memref<4xi64, 1>
%v = memref.load %m[%c0] {nontemporal = true} : memref<4xi64, 1>
memref.store %c1, %m[%c0] {nontemporal = true} : memref<4xi64, 1>
return
}

mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir

Show First 20 Lines • Show All 496 Lines • ▼ Show 20 Linesfunc.func @subview_of_subview_rank_reducing(%m: memref<?x?x?xf32>,
%0 = memref.subview %m[3, 1, 8] [1, %sz, 1] [1, 1, 1] %0 = memref.subview %m[3, 1, 8] [1, %sz, 1] [1, 1, 1]
: memref<?x?x?xf32> : memref<?x?x?xf32>
to memref<?xf32, strided<[1], offset: ?>> to memref<?xf32, strided<[1], offset: ?>>
%1 = memref.subview %0[6] [1] [1] %1 = memref.subview %0[6] [1] [1]
: memref<?xf32, strided<[1], offset: ?>> : memref<?xf32, strided<[1], offset: ?>>
to memref<f32, strided<[], offset: ?>> to memref<f32, strided<[], offset: ?>>
return %1 : memref<f32, strided<[], offset: ?>> return %1 : memref<f32, strided<[], offset: ?>>
} }
// -----
// CHECK-LABEL: func @fold_load_keep_nontemporal(
// CHECK: memref.load %{{.+}}[%{{.+}}, %{{.+}}] {nontemporal = true}
func.func @fold_load_keep_nontemporal(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index) -> f32 {
%0 = memref.subview %arg0[%arg1, %arg2][4, 4][2, 3] : memref<12x32xf32> to memref<4x4xf32, strided<[64, 3], offset: ?>>
%1 = memref.load %0[%arg3, %arg4] {nontemporal = true }: memref<4x4xf32, strided<[64, 3], offset: ?>>
return %1 : f32
}
// -----
// CHECK-LABEL: func @fold_store_keep_nontemporal(
// CHECK: memref.store %{{.+}}, %{{.+}}[%{{.+}}, %{{.+}}] {nontemporal = true} : memref<12x32xf32>
func.func @fold_store_keep_nontemporal(%arg0 : memref<12x32xf32>, %arg1 : index, %arg2 : index, %arg3 : index, %arg4 : index, %arg5 : f32) {
%0 = memref.subview %arg0[%arg1, %arg2][4, 4][2, 3] :
memref<12x32xf32> to memref<4x4xf32, strided<[64, 3], offset: ?>>
memref.store %arg5, %0[%arg3, %arg4] {nontemporal=true}: memref<4x4xf32, strided<[64, 3], offset: ?>>
return
}
No newline at end of file
nicolasvasilacheUnsubmitted
Not Done
ReplyInline Actions

nit: nl

nicolasvasilache: nit: nl