[MLIR][OpenMP] Lowering nontemporal clause to LLVM IR for SIMD directive #118751
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,84 @@ | ||
| //===- LowerNontemporal.cpp -------------------------------------------===// | ||
| // | ||
| // 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 | ||
| // | ||
| //===----------------------------------------------------------------------===// | ||
| // | ||
| // Add nontemporal attributes to load and stores of variables marked as | ||
| // nontemporal. | ||
| // | ||
| //===----------------------------------------------------------------------===// | ||
|
|
||
| #include "flang/Optimizer/Dialect/FIRCG/CGOps.h" | ||
| #include "flang/Optimizer/Dialect/FIROpsSupport.h" | ||
| #include "flang/Optimizer/OpenMP/Passes.h" | ||
| #include "mlir/Dialect/OpenMP/OpenMPDialect.h" | ||
| #include "llvm/ADT/TypeSwitch.h" | ||
|
|
||
| using namespace mlir; | ||
|
|
||
| namespace flangomp { | ||
| #define GEN_PASS_DEF_LOWERNONTEMPORALPASS | ||
| #include "flang/Optimizer/OpenMP/Passes.h.inc" | ||
| } // namespace flangomp | ||
|
|
||
| namespace { | ||
| class LowerNontemporalPass | ||
| : public flangomp::impl::LowerNontemporalPassBase<LowerNontemporalPass> { | ||
| void addNonTemporalAttr(omp::SimdOp simdOp) { | ||
| if (simdOp.getNontemporalVars().empty()) | ||
| return; | ||
|
|
||
| std::function<mlir::Value(mlir::Value)> getBaseOperand = | ||
| [&](mlir::Value operand) -> mlir::Value { | ||
| auto *defOp = operand.getDefiningOp(); | ||
| while (defOp) { | ||
| llvm::TypeSwitch<Operation *>(defOp) | ||
| .Case<fir::ArrayCoorOp, fir::cg::XArrayCoorOp, fir::LoadOp>( | ||
| [&](auto op) { | ||
| operand = op.getMemref(); | ||
| defOp = operand.getDefiningOp(); | ||
| }) | ||
| .Case<fir::BoxAddrOp>([&](auto op) { | ||
| operand = op.getVal(); | ||
| defOp = operand.getDefiningOp(); | ||
| }) | ||
| .Default([&](auto op) { defOp = nullptr; }); | ||
| } | ||
| return operand; | ||
| }; | ||
|
|
||
| // walk through the operations and mark the load and store as nontemporal | ||
| simdOp->walk([&](Operation *op) { | ||
| mlir::Value operand = nullptr; | ||
|
|
||
| if (auto loadOp = llvm::dyn_cast<fir::LoadOp>(op)) | ||
| operand = loadOp.getMemref(); | ||
| else if (auto storeOp = llvm::dyn_cast<fir::StoreOp>(op)) | ||
| operand = storeOp.getMemref(); | ||
|
|
||
| // Skip load and store operations involving boxes (allocatable or pointer | ||
| // types). | ||
| if (operand && !(fir::isAllocatableType(operand.getType()) || | ||
| fir::isPointerType((operand.getType())))) { | ||
|
Comment on lines
+64
to
+65
There was a problem hiding this comment. If the Allocatable and Pointer checks are for skipping the load/store of the boxes then a comment will be useful. There was a problem hiding this comment. Thanks for the comments. |
||
| operand = getBaseOperand(operand); | ||
|
|
||
| // TODO : Handling of nontemporal clause inside atomic construct | ||
| if (llvm::is_contained(simdOp.getNontemporalVars(), operand)) { | ||
| if (auto loadOp = llvm::dyn_cast<fir::LoadOp>(op)) | ||
| loadOp.setNontemporal(true); | ||
| else if (auto storeOp = llvm::dyn_cast<fir::StoreOp>(op)) | ||
| storeOp.setNontemporal(true); | ||
tblah marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| } | ||
| } | ||
| }); | ||
| } | ||
|
|
||
| void runOnOperation() override { | ||
| Operation *op = getOperation(); | ||
| op->walk([&](omp::SimdOp simdOp) { addNonTemporalAttr(simdOp); }); | ||
| } | ||
| }; | ||
| } // namespace | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,111 @@ | ||
| // Test lower-nontemporal pass | ||
| // RUN: fir-opt --fir-to-llvm-ir %s | FileCheck %s --check-prefixes=CHECK-LABEL,CHECK | ||
|
|
||
| // CHECK-LABEL: llvm.func @_QPtest() | ||
| // CHECK: %[[CONST_VAL:.*]] = llvm.mlir.constant(1 : i64) : i64 | ||
| // CHECK: %[[VAL1:.*]] = llvm.alloca %[[CONST_VAL]] x i32 {bindc_name = "n"} : (i64) -> !llvm.ptr | ||
| // CHECK: %[[CONST_VAL1:.*]] = llvm.mlir.constant(1 : i64) : i64 | ||
| // CHECK: %[[VAL2:.*]] = llvm.alloca %[[CONST_VAL1]] x i32 {bindc_name = "i"} : (i64) -> !llvm.ptr | ||
| // CHECK: %[[CONST_VAL2:.*]] = llvm.mlir.constant(1 : i64) : i64 | ||
| // CHECK: %[[VAL3:.*]] = llvm.alloca %[[CONST_VAL2]] x i32 {bindc_name = "c"} : (i64) -> !llvm.ptr | ||
| // CHECK: %[[CONST_VAL3:.*]] = llvm.mlir.constant(1 : i64) : i64 | ||
| // CHECK: %[[VAL4:.*]] = llvm.alloca %[[CONST_VAL3]] x i32 {bindc_name = "b"} : (i64) -> !llvm.ptr | ||
| // CHECK: %[[CONST_VAL4:.*]] = llvm.mlir.constant(1 : i64) : i64 | ||
| // CHECK: %[[VAL5:.*]] = llvm.alloca %[[CONST_VAL4]] x i32 {bindc_name = "a"} : (i64) -> !llvm.ptr | ||
| // CHECK: %[[CONST_VAL5:.*]] = llvm.mlir.constant(1 : i32) : i32 | ||
| // CHECK: %[[VAL6:.*]] = llvm.load %[[VAL1]] : !llvm.ptr -> i32 | ||
| // CHECK: omp.simd nontemporal(%[[VAL5]], %[[VAL3]] : !llvm.ptr, !llvm.ptr) private(@_QFtestEi_private_i32 %[[VAL2]] -> %arg0 : !llvm.ptr) { | ||
| // CHECK: omp.loop_nest (%{{.*}}) : i32 = (%[[CONST_VAL5]]) to (%[[VAL6]]) inclusive step (%[[CONST_VAL5]]) { | ||
| // CHECK: llvm.store %{{.*}}, %{{.*}} : i32, !llvm.ptr | ||
| // CHECK: %[[VAL8:.*]] = llvm.load %[[VAL5]] {nontemporal} : !llvm.ptr -> i32 | ||
| // CHECK: %[[VAL9:.*]] = llvm.load %[[VAL4]] : !llvm.ptr -> i32 | ||
| // CHECK: %[[VAL10:.*]] = llvm.add %[[VAL8]], %[[VAL9]] : i32 | ||
| // CHECK: llvm.store %[[VAL10]], %[[VAL3]] {nontemporal} : i32, !llvm.ptr | ||
| // CHECK: omp.yield | ||
| // CHECK: } | ||
| // CHECK: } | ||
|
|
||
| func.func @_QPtest() { | ||
| %c1_i32 = arith.constant 1 : i32 | ||
| %0 = fir.alloca i32 {bindc_name = "a", uniq_name = "_QFtestEa"} | ||
| %1 = fir.alloca i32 {bindc_name = "b", uniq_name = "_QFtestEb"} | ||
| %2 = fir.alloca i32 {bindc_name = "c", uniq_name = "_QFtestEc"} | ||
| %3 = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFtestEi"} | ||
| %4 = fir.alloca i32 {bindc_name = "n", uniq_name = "_QFtestEn"} | ||
| %5 = fir.load %4 : !fir.ref<i32> | ||
| omp.simd nontemporal(%0, %2 : !fir.ref<i32>, !fir.ref<i32>) private(@_QFtestEi_private_i32 %3 -> %arg0 : !fir.ref<i32>) { | ||
| omp.loop_nest (%arg1) : i32 = (%c1_i32) to (%5) inclusive step (%c1_i32) { | ||
| fir.store %arg1 to %arg0 : !fir.ref<i32> | ||
| %6 = fir.load %0 {nontemporal}: !fir.ref<i32> | ||
| %7 = fir.load %1 : !fir.ref<i32> | ||
| %8 = arith.addi %6, %7 : i32 | ||
| fir.store %8 to %2 {nontemporal} : !fir.ref<i32> | ||
| omp.yield | ||
| } | ||
| } | ||
| return | ||
| } | ||
|
|
||
| // CHECK-LABEL: llvm.func @_QPsimd_nontemporal_allocatable | ||
| // CHECK: %[[CONST_VAL:.*]] = llvm.mlir.constant(1 : i64) : i64 | ||
| // CHECK: %[[ALLOCA2:.*]] = llvm.alloca %[[CONST_VAL]] x i32 {bindc_name = "i"} : (i64) -> !llvm.ptr | ||
| // CHECK: %[[IDX_VAL:.*]] = llvm.mlir.constant(1 : i32) : i32 | ||
| // CHECK: %[[CONST_VAL1:.*]] = llvm.mlir.constant(0 : index) : i64 | ||
| // CHECK: %[[END_IDX:.*]] = llvm.mlir.constant(100 : i32) : i32 | ||
| // CHECK: omp.simd nontemporal(%[[ARG0:.*]] : !llvm.ptr) private(@_QFsimd_nontemporal_allocatableEi_private_i32 %[[ALLOCA2]] -> %[[ARG2:.*]] : !llvm.ptr) { | ||
| // CHECK: omp.loop_nest (%[[ARG3:.*]]) : i32 = (%[[IDX_VAL]]) to (%[[END_IDX]]) inclusive step (%[[IDX_VAL]]) { | ||
| // CHECK: llvm.store %[[ARG3]], %[[ARG2]] : i32, !llvm.ptr | ||
| // CHECK: %[[CONST_VAL2:.*]] = llvm.mlir.constant(48 : i32) : i32 | ||
| // CHECK: "llvm.intr.memcpy"(%[[ALLOCA1:.*]], %[[ARG0]], %[[CONST_VAL2]]) <{isVolatile = false}> : (!llvm.ptr, !llvm.ptr, i32) -> () | ||
| // CHECK: %[[VAL1:.*]] = llvm.load %[[ARG2]] : !llvm.ptr -> i32 | ||
| // CHECK: %[[VAL2:.*]] = llvm.sext %[[VAL1]] : i32 to i64 | ||
| // CHECK: %[[VAL3:.*]] = llvm.getelementptr %[[ALLOCA1]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(ptr, i64, i32, i8, i8, i8, i8, array<1 x array<3 x i64>>)> | ||
| // CHECK: %[[VAL4:.*]] = llvm.load %[[VAL3]] : !llvm.ptr -> !llvm.ptr | ||
| // CHECK: %[[VAL5:.*]] = llvm.getelementptr %[[ALLOCA1]][0, 7, %[[CONST_VAL1]], 0] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.struct<(ptr, i64, i32, i8, i8, i8, i8, array<1 x array<3 x i64>>)> | ||
| // CHECK: %[[VAL6:.*]] = llvm.load %[[VAL5]] : !llvm.ptr -> i64 | ||
| // CHECK: %[[VAL7:.*]] = llvm.getelementptr %[[ALLOCA1]][0, 7, %[[CONST_VAL1]], 1] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.struct<(ptr, i64, i32, i8, i8, i8, i8, array<1 x array<3 x i64>>)> | ||
| // CHECK: %[[VAL8:.*]] = llvm.load %[[VAL7]] : !llvm.ptr -> i64 | ||
| // CHECK: %[[VAL10:.*]] = llvm.mlir.constant(1 : i64) : i64 | ||
| // CHECK: %[[VAL11:.*]] = llvm.mlir.constant(0 : i64) : i64 | ||
| // CHECK: %[[VAL12:.*]] = llvm.sub %[[VAL2]], %[[VAL6]] overflow<nsw> : i64 | ||
| // CHECK: %[[VAL13:.*]] = llvm.mul %[[VAL12]], %[[VAL10]] overflow<nsw> : i64 | ||
| // CHECK: %[[VAL14:.*]] = llvm.mul %[[VAL13]], %[[VAL10]] overflow<nsw> : i64 | ||
| // CHECK: %[[VAL15:.*]] = llvm.add %[[VAL14]], %[[VAL11]] overflow<nsw> : i64 | ||
| // CHECK: %[[VAL16:.*]] = llvm.mul %[[VAL10]], %[[VAL8]] overflow<nsw> : i64 | ||
| // CHECK: %[[VAL17:.*]] = llvm.getelementptr %[[VAL4]][%[[VAL15]]] : (!llvm.ptr, i64) -> !llvm.ptr, i32 | ||
| // CHECK: %[[VAL18:.*]] = llvm.load %[[VAL17]] {nontemporal} : !llvm.ptr -> i32 | ||
| // CHECK: %[[VAL19:.*]] = llvm.load %{{.*}} : !llvm.ptr -> i32 | ||
| // CHECK: %[[VAL20:.*]] = llvm.add %[[VAL18]], %[[VAL19]] : i32 | ||
| // CHECK: llvm.store %[[VAL20]], %[[VAL17]] {nontemporal} : i32, !llvm.ptr | ||
| // CHECK: omp.yield | ||
| // CHECK: } | ||
| // CHECK: } | ||
| // CHECK: llvm.return | ||
|
|
||
| func.func @_QPsimd_nontemporal_allocatable(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>> {fir.bindc_name = "x"}, %arg1: !fir.ref<i32> {fir.bindc_name = "y"}) { | ||
| %c100 = arith.constant 100 : index | ||
| %c1_i32 = arith.constant 1 : i32 | ||
| %c0 = arith.constant 0 : index | ||
| %c100_i32 = arith.constant 100 : i32 | ||
| %0 = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFsimd_nontemporal_allocatableEi"} | ||
| %1 = fir.allocmem !fir.array<?xi32>, %c100 {fir.must_be_heap = true, uniq_name = "_QFsimd_nontemporal_allocatableEx.alloc"} | ||
| %2 = fircg.ext_embox %1(%c100) : (!fir.heap<!fir.array<?xi32>>, index) -> !fir.box<!fir.heap<!fir.array<?xi32>>> | ||
| fir.store %2 to %arg0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>> | ||
| omp.simd nontemporal(%arg0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>) private(@_QFsimd_nontemporal_allocatableEi_private_i32 %0 -> %arg2 : !fir.ref<i32>) { | ||
| omp.loop_nest (%arg3) : i32 = (%c1_i32) to (%c100_i32) inclusive step (%c1_i32) { | ||
| fir.store %arg3 to %arg2 : !fir.ref<i32> | ||
| %7 = fir.load %arg0 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>> | ||
| %8 = fir.load %arg2 : !fir.ref<i32> | ||
| %9 = fir.convert %8 : (i32) -> i64 | ||
| %10 = fir.box_addr %7 : (!fir.box<!fir.heap<!fir.array<?xi32>>>) -> !fir.heap<!fir.array<?xi32>> | ||
| %11:3 = fir.box_dims %7, %c0 : (!fir.box<!fir.heap<!fir.array<?xi32>>>, index) -> (index, index, index) | ||
| %12 = fircg.ext_array_coor %10(%11#1) origin %11#0<%9> : (!fir.heap<!fir.array<?xi32>>, index, index, i64) -> !fir.ref<i32> | ||
| %13 = fir.load %12 {nontemporal} : !fir.ref<i32> | ||
| %14 = fir.load %arg1 : !fir.ref<i32> | ||
| %15 = arith.addi %13, %14 : i32 | ||
| fir.store %15 to %12 {nontemporal} : !fir.ref<i32> | ||
| omp.yield | ||
| } | ||
| } | ||
| return | ||
| } |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why not run this on
omp::SimdOp? Running it on functions is wasteful because many functions will not contain a simd op and runs the risk of missing any simd ops which are not nested inside of functions (I can't see how that would happen currently but we have had problems before lowering non-function (but function-like) operations such asomp.declare_reductionbecause subsequent passes incorrectly assumed all operations were nested inside of functions.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I am applying this pass specifically to omp::simdOp, not to all operations within the function. All the nontemporal examples that I referred so far have omp.simd nested inside a function, so I run the pass at the function level, look for omp::simdOp and only add the attribute when such an operation(omp::simdop) is present.
Could you please share an example that explains the scenario "omp.simd" op is not nested inside function?
That would help me to understand better.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It was only a hypothetical example - I can't think of a case where we would do simd not inside of a function either.
Still, changing this to
should be a better fit with the intention of the pass, because it will not modify anything outside of the simd op
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Theoretically, it could be inside one of the declaration operations (Private, Reduction operations).
BTW, why was this change not made?
Uh oh!
There was an error while loading. Please reload this page.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Re: declaration operations: It could be but I thought that would be so unlikely that it would be better to only make that (minimal) change if it was ever needed.
Re: this thread: Kaviya and I were unable to get it working. It turns out that the MLIR pass manager doesn't allow for arbitrary nesting: only 1 level deep. So you can run a func.func pass on a pass manager for a module because the func.func is a direct descendant. But it couldn't be done for simd because that could be anywhere in the operation tree. I should have posted an update here.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks, Tom!
@kiranchandramohan -- Tom has covered all the points here.