diff options
Diffstat (limited to 'mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp')
| -rw-r--r-- | mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp | 172 |
1 files changed, 85 insertions, 87 deletions
diff --git a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp index 642ced9..dcfe2c7 100644 --- a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp +++ b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp @@ -40,6 +40,16 @@ static bool isScalarLikeType(Type type) { return type.isIntOrIndexOrFloat() || isa<ComplexType>(type); } +/// Helper function to attach the `VarName` attribute to an operation +/// if a variable name is provided. +static void attachVarNameAttr(Operation *op, OpBuilder &builder, + StringRef varName) { + if (!varName.empty()) { + auto varNameAttr = acc::VarNameAttr::get(builder.getContext(), varName); + op->setAttr(acc::getVarNameAttrName(), varNameAttr); + } +} + struct MemRefPointerLikeModel : public PointerLikeType::ExternalModel<MemRefPointerLikeModel, MemRefType> { @@ -83,7 +93,9 @@ struct MemRefPointerLikeModel // then we can generate an alloca operation. if (memrefTy.hasStaticShape()) { needsFree = false; // alloca doesn't need deallocation - return memref::AllocaOp::create(builder, loc, memrefTy).getResult(); + auto allocaOp = memref::AllocaOp::create(builder, loc, memrefTy); + attachVarNameAttr(allocaOp, builder, varName); + return allocaOp.getResult(); } // For dynamic memrefs, extract sizes from the original variable if @@ -103,8 +115,10 @@ struct MemRefPointerLikeModel // Static dimensions are handled automatically by AllocOp } needsFree = true; // alloc needs deallocation - return memref::AllocOp::create(builder, loc, memrefTy, dynamicSizes) - .getResult(); + auto allocOp = + memref::AllocOp::create(builder, loc, memrefTy, dynamicSizes); + attachVarNameAttr(allocOp, builder, varName); + return allocOp.getResult(); } // TODO: Unranked not yet supported. @@ -1016,12 +1030,12 @@ struct RemoveConstantIfConditionWithRegion : public OpRewritePattern<OpTy> { //===----------------------------------------------------------------------===// /// Create and populate an init region for privatization recipes. -/// Returns the init block on success, or nullptr on failure. +/// Returns success if the region is populated, failure otherwise. /// Sets needsFree to indicate if the allocated memory requires deallocation. -static std::unique_ptr<Block> createInitRegion(OpBuilder &builder, Location loc, - Type varType, StringRef varName, - ValueRange bounds, - bool &needsFree) { +static LogicalResult createInitRegion(OpBuilder &builder, Location loc, + Region &initRegion, Type varType, + StringRef varName, ValueRange bounds, + bool &needsFree) { // Create init block with arguments: original value + bounds SmallVector<Type> argTypes{varType}; SmallVector<Location> argLocs{loc}; @@ -1030,9 +1044,9 @@ static std::unique_ptr<Block> createInitRegion(OpBuilder &builder, Location loc, argLocs.push_back(loc); } - auto initBlock = std::make_unique<Block>(); + Block *initBlock = builder.createBlock(&initRegion); initBlock->addArguments(argTypes, argLocs); - builder.setInsertionPointToStart(initBlock.get()); + builder.setInsertionPointToStart(initBlock); Value privatizedValue; @@ -1046,7 +1060,7 @@ static std::unique_ptr<Block> createInitRegion(OpBuilder &builder, Location loc, privatizedValue = mappableTy.generatePrivateInit( builder, loc, typedVar, varName, bounds, {}, needsFree); if (!privatizedValue) - return nullptr; + return failure(); } else { assert(isa<PointerLikeType>(varType) && "Expected PointerLikeType"); auto pointerLikeTy = cast<PointerLikeType>(varType); @@ -1054,21 +1068,21 @@ static std::unique_ptr<Block> createInitRegion(OpBuilder &builder, Location loc, privatizedValue = pointerLikeTy.genAllocate(builder, loc, varName, varType, blockArgVar, needsFree); if (!privatizedValue) - return nullptr; + return failure(); } // Add yield operation to init block acc::YieldOp::create(builder, loc, privatizedValue); - return initBlock; + return success(); } /// Create and populate a copy region for firstprivate recipes. -/// Returns the copy block on success, or nullptr on failure. +/// Returns success if the region is populated, failure otherwise. /// TODO: Handle MappableType - it does not yet have a copy API. -static std::unique_ptr<Block> createCopyRegion(OpBuilder &builder, Location loc, - Type varType, - ValueRange bounds) { +static LogicalResult createCopyRegion(OpBuilder &builder, Location loc, + Region ©Region, Type varType, + ValueRange bounds) { // Create copy block with arguments: original value + privatized value + // bounds SmallVector<Type> copyArgTypes{varType, varType}; @@ -1078,16 +1092,16 @@ static std::unique_ptr<Block> createCopyRegion(OpBuilder &builder, Location loc, copyArgLocs.push_back(loc); } - auto copyBlock = std::make_unique<Block>(); + Block *copyBlock = builder.createBlock(©Region); copyBlock->addArguments(copyArgTypes, copyArgLocs); - builder.setInsertionPointToStart(copyBlock.get()); + builder.setInsertionPointToStart(copyBlock); bool isMappable = isa<MappableType>(varType); bool isPointerLike = isa<PointerLikeType>(varType); // TODO: Handle MappableType - it does not yet have a copy API. // Otherwise, for now just fallback to pointer-like behavior. if (isMappable && !isPointerLike) - return nullptr; + return failure(); // Generate copy region body based on variable type if (isPointerLike) { @@ -1099,21 +1113,20 @@ static std::unique_ptr<Block> createCopyRegion(OpBuilder &builder, Location loc, if (!pointerLikeTy.genCopy( builder, loc, cast<TypedValue<PointerLikeType>>(privatizedArg), cast<TypedValue<PointerLikeType>>(originalArg), varType)) - return nullptr; + return failure(); } // Add terminator to copy block acc::TerminatorOp::create(builder, loc); - return copyBlock; + return success(); } /// Create and populate a destroy region for privatization recipes. -/// Returns the destroy block on success, or nullptr if not needed. -static std::unique_ptr<Block> createDestroyRegion(OpBuilder &builder, - Location loc, Type varType, - Value allocRes, - ValueRange bounds) { +/// Returns success if the region is populated, failure otherwise. +static LogicalResult createDestroyRegion(OpBuilder &builder, Location loc, + Region &destroyRegion, Type varType, + Value allocRes, ValueRange bounds) { // Create destroy block with arguments: original value + privatized value + // bounds SmallVector<Type> destroyArgTypes{varType, varType}; @@ -1123,28 +1136,25 @@ static std::unique_ptr<Block> createDestroyRegion(OpBuilder &builder, destroyArgLocs.push_back(loc); } - auto destroyBlock = std::make_unique<Block>(); + Block *destroyBlock = builder.createBlock(&destroyRegion); destroyBlock->addArguments(destroyArgTypes, destroyArgLocs); - builder.setInsertionPointToStart(destroyBlock.get()); + builder.setInsertionPointToStart(destroyBlock); - bool isMappable = isa<MappableType>(varType); - bool isPointerLike = isa<PointerLikeType>(varType); - // TODO: Handle MappableType - it does not yet have a deallocation API. - // Otherwise, for now just fallback to pointer-like behavior. - if (isMappable && !isPointerLike) - return nullptr; - - assert(isa<PointerLikeType>(varType) && "Expected PointerLikeType"); - auto pointerLikeTy = cast<PointerLikeType>(varType); - auto privatizedArg = + auto varToFree = cast<TypedValue<PointerLikeType>>(destroyBlock->getArgument(1)); - // Pass allocRes to help determine the allocation type - if (!pointerLikeTy.genFree(builder, loc, privatizedArg, allocRes, varType)) - return nullptr; + if (isa<MappableType>(varType)) { + auto mappableTy = cast<MappableType>(varType); + if (!mappableTy.generatePrivateDestroy(builder, loc, varToFree)) + return failure(); + } else { + assert(isa<PointerLikeType>(varType) && "Expected PointerLikeType"); + auto pointerLikeTy = cast<PointerLikeType>(varType); + if (!pointerLikeTy.genFree(builder, loc, varToFree, allocRes, varType)) + return failure(); + } acc::TerminatorOp::create(builder, loc); - - return destroyBlock; + return success(); } } // namespace @@ -1206,40 +1216,33 @@ PrivateRecipeOp::createAndPopulate(OpBuilder &builder, Location loc, if (!isMappable && !isPointerLike) return std::nullopt; - // Create init and destroy blocks using shared helpers OpBuilder::InsertionGuard guard(builder); - // Save the original insertion point for creating the recipe operation later - auto originalInsertionPoint = builder.saveInsertionPoint(); + // Create the recipe operation first so regions have proper parent context + auto recipe = PrivateRecipeOp::create(builder, loc, recipeName, varType); + // Populate the init region bool needsFree = false; - auto initBlock = - createInitRegion(builder, loc, varType, varName, bounds, needsFree); - if (!initBlock) + if (failed(createInitRegion(builder, loc, recipe.getInitRegion(), varType, + varName, bounds, needsFree))) { + recipe.erase(); return std::nullopt; + } // Only create destroy region if the allocation needs deallocation - std::unique_ptr<Block> destroyBlock; if (needsFree) { // Extract the allocated value from the init block's yield operation - auto yieldOp = cast<acc::YieldOp>(initBlock->getTerminator()); + auto yieldOp = + cast<acc::YieldOp>(recipe.getInitRegion().front().getTerminator()); Value allocRes = yieldOp.getOperand(0); - destroyBlock = createDestroyRegion(builder, loc, varType, allocRes, bounds); - if (!destroyBlock) + if (failed(createDestroyRegion(builder, loc, recipe.getDestroyRegion(), + varType, allocRes, bounds))) { + recipe.erase(); return std::nullopt; + } } - // Now create the recipe operation at the original insertion point and attach - // the blocks - builder.restoreInsertionPoint(originalInsertionPoint); - auto recipe = PrivateRecipeOp::create(builder, loc, recipeName, varType); - - // Move the blocks into the recipe's regions - recipe.getInitRegion().push_back(initBlock.release()); - if (destroyBlock) - recipe.getDestroyRegion().push_back(destroyBlock.release()); - return recipe; } @@ -1285,45 +1288,40 @@ FirstprivateRecipeOp::createAndPopulate(OpBuilder &builder, Location loc, if (!isMappable && !isPointerLike) return std::nullopt; - // Create init, copy, and destroy blocks using shared helpers OpBuilder::InsertionGuard guard(builder); - // Save the original insertion point for creating the recipe operation later - auto originalInsertionPoint = builder.saveInsertionPoint(); + // Create the recipe operation first so regions have proper parent context + auto recipe = FirstprivateRecipeOp::create(builder, loc, recipeName, varType); + // Populate the init region bool needsFree = false; - auto initBlock = - createInitRegion(builder, loc, varType, varName, bounds, needsFree); - if (!initBlock) + if (failed(createInitRegion(builder, loc, recipe.getInitRegion(), varType, + varName, bounds, needsFree))) { + recipe.erase(); return std::nullopt; + } - auto copyBlock = createCopyRegion(builder, loc, varType, bounds); - if (!copyBlock) + // Populate the copy region + if (failed(createCopyRegion(builder, loc, recipe.getCopyRegion(), varType, + bounds))) { + recipe.erase(); return std::nullopt; + } // Only create destroy region if the allocation needs deallocation - std::unique_ptr<Block> destroyBlock; if (needsFree) { // Extract the allocated value from the init block's yield operation - auto yieldOp = cast<acc::YieldOp>(initBlock->getTerminator()); + auto yieldOp = + cast<acc::YieldOp>(recipe.getInitRegion().front().getTerminator()); Value allocRes = yieldOp.getOperand(0); - destroyBlock = createDestroyRegion(builder, loc, varType, allocRes, bounds); - if (!destroyBlock) + if (failed(createDestroyRegion(builder, loc, recipe.getDestroyRegion(), + varType, allocRes, bounds))) { + recipe.erase(); return std::nullopt; + } } - // Now create the recipe operation at the original insertion point and attach - // the blocks - builder.restoreInsertionPoint(originalInsertionPoint); - auto recipe = FirstprivateRecipeOp::create(builder, loc, recipeName, varType); - - // Move the blocks into the recipe's regions - recipe.getInitRegion().push_back(initBlock.release()); - recipe.getCopyRegion().push_back(copyBlock.release()); - if (destroyBlock) - recipe.getDestroyRegion().push_back(destroyBlock.release()); - return recipe; } |