[mlir][acc] Fix async only api on data entry operations #122818
Merged
Conversation
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
Data entry operations which are created from constructs with async clause that has no value (aka `acc data copyin(var) async`) end up holding an attribute array named to keep track of this information. However, in cases where `async` clause is not used, calling `hasAsyncOnly` ends up crashing since this attribute is not set. Thus, to fix this issue, ensure that we check for this attribute before trying to walk the attribute array.
|
@llvm/pr-subscribers-openacc @llvm/pr-subscribers-mlir Author: Razvan Lupusoru (razvanlupusoru) ChangesData entry operations which are created from constructs with async clause that has no value (aka Thus, to fix this issue, ensure that we check for this attribute before trying to walk the attribute array. Full diff: https://github.com/llvm/llvm-project/pull/122818.diff 2 Files Affected:
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
index a47f70b168066e..c60eb5cc620a7d 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
@@ -445,7 +445,10 @@ class OpenACC_DataEntryOp<string mnemonic, string clause, string extraDescriptio
}
/// Return true if the op has the async attribute for the given device_type.
bool hasAsyncOnly(mlir::acc::DeviceType deviceType) {
- for (auto attr : getAsyncOnlyAttr()) {
+ mlir::ArrayAttr asyncOnly = getAsyncOnlyAttr();
+ if (!asyncOnly)
+ return false;
+ for (auto attr : asyncOnly) {
auto deviceTypeAttr = mlir::dyn_cast<mlir::acc::DeviceTypeAttr>(attr);
if (deviceTypeAttr.getValue() == deviceType)
return true;
@@ -817,7 +820,10 @@ class OpenACC_DataExitOp<string mnemonic, string clause, string extraDescription
}
/// Return true if the op has the async attribute for the given device_type.
bool hasAsyncOnly(mlir::acc::DeviceType deviceType) {
- for (auto attr : getAsyncOnlyAttr()) {
+ mlir::ArrayAttr asyncOnly = getAsyncOnlyAttr();
+ if (!asyncOnly)
+ return false;
+ for (auto attr : asyncOnly) {
auto deviceTypeAttr = mlir::dyn_cast<mlir::acc::DeviceTypeAttr>(attr);
if (deviceTypeAttr.getValue() == deviceType)
return true;
diff --git a/mlir/unittests/Dialect/OpenACC/OpenACCOpsTest.cpp b/mlir/unittests/Dialect/OpenACC/OpenACCOpsTest.cpp
index cfb8aa767b6f86..aa16421cbec512 100644
--- a/mlir/unittests/Dialect/OpenACC/OpenACCOpsTest.cpp
+++ b/mlir/unittests/Dialect/OpenACC/OpenACCOpsTest.cpp
@@ -77,6 +77,54 @@ TEST_F(OpenACCOpsTest, asyncOnlyTest) {
testAsyncOnly<SerialOp>(b, context, loc, dtypes);
}
+template <typename Op>
+void testAsyncOnlyDataEntry(OpBuilder &b, MLIRContext &context, Location loc,
+ llvm::SmallVector<DeviceType> &dtypes) {
+ auto memrefTy = MemRefType::get({}, b.getI32Type());
+ OwningOpRef<memref::AllocaOp> varPtrOp =
+ b.create<memref::AllocaOp>(loc, memrefTy);
+
+ TypedValue<PointerLikeType> varPtr =
+ cast<TypedValue<PointerLikeType>>(varPtrOp->getResult());
+ OwningOpRef<Op> op = b.create<Op>(loc, varPtr,
+ /*structured=*/true, /*implicit=*/true);
+
+ EXPECT_FALSE(op->hasAsyncOnly());
+ for (auto d : dtypes)
+ EXPECT_FALSE(op->hasAsyncOnly(d));
+
+ auto dtypeNone = DeviceTypeAttr::get(&context, DeviceType::None);
+ op->setAsyncOnlyAttr(b.getArrayAttr({dtypeNone}));
+ EXPECT_TRUE(op->hasAsyncOnly());
+ EXPECT_TRUE(op->hasAsyncOnly(DeviceType::None));
+ op->removeAsyncOnlyAttr();
+
+ auto dtypeHost = DeviceTypeAttr::get(&context, DeviceType::Host);
+ op->setAsyncOnlyAttr(b.getArrayAttr({dtypeHost}));
+ EXPECT_TRUE(op->hasAsyncOnly(DeviceType::Host));
+ EXPECT_FALSE(op->hasAsyncOnly());
+ op->removeAsyncOnlyAttr();
+
+ auto dtypeStar = DeviceTypeAttr::get(&context, DeviceType::Star);
+ op->setAsyncOnlyAttr(b.getArrayAttr({dtypeHost, dtypeStar}));
+ EXPECT_TRUE(op->hasAsyncOnly(DeviceType::Star));
+ EXPECT_TRUE(op->hasAsyncOnly(DeviceType::Host));
+ EXPECT_FALSE(op->hasAsyncOnly());
+
+ op->removeAsyncOnlyAttr();
+}
+
+TEST_F(OpenACCOpsTest, asyncOnlyTestDataEntry) {
+ testAsyncOnlyDataEntry<DevicePtrOp>(b, context, loc, dtypes);
+ testAsyncOnlyDataEntry<PresentOp>(b, context, loc, dtypes);
+ testAsyncOnlyDataEntry<CopyinOp>(b, context, loc, dtypes);
+ testAsyncOnlyDataEntry<CreateOp>(b, context, loc, dtypes);
+ testAsyncOnlyDataEntry<NoCreateOp>(b, context, loc, dtypes);
+ testAsyncOnlyDataEntry<AttachOp>(b, context, loc, dtypes);
+ testAsyncOnlyDataEntry<UpdateDeviceOp>(b, context, loc, dtypes);
+ testAsyncOnlyDataEntry<UseDeviceOp>(b, context, loc, dtypes);
+}
+
template <typename Op>
void testAsyncValue(OpBuilder &b, MLIRContext &context, Location loc,
llvm::SmallVector<DeviceType> &dtypes) {
@@ -105,6 +153,46 @@ TEST_F(OpenACCOpsTest, asyncValueTest) {
testAsyncValue<SerialOp>(b, context, loc, dtypes);
}
+template <typename Op>
+void testAsyncValueDataEntry(OpBuilder &b, MLIRContext &context, Location loc,
+ llvm::SmallVector<DeviceType> &dtypes) {
+ auto memrefTy = MemRefType::get({}, b.getI32Type());
+ OwningOpRef<memref::AllocaOp> varPtrOp =
+ b.create<memref::AllocaOp>(loc, memrefTy);
+
+ TypedValue<PointerLikeType> varPtr =
+ cast<TypedValue<PointerLikeType>>(varPtrOp->getResult());
+ OwningOpRef<Op> op = b.create<Op>(loc, varPtr,
+ /*structured=*/true, /*implicit=*/true);
+
+ mlir::Value empty;
+ EXPECT_EQ(op->getAsyncValue(), empty);
+ for (auto d : dtypes)
+ EXPECT_EQ(op->getAsyncValue(d), empty);
+
+ OwningOpRef<arith::ConstantIndexOp> val =
+ b.create<arith::ConstantIndexOp>(loc, 1);
+ auto dtypeNvidia = DeviceTypeAttr::get(&context, DeviceType::Nvidia);
+ op->setAsyncOperandsDeviceTypeAttr(b.getArrayAttr({dtypeNvidia}));
+ op->getAsyncOperandsMutable().assign(val->getResult());
+ EXPECT_EQ(op->getAsyncValue(), empty);
+ EXPECT_EQ(op->getAsyncValue(DeviceType::Nvidia), val->getResult());
+
+ op->getAsyncOperandsMutable().clear();
+ op->removeAsyncOperandsDeviceTypeAttr();
+}
+
+TEST_F(OpenACCOpsTest, asyncValueTestDataEntry) {
+ testAsyncValueDataEntry<DevicePtrOp>(b, context, loc, dtypes);
+ testAsyncValueDataEntry<PresentOp>(b, context, loc, dtypes);
+ testAsyncValueDataEntry<CopyinOp>(b, context, loc, dtypes);
+ testAsyncValueDataEntry<CreateOp>(b, context, loc, dtypes);
+ testAsyncValueDataEntry<NoCreateOp>(b, context, loc, dtypes);
+ testAsyncValueDataEntry<AttachOp>(b, context, loc, dtypes);
+ testAsyncValueDataEntry<UpdateDeviceOp>(b, context, loc, dtypes);
+ testAsyncValueDataEntry<UseDeviceOp>(b, context, loc, dtypes);
+}
+
template <typename Op>
void testNumGangsValues(OpBuilder &b, MLIRContext &context, Location loc,
llvm::SmallVector<DeviceType> &dtypes,
|
clementval
approved these changes
Jan 13, 2025
|
The |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Data entry operations which are created from constructs with async clause that has no value (aka
acc data copyin(var) async) end up holding an attribute array named to keep track of this information. However, in cases whereasyncclause is not used, callinghasAsyncOnlyends up crashing since this attribute is not set.Thus, to fix this issue, ensure that we check for this attribute before trying to walk the attribute array.