-
Notifications
You must be signed in to change notification settings - Fork 15.2k
[mlir][acc] Extend PointerLikeType to provide alloc, dealloc, copy #162328
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
A variable in an acc data clause operation must have a type that implements either PointerLikeType or a MappableType interface. These interfaces provide the contract that allows acc dialect and its transform passes to interact with a source dialect. One of these requirements is ability to generate code that creates memory for a private copy and ability to initialize that copy from another variable. Thus, update the PointerLikeType API to provide the means to create allocation, deallocation, and copy. This will be used as a way to fill in privatization and firstprivatization recipes. This new API was implemented for memref along with testing to exercise it via the implementation of a testing pass.
@llvm/pr-subscribers-mlir-openacc @llvm/pr-subscribers-openacc Author: Razvan Lupusoru (razvanlupusoru) ChangesA variable in an acc data clause operation must have a type that implements either PointerLikeType or a MappableType interface. These interfaces provide the contract that allows acc dialect and its transform passes to interact with a source dialect. One of these requirements is ability to generate code that creates memory for a private copy and ability to initialize that copy from another variable. Thus, update the PointerLikeType API to provide the means to create allocation, deallocation, and copy. This will be used as a way to fill in privatization and firstprivatization recipes. This new API was implemented for memref along with testing to exercise it via the implementation of a testing pass. Patch is 27.53 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/162328.diff 11 Files Affected:
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCTypeInterfaces.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCTypeInterfaces.td
index 9123ac34af67d..9a89f8ede191e 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCTypeInterfaces.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCTypeInterfaces.td
@@ -70,6 +70,85 @@ def OpenACC_PointerLikeTypeInterface : TypeInterface<"PointerLikeType"> {
return ::mlir::acc::VariableTypeCategory::uncategorized;
}]
>,
+ InterfaceMethod<
+ /*description=*/[{
+ Generates allocation operations for the pointer-like type. It will create
+ an allocate that produces memory space for an instance of the current type.
+
+ The `varName` parameter is optional and can be used to provide a name
+ for the allocated variable. If the current type is represented
+ in a way that it does not capture the pointee type, `varType` must be
+ passed in to provide the necessary type information.
+
+ The `originalVar` parameter is optional but enables support for dynamic
+ types (e.g., dynamic memrefs). When provided, implementations can extract
+ runtime dimension information from the original variable to create
+ allocations with matching dynamic sizes.
+
+ Returns true if allocation was successfully generated, false otherwise.
+ }],
+ /*retTy=*/"bool",
+ /*methodName=*/"genAllocate",
+ /*args=*/(ins "::mlir::OpBuilder &":$builder,
+ "::mlir::Location":$loc,
+ "::llvm::StringRef":$varName,
+ "::mlir::Type":$varType,
+ "::mlir::Value":$originalVar),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return false;
+ }]
+ >,
+ InterfaceMethod<
+ /*description=*/[{
+ Generates deallocation operations for the pointer-like type. It deallocates
+ the instance provided.
+
+ The `varPtr` parameter is required and must represent an instance that was
+ previously allocated. If the current type is represented in a way that it
+ does not capture the pointee type, `varType` must be passed in to provide
+ the necessary type information. Nothing is generated in case the allocate
+ is `alloca`-like.
+
+ Returns true if deallocation was successfully generated or successfully
+ deemed as not needed to be generated, false otherwise.
+ }],
+ /*retTy=*/"bool",
+ /*methodName=*/"genFree",
+ /*args=*/(ins "::mlir::OpBuilder &":$builder,
+ "::mlir::Location":$loc,
+ "::mlir::TypedValue<::mlir::acc::PointerLikeType>":$varPtr,
+ "::mlir::Type":$varType),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return false;
+ }]
+ >,
+ InterfaceMethod<
+ /*description=*/[{
+ Generates copy operations for the pointer-like type. It copies the memory
+ from the source to the destination. Typically used to initialize one
+ variable of this type from another.
+
+ The `destination` and `source` parameters represent the target and source
+ instances respectively. If the current type is represented in a way that it
+ does not capture the pointee type, `varType` must be passed in to provide
+ the necessary type information.
+
+ Returns true if copy was successfully generated, false otherwise.
+ }],
+ /*retTy=*/"bool",
+ /*methodName=*/"genCopy",
+ /*args=*/(ins "::mlir::OpBuilder &":$builder,
+ "::mlir::Location":$loc,
+ "::mlir::TypedValue<::mlir::acc::PointerLikeType>":$destination,
+ "::mlir::TypedValue<::mlir::acc::PointerLikeType>":$source,
+ "::mlir::Type":$varType),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return false;
+ }]
+ >,
];
}
diff --git a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
index 6598ac141008f..c2a21acb4d62c 100644
--- a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
+++ b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
@@ -7,6 +7,7 @@
// =============================================================================
#include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
@@ -44,6 +45,7 @@ struct MemRefPointerLikeModel
Type getElementType(Type pointer) const {
return cast<MemRefType>(pointer).getElementType();
}
+
mlir::acc::VariableTypeCategory
getPointeeTypeCategory(Type pointer, TypedValue<PointerLikeType> varPtr,
Type varType) const {
@@ -70,6 +72,117 @@ struct MemRefPointerLikeModel
assert(memrefTy.getRank() > 0 && "rank expected to be positive");
return mlir::acc::VariableTypeCategory::array;
}
+
+ bool genAllocate(Type pointer, OpBuilder &builder, Location loc,
+ StringRef varName, Type varType, Value originalVar) const {
+ auto memrefTy = cast<MemRefType>(pointer);
+
+ // Check if this is a static memref (all dimensions are known) - if yes
+ // then we can generate an alloca operation.
+ if (memrefTy.hasStaticShape()) {
+ memref::AllocaOp::create(builder, loc, memrefTy);
+ return true;
+ }
+
+ // For dynamic memrefs, extract sizes from the original variable if provided.
+ // Otherwise they cannot be handled.
+ if (originalVar && originalVar.getType() == memrefTy &&
+ memrefTy.hasRank()) {
+ SmallVector<Value> dynamicSizes;
+ for (int64_t i = 0; i < memrefTy.getRank(); ++i) {
+ if (memrefTy.isDynamicDim(i)) {
+ // Extract the size of dimension i from the original variable
+ auto indexValue = arith::ConstantIndexOp::create(builder, loc, i);
+ auto dimSize =
+ memref::DimOp::create(builder, loc, originalVar, indexValue);
+ dynamicSizes.push_back(dimSize);
+ }
+ // Note: We only add dynamic sizes to the dynamicSizes array
+ // Static dimensions are handled automatically by AllocOp
+ }
+ memref::AllocOp::create(builder, loc, memrefTy, dynamicSizes);
+ return true;
+ }
+
+
+ // TODO: Unranked not yet supported.
+ return false;
+ }
+
+ bool genFree(Type pointer, OpBuilder &builder, Location loc,
+ TypedValue<PointerLikeType> varPtr, Type varType) const {
+ if (auto memrefValue = dyn_cast<TypedValue<MemRefType>>(varPtr)) {
+ // Walk through casts to find the original allocation
+ Value currentValue = memrefValue;
+ Operation *originalAlloc = nullptr;
+
+ // Follow the chain of operations to find the original allocation
+ // even if a casted result is provided.
+ while (currentValue) {
+ if (auto *definingOp = currentValue.getDefiningOp()) {
+ // Check if this is an allocation operation
+ if (isa<memref::AllocOp, memref::AllocaOp>(definingOp)) {
+ originalAlloc = definingOp;
+ break;
+ }
+
+ // Check if this is a cast operation we can look through
+ if (auto castOp = dyn_cast<memref::CastOp>(definingOp)) {
+ currentValue = castOp.getSource();
+ continue;
+ }
+
+ // Check for other cast-like operations
+ if (auto reinterpretCastOp =
+ dyn_cast<memref::ReinterpretCastOp>(definingOp)) {
+ currentValue = reinterpretCastOp.getSource();
+ continue;
+ }
+
+ // If we can't look through this operation, stop
+ break;
+ }
+ // This is a block argument or similar - can't trace further.
+ break;
+ }
+
+ if (originalAlloc) {
+ if (isa<memref::AllocaOp>(originalAlloc)) {
+ // This is an alloca - no dealloc needed, but return true (success)
+ return true;
+ }
+ if (isa<memref::AllocOp>(originalAlloc)) {
+ // This is an alloc - generate dealloc
+ memref::DeallocOp::create(builder, loc, memrefValue);
+ return true;
+ }
+ }
+ }
+
+ return false;
+ }
+
+ bool genCopy(Type pointer, OpBuilder &builder, Location loc,
+ TypedValue<PointerLikeType> destination,
+ TypedValue<PointerLikeType> source, Type varType) const {
+ // Generate a copy operation between two memrefs
+ auto destMemref = dyn_cast_if_present<TypedValue<MemRefType>>(destination);
+ auto srcMemref = dyn_cast_if_present<TypedValue<MemRefType>>(source);
+
+ // As per memref documentation, source and destination must have same
+ // element type and shape in order to be compatible. We do not want to fail
+ // with an IR verification error - thus check that before generating the
+ // copy operation.
+ if (destMemref && srcMemref &&
+ destMemref.getType().getElementType() ==
+ srcMemref.getType().getElementType() &&
+ destMemref.getType().getShape() == srcMemref.getType().getShape()) {
+ memref::CopyOp::create(builder, loc, srcMemref, destMemref);
+ return true;
+ }
+
+ return false;
+ }
};
struct LLVMPointerPointerLikeModel
diff --git a/mlir/test/Dialect/OpenACC/pointer-like-interface-alloc.mlir b/mlir/test/Dialect/OpenACC/pointer-like-interface-alloc.mlir
new file mode 100644
index 0000000000000..603ace85072ac
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/pointer-like-interface-alloc.mlir
@@ -0,0 +1,24 @@
+// RUN: mlir-opt %s --split-input-file --pass-pipeline="builtin.module(func.func(test-acc-pointer-like-interface{test-mode=alloc}))" 2>&1 | FileCheck %s
+
+func.func @test_static_memref_alloc() {
+ %0 = memref.alloca() {test.ptr} : memref<10x20xf32>
+ // CHECK: Successfully generated alloc for operation: %[[ORIG:.*]] = memref.alloca() {test.ptr} : memref<10x20xf32>
+ // CHECK: Generated: %{{.*}} = memref.alloca() : memref<10x20xf32>
+ return
+}
+
+// -----
+
+func.func @test_dynamic_memref_alloc() {
+ %c10 = arith.constant 10 : index
+ %c20 = arith.constant 20 : index
+ %orig = memref.alloc(%c10, %c20) {test.ptr} : memref<?x?xf32>
+
+ // CHECK: Successfully generated alloc for operation: %[[ORIG:.*]] = memref.alloc(%[[C10:.*]], %[[C20:.*]]) {test.ptr} : memref<?x?xf32>
+ // CHECK: Generated: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK: Generated: %[[DIM0:.*]] = memref.dim %[[ORIG]], %[[C0]] : memref<?x?xf32>
+ // CHECK: Generated: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK: Generated: %[[DIM1:.*]] = memref.dim %[[ORIG]], %[[C1]] : memref<?x?xf32>
+ // CHECK: Generated: %{{.*}} = memref.alloc(%[[DIM0]], %[[DIM1]]) : memref<?x?xf32>
+ return
+}
diff --git a/mlir/test/Dialect/OpenACC/pointer-like-interface-copy.mlir b/mlir/test/Dialect/OpenACC/pointer-like-interface-copy.mlir
new file mode 100644
index 0000000000000..9220d84a29f94
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/pointer-like-interface-copy.mlir
@@ -0,0 +1,23 @@
+// RUN: mlir-opt %s --split-input-file --pass-pipeline="builtin.module(func.func(test-acc-pointer-like-interface{test-mode=copy}))" 2>&1 | FileCheck %s
+
+func.func @test_copy_static() {
+ %src = memref.alloca() {test.src_ptr} : memref<10x20xf32>
+ %dest = memref.alloca() {test.dest_ptr} : memref<10x20xf32>
+
+ // CHECK: Successfully generated copy from source: %[[SRC:.*]] = memref.alloca() {test.src_ptr} : memref<10x20xf32> to destination: %[[DEST:.*]] = memref.alloca() {test.dest_ptr} : memref<10x20xf32>
+ // CHECK: Generated: memref.copy %[[SRC]], %[[DEST]] : memref<10x20xf32> to memref<10x20xf32>
+ return
+}
+
+// -----
+
+func.func @test_copy_dynamic() {
+ %c10 = arith.constant 10 : index
+ %c20 = arith.constant 20 : index
+ %src = memref.alloc(%c10, %c20) {test.src_ptr} : memref<?x?xf32>
+ %dest = memref.alloc(%c10, %c20) {test.dest_ptr} : memref<?x?xf32>
+
+ // CHECK: Successfully generated copy from source: %[[SRC:.*]] = memref.alloc(%[[C10:.*]], %[[C20:.*]]) {test.src_ptr} : memref<?x?xf32> to destination: %[[DEST:.*]] = memref.alloc(%[[C10]], %[[C20]]) {test.dest_ptr} : memref<?x?xf32>
+ // CHECK: Generated: memref.copy %[[SRC]], %[[DEST]] : memref<?x?xf32> to memref<?x?xf32>
+ return
+}
diff --git a/mlir/test/Dialect/OpenACC/pointer-like-interface-free.mlir b/mlir/test/Dialect/OpenACC/pointer-like-interface-free.mlir
new file mode 100644
index 0000000000000..ecf4f75fe06d1
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/pointer-like-interface-free.mlir
@@ -0,0 +1,31 @@
+// RUN: mlir-opt %s --split-input-file --pass-pipeline="builtin.module(func.func(test-acc-pointer-like-interface{test-mode=free}))" 2>&1 | FileCheck %s
+
+func.func @test_static_memref_free() {
+ %0 = memref.alloca() {test.ptr} : memref<10x20xf32>
+ // CHECK: Successfully generated free for operation: %[[ORIG:.*]] = memref.alloca() {test.ptr} : memref<10x20xf32>
+ // CHECK-NOT: Generated
+ return
+}
+
+// -----
+
+func.func @test_dynamic_memref_free() {
+ %c10 = arith.constant 10 : index
+ %c20 = arith.constant 20 : index
+ %orig = memref.alloc(%c10, %c20) {test.ptr} : memref<?x?xf32>
+
+ // CHECK: Successfully generated free for operation: %[[ORIG:.*]] = memref.alloc(%[[C10:.*]], %[[C20:.*]]) {test.ptr} : memref<?x?xf32>
+ // CHECK: Generated: memref.dealloc %[[ORIG]] : memref<?x?xf32>
+ return
+}
+
+// -----
+
+func.func @test_cast_walking_free() {
+ %0 = memref.alloca() : memref<10x20xf32>
+ %1 = memref.cast %0 {test.ptr} : memref<10x20xf32> to memref<?x?xf32>
+
+ // CHECK: Successfully generated free for operation: %[[CAST:.*]] = memref.cast %[[ALLOCA:.*]] {test.ptr} : memref<10x20xf32> to memref<?x?xf32>
+ // CHECK-NOT: Generated
+ return
+}
diff --git a/mlir/test/lib/Dialect/CMakeLists.txt b/mlir/test/lib/Dialect/CMakeLists.txt
index 3b7bd9b9637a8..e31140a6f6ba7 100644
--- a/mlir/test/lib/Dialect/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/CMakeLists.txt
@@ -12,6 +12,7 @@ add_subdirectory(Math)
add_subdirectory(MemRef)
add_subdirectory(Shard)
add_subdirectory(NVGPU)
+add_subdirectory(OpenACC)
add_subdirectory(SCF)
add_subdirectory(Shape)
add_subdirectory(SPIRV)
diff --git a/mlir/test/lib/Dialect/OpenACC/CMakeLists.txt b/mlir/test/lib/Dialect/OpenACC/CMakeLists.txt
new file mode 100644
index 0000000000000..f84055df1b6c6
--- /dev/null
+++ b/mlir/test/lib/Dialect/OpenACC/CMakeLists.txt
@@ -0,0 +1,16 @@
+add_mlir_library(MLIROpenACCTestPasses
+ TestOpenACC.cpp
+ TestPointerLikeTypeInterface.cpp
+
+ EXCLUDE_FROM_LIBMLIR
+)
+mlir_target_link_libraries(MLIROpenACCTestPasses PUBLIC
+ MLIRIR
+ MLIRArithDialect
+ MLIRFuncDialect
+ MLIRMemRefDialect
+ MLIROpenACCDialect
+ MLIRPass
+ MLIRSupport
+)
+
diff --git a/mlir/test/lib/Dialect/OpenACC/TestOpenACC.cpp b/mlir/test/lib/Dialect/OpenACC/TestOpenACC.cpp
new file mode 100644
index 0000000000000..6f7bdde176aa7
--- /dev/null
+++ b/mlir/test/lib/Dialect/OpenACC/TestOpenACC.cpp
@@ -0,0 +1,26 @@
+//===- TestOpenACC.cpp - OpenACC Test Registration ------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains unified registration for all OpenACC test passes.
+//
+//===----------------------------------------------------------------------===//
+
+namespace mlir {
+namespace test {
+
+// Forward declarations of individual test pass registration functions
+void registerTestPointerLikeTypeInterfacePass();
+
+// Unified registration function for all OpenACC tests
+void registerTestOpenACC() {
+ registerTestPointerLikeTypeInterfacePass();
+}
+
+} // namespace test
+} // namespace mlir
+
diff --git a/mlir/test/lib/Dialect/OpenACC/TestPointerLikeTypeInterface.cpp b/mlir/test/lib/Dialect/OpenACC/TestPointerLikeTypeInterface.cpp
new file mode 100644
index 0000000000000..fc573d4318b5b
--- /dev/null
+++ b/mlir/test/lib/Dialect/OpenACC/TestPointerLikeTypeInterface.cpp
@@ -0,0 +1,310 @@
+//===- TestPointerLikeTypeInterface.cpp - Test PointerLikeType interface -===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains test passes for testing the OpenACC PointerLikeType
+// interface methods.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/Support/CommandLine.h"
+
+using namespace mlir;
+using namespace mlir::acc;
+
+namespace {
+
+struct OperationTracker : public OpBuilder::Listener {
+ SmallVector<Operation *> insertedOps;
+
+ void notifyOperationInserted(Operation *op,
+ OpBuilder::InsertPoint previous) override {
+ insertedOps.push_back(op);
+ }
+};
+
+struct TestPointerLikeTypeInterfacePass
+ : public PassWrapper<TestPointerLikeTypeInterfacePass,
+ OperationPass<func::FuncOp>> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestPointerLikeTypeInterfacePass)
+
+ TestPointerLikeTypeInterfacePass() = default;
+ TestPointerLikeTypeInterfacePass(const TestPointerLikeTypeInterfacePass &pass)
+ : PassWrapper(pass) {
+ testMode = pass.testMode;
+ }
+
+ Pass::Option<std::string> testMode{
+ *this, "test-mode",
+ llvm::cl::desc("Test mode: walk, alloc, copy, or free"),
+ llvm::cl::init("walk")};
+
+ StringRef getArgument() const override {
+ return "test-acc-pointer-like-interface";
+ }
+
+ StringRef getDescription() const override {
+ return "Test OpenACC PointerLikeType interface methods on any implementing "
+ "type";
+ }
+
+ void runOnOperation() override;
+
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<acc::OpenACCDialect>();
+ registry.insert<arith::ArithDialect>();
+ registry.insert<memref::MemRefDialect>();
+ }
+
+private:
+ void walkAndPrint();
+ void testGenAllocate(Operation *op, Value result, PointerLikeType pointerType,
+ OpBuilder &builder);
+ void testGenFree(Operation *op, Value result, PointerLikeType pointerType,
+ OpBuilder &builder);
+ void testGenCopy(Operation *srcOp, Operation *destOp, Value srcResult, Value destResult,
+ PointerLikeType pointerType, OpBuilder &builder);
+
+ struct PointerCandidate {
+ Operation *op;
+ Value result;
+ PointerLikeType pointerType;
+ };
+};
+
+void TestPointerLikeTypeInterfacePass::runOnOperation() {
+ if (testMode == "walk") {
+ walkAndPrint();
+ return;
+ }
+
+ auto func = getOperation();
+ OpBuilder builder(&getContext());
+
+ if (testMode == "alloc" || testMode == "free") {
+ // Collect all candidates first
+ SmallVector<PointerCandidate> candidates;
+ func.walk([&](Operation *op) {
+ if (op->hasAttr("test.ptr")) {
+ for (auto result : op->getResults()) {
+ if (isa<PointerLikeType>(result.getType())) {
+ candidates.push_back(
+ {op, result, cast<PointerLikeType>(result.getType())});
+ break; // Only take the first PointerLikeType result
+ }
+ }
+ }
+ });
+
+ // Now test all candidates
+ for (const auto &candidate : candidates) {
+ if (testMode == "alloc") {
+ testGenAllocate(candidate.op, candidate.result, candidate.pointerType,
+ builder);
+ } else if (testMode == "free") {
+ testGenFree(candidate.op, candidate.result, candidate.pointerType,
+ builder);
+ }
+ }
+ } else if (testMode == "copy") {
+ // Collect all source and destination candidates
+ SmallVector<PointerCandidate> sources, destinations;
+
+ func.walk([&](Operation *op) {
+ if (op->hasAttr("test.src_ptr")) {
+ for (auto resul...
[truncated]
|
@llvm/pr-subscribers-mlir Author: Razvan Lupusoru (razvanlupusoru) ChangesA variable in an acc data clause operation must have a type that implements either PointerLikeType or a MappableType interface. These interfaces provide the contract that allows acc dialect and its transform passes to interact with a source dialect. One of these requirements is ability to generate code that creates memory for a private copy and ability to initialize that copy from another variable. Thus, update the PointerLikeType API to provide the means to create allocation, deallocation, and copy. This will be used as a way to fill in privatization and firstprivatization recipes. This new API was implemented for memref along with testing to exercise it via the implementation of a testing pass. Patch is 27.53 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/162328.diff 11 Files Affected:
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCTypeInterfaces.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCTypeInterfaces.td
index 9123ac34af67d..9a89f8ede191e 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCTypeInterfaces.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCTypeInterfaces.td
@@ -70,6 +70,85 @@ def OpenACC_PointerLikeTypeInterface : TypeInterface<"PointerLikeType"> {
return ::mlir::acc::VariableTypeCategory::uncategorized;
}]
>,
+ InterfaceMethod<
+ /*description=*/[{
+ Generates allocation operations for the pointer-like type. It will create
+ an allocate that produces memory space for an instance of the current type.
+
+ The `varName` parameter is optional and can be used to provide a name
+ for the allocated variable. If the current type is represented
+ in a way that it does not capture the pointee type, `varType` must be
+ passed in to provide the necessary type information.
+
+ The `originalVar` parameter is optional but enables support for dynamic
+ types (e.g., dynamic memrefs). When provided, implementations can extract
+ runtime dimension information from the original variable to create
+ allocations with matching dynamic sizes.
+
+ Returns true if allocation was successfully generated, false otherwise.
+ }],
+ /*retTy=*/"bool",
+ /*methodName=*/"genAllocate",
+ /*args=*/(ins "::mlir::OpBuilder &":$builder,
+ "::mlir::Location":$loc,
+ "::llvm::StringRef":$varName,
+ "::mlir::Type":$varType,
+ "::mlir::Value":$originalVar),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return false;
+ }]
+ >,
+ InterfaceMethod<
+ /*description=*/[{
+ Generates deallocation operations for the pointer-like type. It deallocates
+ the instance provided.
+
+ The `varPtr` parameter is required and must represent an instance that was
+ previously allocated. If the current type is represented in a way that it
+ does not capture the pointee type, `varType` must be passed in to provide
+ the necessary type information. Nothing is generated in case the allocate
+ is `alloca`-like.
+
+ Returns true if deallocation was successfully generated or successfully
+ deemed as not needed to be generated, false otherwise.
+ }],
+ /*retTy=*/"bool",
+ /*methodName=*/"genFree",
+ /*args=*/(ins "::mlir::OpBuilder &":$builder,
+ "::mlir::Location":$loc,
+ "::mlir::TypedValue<::mlir::acc::PointerLikeType>":$varPtr,
+ "::mlir::Type":$varType),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return false;
+ }]
+ >,
+ InterfaceMethod<
+ /*description=*/[{
+ Generates copy operations for the pointer-like type. It copies the memory
+ from the source to the destination. Typically used to initialize one
+ variable of this type from another.
+
+ The `destination` and `source` parameters represent the target and source
+ instances respectively. If the current type is represented in a way that it
+ does not capture the pointee type, `varType` must be passed in to provide
+ the necessary type information.
+
+ Returns true if copy was successfully generated, false otherwise.
+ }],
+ /*retTy=*/"bool",
+ /*methodName=*/"genCopy",
+ /*args=*/(ins "::mlir::OpBuilder &":$builder,
+ "::mlir::Location":$loc,
+ "::mlir::TypedValue<::mlir::acc::PointerLikeType>":$destination,
+ "::mlir::TypedValue<::mlir::acc::PointerLikeType>":$source,
+ "::mlir::Type":$varType),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return false;
+ }]
+ >,
];
}
diff --git a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
index 6598ac141008f..c2a21acb4d62c 100644
--- a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
+++ b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
@@ -7,6 +7,7 @@
// =============================================================================
#include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
@@ -44,6 +45,7 @@ struct MemRefPointerLikeModel
Type getElementType(Type pointer) const {
return cast<MemRefType>(pointer).getElementType();
}
+
mlir::acc::VariableTypeCategory
getPointeeTypeCategory(Type pointer, TypedValue<PointerLikeType> varPtr,
Type varType) const {
@@ -70,6 +72,117 @@ struct MemRefPointerLikeModel
assert(memrefTy.getRank() > 0 && "rank expected to be positive");
return mlir::acc::VariableTypeCategory::array;
}
+
+ bool genAllocate(Type pointer, OpBuilder &builder, Location loc,
+ StringRef varName, Type varType, Value originalVar) const {
+ auto memrefTy = cast<MemRefType>(pointer);
+
+ // Check if this is a static memref (all dimensions are known) - if yes
+ // then we can generate an alloca operation.
+ if (memrefTy.hasStaticShape()) {
+ memref::AllocaOp::create(builder, loc, memrefTy);
+ return true;
+ }
+
+ // For dynamic memrefs, extract sizes from the original variable if provided.
+ // Otherwise they cannot be handled.
+ if (originalVar && originalVar.getType() == memrefTy &&
+ memrefTy.hasRank()) {
+ SmallVector<Value> dynamicSizes;
+ for (int64_t i = 0; i < memrefTy.getRank(); ++i) {
+ if (memrefTy.isDynamicDim(i)) {
+ // Extract the size of dimension i from the original variable
+ auto indexValue = arith::ConstantIndexOp::create(builder, loc, i);
+ auto dimSize =
+ memref::DimOp::create(builder, loc, originalVar, indexValue);
+ dynamicSizes.push_back(dimSize);
+ }
+ // Note: We only add dynamic sizes to the dynamicSizes array
+ // Static dimensions are handled automatically by AllocOp
+ }
+ memref::AllocOp::create(builder, loc, memrefTy, dynamicSizes);
+ return true;
+ }
+
+
+ // TODO: Unranked not yet supported.
+ return false;
+ }
+
+ bool genFree(Type pointer, OpBuilder &builder, Location loc,
+ TypedValue<PointerLikeType> varPtr, Type varType) const {
+ if (auto memrefValue = dyn_cast<TypedValue<MemRefType>>(varPtr)) {
+ // Walk through casts to find the original allocation
+ Value currentValue = memrefValue;
+ Operation *originalAlloc = nullptr;
+
+ // Follow the chain of operations to find the original allocation
+ // even if a casted result is provided.
+ while (currentValue) {
+ if (auto *definingOp = currentValue.getDefiningOp()) {
+ // Check if this is an allocation operation
+ if (isa<memref::AllocOp, memref::AllocaOp>(definingOp)) {
+ originalAlloc = definingOp;
+ break;
+ }
+
+ // Check if this is a cast operation we can look through
+ if (auto castOp = dyn_cast<memref::CastOp>(definingOp)) {
+ currentValue = castOp.getSource();
+ continue;
+ }
+
+ // Check for other cast-like operations
+ if (auto reinterpretCastOp =
+ dyn_cast<memref::ReinterpretCastOp>(definingOp)) {
+ currentValue = reinterpretCastOp.getSource();
+ continue;
+ }
+
+ // If we can't look through this operation, stop
+ break;
+ }
+ // This is a block argument or similar - can't trace further.
+ break;
+ }
+
+ if (originalAlloc) {
+ if (isa<memref::AllocaOp>(originalAlloc)) {
+ // This is an alloca - no dealloc needed, but return true (success)
+ return true;
+ }
+ if (isa<memref::AllocOp>(originalAlloc)) {
+ // This is an alloc - generate dealloc
+ memref::DeallocOp::create(builder, loc, memrefValue);
+ return true;
+ }
+ }
+ }
+
+ return false;
+ }
+
+ bool genCopy(Type pointer, OpBuilder &builder, Location loc,
+ TypedValue<PointerLikeType> destination,
+ TypedValue<PointerLikeType> source, Type varType) const {
+ // Generate a copy operation between two memrefs
+ auto destMemref = dyn_cast_if_present<TypedValue<MemRefType>>(destination);
+ auto srcMemref = dyn_cast_if_present<TypedValue<MemRefType>>(source);
+
+ // As per memref documentation, source and destination must have same
+ // element type and shape in order to be compatible. We do not want to fail
+ // with an IR verification error - thus check that before generating the
+ // copy operation.
+ if (destMemref && srcMemref &&
+ destMemref.getType().getElementType() ==
+ srcMemref.getType().getElementType() &&
+ destMemref.getType().getShape() == srcMemref.getType().getShape()) {
+ memref::CopyOp::create(builder, loc, srcMemref, destMemref);
+ return true;
+ }
+
+ return false;
+ }
};
struct LLVMPointerPointerLikeModel
diff --git a/mlir/test/Dialect/OpenACC/pointer-like-interface-alloc.mlir b/mlir/test/Dialect/OpenACC/pointer-like-interface-alloc.mlir
new file mode 100644
index 0000000000000..603ace85072ac
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/pointer-like-interface-alloc.mlir
@@ -0,0 +1,24 @@
+// RUN: mlir-opt %s --split-input-file --pass-pipeline="builtin.module(func.func(test-acc-pointer-like-interface{test-mode=alloc}))" 2>&1 | FileCheck %s
+
+func.func @test_static_memref_alloc() {
+ %0 = memref.alloca() {test.ptr} : memref<10x20xf32>
+ // CHECK: Successfully generated alloc for operation: %[[ORIG:.*]] = memref.alloca() {test.ptr} : memref<10x20xf32>
+ // CHECK: Generated: %{{.*}} = memref.alloca() : memref<10x20xf32>
+ return
+}
+
+// -----
+
+func.func @test_dynamic_memref_alloc() {
+ %c10 = arith.constant 10 : index
+ %c20 = arith.constant 20 : index
+ %orig = memref.alloc(%c10, %c20) {test.ptr} : memref<?x?xf32>
+
+ // CHECK: Successfully generated alloc for operation: %[[ORIG:.*]] = memref.alloc(%[[C10:.*]], %[[C20:.*]]) {test.ptr} : memref<?x?xf32>
+ // CHECK: Generated: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK: Generated: %[[DIM0:.*]] = memref.dim %[[ORIG]], %[[C0]] : memref<?x?xf32>
+ // CHECK: Generated: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK: Generated: %[[DIM1:.*]] = memref.dim %[[ORIG]], %[[C1]] : memref<?x?xf32>
+ // CHECK: Generated: %{{.*}} = memref.alloc(%[[DIM0]], %[[DIM1]]) : memref<?x?xf32>
+ return
+}
diff --git a/mlir/test/Dialect/OpenACC/pointer-like-interface-copy.mlir b/mlir/test/Dialect/OpenACC/pointer-like-interface-copy.mlir
new file mode 100644
index 0000000000000..9220d84a29f94
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/pointer-like-interface-copy.mlir
@@ -0,0 +1,23 @@
+// RUN: mlir-opt %s --split-input-file --pass-pipeline="builtin.module(func.func(test-acc-pointer-like-interface{test-mode=copy}))" 2>&1 | FileCheck %s
+
+func.func @test_copy_static() {
+ %src = memref.alloca() {test.src_ptr} : memref<10x20xf32>
+ %dest = memref.alloca() {test.dest_ptr} : memref<10x20xf32>
+
+ // CHECK: Successfully generated copy from source: %[[SRC:.*]] = memref.alloca() {test.src_ptr} : memref<10x20xf32> to destination: %[[DEST:.*]] = memref.alloca() {test.dest_ptr} : memref<10x20xf32>
+ // CHECK: Generated: memref.copy %[[SRC]], %[[DEST]] : memref<10x20xf32> to memref<10x20xf32>
+ return
+}
+
+// -----
+
+func.func @test_copy_dynamic() {
+ %c10 = arith.constant 10 : index
+ %c20 = arith.constant 20 : index
+ %src = memref.alloc(%c10, %c20) {test.src_ptr} : memref<?x?xf32>
+ %dest = memref.alloc(%c10, %c20) {test.dest_ptr} : memref<?x?xf32>
+
+ // CHECK: Successfully generated copy from source: %[[SRC:.*]] = memref.alloc(%[[C10:.*]], %[[C20:.*]]) {test.src_ptr} : memref<?x?xf32> to destination: %[[DEST:.*]] = memref.alloc(%[[C10]], %[[C20]]) {test.dest_ptr} : memref<?x?xf32>
+ // CHECK: Generated: memref.copy %[[SRC]], %[[DEST]] : memref<?x?xf32> to memref<?x?xf32>
+ return
+}
diff --git a/mlir/test/Dialect/OpenACC/pointer-like-interface-free.mlir b/mlir/test/Dialect/OpenACC/pointer-like-interface-free.mlir
new file mode 100644
index 0000000000000..ecf4f75fe06d1
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/pointer-like-interface-free.mlir
@@ -0,0 +1,31 @@
+// RUN: mlir-opt %s --split-input-file --pass-pipeline="builtin.module(func.func(test-acc-pointer-like-interface{test-mode=free}))" 2>&1 | FileCheck %s
+
+func.func @test_static_memref_free() {
+ %0 = memref.alloca() {test.ptr} : memref<10x20xf32>
+ // CHECK: Successfully generated free for operation: %[[ORIG:.*]] = memref.alloca() {test.ptr} : memref<10x20xf32>
+ // CHECK-NOT: Generated
+ return
+}
+
+// -----
+
+func.func @test_dynamic_memref_free() {
+ %c10 = arith.constant 10 : index
+ %c20 = arith.constant 20 : index
+ %orig = memref.alloc(%c10, %c20) {test.ptr} : memref<?x?xf32>
+
+ // CHECK: Successfully generated free for operation: %[[ORIG:.*]] = memref.alloc(%[[C10:.*]], %[[C20:.*]]) {test.ptr} : memref<?x?xf32>
+ // CHECK: Generated: memref.dealloc %[[ORIG]] : memref<?x?xf32>
+ return
+}
+
+// -----
+
+func.func @test_cast_walking_free() {
+ %0 = memref.alloca() : memref<10x20xf32>
+ %1 = memref.cast %0 {test.ptr} : memref<10x20xf32> to memref<?x?xf32>
+
+ // CHECK: Successfully generated free for operation: %[[CAST:.*]] = memref.cast %[[ALLOCA:.*]] {test.ptr} : memref<10x20xf32> to memref<?x?xf32>
+ // CHECK-NOT: Generated
+ return
+}
diff --git a/mlir/test/lib/Dialect/CMakeLists.txt b/mlir/test/lib/Dialect/CMakeLists.txt
index 3b7bd9b9637a8..e31140a6f6ba7 100644
--- a/mlir/test/lib/Dialect/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/CMakeLists.txt
@@ -12,6 +12,7 @@ add_subdirectory(Math)
add_subdirectory(MemRef)
add_subdirectory(Shard)
add_subdirectory(NVGPU)
+add_subdirectory(OpenACC)
add_subdirectory(SCF)
add_subdirectory(Shape)
add_subdirectory(SPIRV)
diff --git a/mlir/test/lib/Dialect/OpenACC/CMakeLists.txt b/mlir/test/lib/Dialect/OpenACC/CMakeLists.txt
new file mode 100644
index 0000000000000..f84055df1b6c6
--- /dev/null
+++ b/mlir/test/lib/Dialect/OpenACC/CMakeLists.txt
@@ -0,0 +1,16 @@
+add_mlir_library(MLIROpenACCTestPasses
+ TestOpenACC.cpp
+ TestPointerLikeTypeInterface.cpp
+
+ EXCLUDE_FROM_LIBMLIR
+)
+mlir_target_link_libraries(MLIROpenACCTestPasses PUBLIC
+ MLIRIR
+ MLIRArithDialect
+ MLIRFuncDialect
+ MLIRMemRefDialect
+ MLIROpenACCDialect
+ MLIRPass
+ MLIRSupport
+)
+
diff --git a/mlir/test/lib/Dialect/OpenACC/TestOpenACC.cpp b/mlir/test/lib/Dialect/OpenACC/TestOpenACC.cpp
new file mode 100644
index 0000000000000..6f7bdde176aa7
--- /dev/null
+++ b/mlir/test/lib/Dialect/OpenACC/TestOpenACC.cpp
@@ -0,0 +1,26 @@
+//===- TestOpenACC.cpp - OpenACC Test Registration ------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains unified registration for all OpenACC test passes.
+//
+//===----------------------------------------------------------------------===//
+
+namespace mlir {
+namespace test {
+
+// Forward declarations of individual test pass registration functions
+void registerTestPointerLikeTypeInterfacePass();
+
+// Unified registration function for all OpenACC tests
+void registerTestOpenACC() {
+ registerTestPointerLikeTypeInterfacePass();
+}
+
+} // namespace test
+} // namespace mlir
+
diff --git a/mlir/test/lib/Dialect/OpenACC/TestPointerLikeTypeInterface.cpp b/mlir/test/lib/Dialect/OpenACC/TestPointerLikeTypeInterface.cpp
new file mode 100644
index 0000000000000..fc573d4318b5b
--- /dev/null
+++ b/mlir/test/lib/Dialect/OpenACC/TestPointerLikeTypeInterface.cpp
@@ -0,0 +1,310 @@
+//===- TestPointerLikeTypeInterface.cpp - Test PointerLikeType interface -===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains test passes for testing the OpenACC PointerLikeType
+// interface methods.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/Support/CommandLine.h"
+
+using namespace mlir;
+using namespace mlir::acc;
+
+namespace {
+
+struct OperationTracker : public OpBuilder::Listener {
+ SmallVector<Operation *> insertedOps;
+
+ void notifyOperationInserted(Operation *op,
+ OpBuilder::InsertPoint previous) override {
+ insertedOps.push_back(op);
+ }
+};
+
+struct TestPointerLikeTypeInterfacePass
+ : public PassWrapper<TestPointerLikeTypeInterfacePass,
+ OperationPass<func::FuncOp>> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestPointerLikeTypeInterfacePass)
+
+ TestPointerLikeTypeInterfacePass() = default;
+ TestPointerLikeTypeInterfacePass(const TestPointerLikeTypeInterfacePass &pass)
+ : PassWrapper(pass) {
+ testMode = pass.testMode;
+ }
+
+ Pass::Option<std::string> testMode{
+ *this, "test-mode",
+ llvm::cl::desc("Test mode: walk, alloc, copy, or free"),
+ llvm::cl::init("walk")};
+
+ StringRef getArgument() const override {
+ return "test-acc-pointer-like-interface";
+ }
+
+ StringRef getDescription() const override {
+ return "Test OpenACC PointerLikeType interface methods on any implementing "
+ "type";
+ }
+
+ void runOnOperation() override;
+
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<acc::OpenACCDialect>();
+ registry.insert<arith::ArithDialect>();
+ registry.insert<memref::MemRefDialect>();
+ }
+
+private:
+ void walkAndPrint();
+ void testGenAllocate(Operation *op, Value result, PointerLikeType pointerType,
+ OpBuilder &builder);
+ void testGenFree(Operation *op, Value result, PointerLikeType pointerType,
+ OpBuilder &builder);
+ void testGenCopy(Operation *srcOp, Operation *destOp, Value srcResult, Value destResult,
+ PointerLikeType pointerType, OpBuilder &builder);
+
+ struct PointerCandidate {
+ Operation *op;
+ Value result;
+ PointerLikeType pointerType;
+ };
+};
+
+void TestPointerLikeTypeInterfacePass::runOnOperation() {
+ if (testMode == "walk") {
+ walkAndPrint();
+ return;
+ }
+
+ auto func = getOperation();
+ OpBuilder builder(&getContext());
+
+ if (testMode == "alloc" || testMode == "free") {
+ // Collect all candidates first
+ SmallVector<PointerCandidate> candidates;
+ func.walk([&](Operation *op) {
+ if (op->hasAttr("test.ptr")) {
+ for (auto result : op->getResults()) {
+ if (isa<PointerLikeType>(result.getType())) {
+ candidates.push_back(
+ {op, result, cast<PointerLikeType>(result.getType())});
+ break; // Only take the first PointerLikeType result
+ }
+ }
+ }
+ });
+
+ // Now test all candidates
+ for (const auto &candidate : candidates) {
+ if (testMode == "alloc") {
+ testGenAllocate(candidate.op, candidate.result, candidate.pointerType,
+ builder);
+ } else if (testMode == "free") {
+ testGenFree(candidate.op, candidate.result, candidate.pointerType,
+ builder);
+ }
+ }
+ } else if (testMode == "copy") {
+ // Collect all source and destination candidates
+ SmallVector<PointerCandidate> sources, destinations;
+
+ func.walk([&](Operation *op) {
+ if (op->hasAttr("test.src_ptr")) {
+ for (auto resul...
[truncated]
|
@llvm/pr-subscribers-mlir-core Author: Razvan Lupusoru (razvanlupusoru) ChangesA variable in an acc data clause operation must have a type that implements either PointerLikeType or a MappableType interface. These interfaces provide the contract that allows acc dialect and its transform passes to interact with a source dialect. One of these requirements is ability to generate code that creates memory for a private copy and ability to initialize that copy from another variable. Thus, update the PointerLikeType API to provide the means to create allocation, deallocation, and copy. This will be used as a way to fill in privatization and firstprivatization recipes. This new API was implemented for memref along with testing to exercise it via the implementation of a testing pass. Patch is 27.53 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/162328.diff 11 Files Affected:
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCTypeInterfaces.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCTypeInterfaces.td
index 9123ac34af67d..9a89f8ede191e 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCTypeInterfaces.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCTypeInterfaces.td
@@ -70,6 +70,85 @@ def OpenACC_PointerLikeTypeInterface : TypeInterface<"PointerLikeType"> {
return ::mlir::acc::VariableTypeCategory::uncategorized;
}]
>,
+ InterfaceMethod<
+ /*description=*/[{
+ Generates allocation operations for the pointer-like type. It will create
+ an allocate that produces memory space for an instance of the current type.
+
+ The `varName` parameter is optional and can be used to provide a name
+ for the allocated variable. If the current type is represented
+ in a way that it does not capture the pointee type, `varType` must be
+ passed in to provide the necessary type information.
+
+ The `originalVar` parameter is optional but enables support for dynamic
+ types (e.g., dynamic memrefs). When provided, implementations can extract
+ runtime dimension information from the original variable to create
+ allocations with matching dynamic sizes.
+
+ Returns true if allocation was successfully generated, false otherwise.
+ }],
+ /*retTy=*/"bool",
+ /*methodName=*/"genAllocate",
+ /*args=*/(ins "::mlir::OpBuilder &":$builder,
+ "::mlir::Location":$loc,
+ "::llvm::StringRef":$varName,
+ "::mlir::Type":$varType,
+ "::mlir::Value":$originalVar),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return false;
+ }]
+ >,
+ InterfaceMethod<
+ /*description=*/[{
+ Generates deallocation operations for the pointer-like type. It deallocates
+ the instance provided.
+
+ The `varPtr` parameter is required and must represent an instance that was
+ previously allocated. If the current type is represented in a way that it
+ does not capture the pointee type, `varType` must be passed in to provide
+ the necessary type information. Nothing is generated in case the allocate
+ is `alloca`-like.
+
+ Returns true if deallocation was successfully generated or successfully
+ deemed as not needed to be generated, false otherwise.
+ }],
+ /*retTy=*/"bool",
+ /*methodName=*/"genFree",
+ /*args=*/(ins "::mlir::OpBuilder &":$builder,
+ "::mlir::Location":$loc,
+ "::mlir::TypedValue<::mlir::acc::PointerLikeType>":$varPtr,
+ "::mlir::Type":$varType),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return false;
+ }]
+ >,
+ InterfaceMethod<
+ /*description=*/[{
+ Generates copy operations for the pointer-like type. It copies the memory
+ from the source to the destination. Typically used to initialize one
+ variable of this type from another.
+
+ The `destination` and `source` parameters represent the target and source
+ instances respectively. If the current type is represented in a way that it
+ does not capture the pointee type, `varType` must be passed in to provide
+ the necessary type information.
+
+ Returns true if copy was successfully generated, false otherwise.
+ }],
+ /*retTy=*/"bool",
+ /*methodName=*/"genCopy",
+ /*args=*/(ins "::mlir::OpBuilder &":$builder,
+ "::mlir::Location":$loc,
+ "::mlir::TypedValue<::mlir::acc::PointerLikeType>":$destination,
+ "::mlir::TypedValue<::mlir::acc::PointerLikeType>":$source,
+ "::mlir::Type":$varType),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return false;
+ }]
+ >,
];
}
diff --git a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
index 6598ac141008f..c2a21acb4d62c 100644
--- a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
+++ b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
@@ -7,6 +7,7 @@
// =============================================================================
#include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
@@ -44,6 +45,7 @@ struct MemRefPointerLikeModel
Type getElementType(Type pointer) const {
return cast<MemRefType>(pointer).getElementType();
}
+
mlir::acc::VariableTypeCategory
getPointeeTypeCategory(Type pointer, TypedValue<PointerLikeType> varPtr,
Type varType) const {
@@ -70,6 +72,117 @@ struct MemRefPointerLikeModel
assert(memrefTy.getRank() > 0 && "rank expected to be positive");
return mlir::acc::VariableTypeCategory::array;
}
+
+ bool genAllocate(Type pointer, OpBuilder &builder, Location loc,
+ StringRef varName, Type varType, Value originalVar) const {
+ auto memrefTy = cast<MemRefType>(pointer);
+
+ // Check if this is a static memref (all dimensions are known) - if yes
+ // then we can generate an alloca operation.
+ if (memrefTy.hasStaticShape()) {
+ memref::AllocaOp::create(builder, loc, memrefTy);
+ return true;
+ }
+
+ // For dynamic memrefs, extract sizes from the original variable if provided.
+ // Otherwise they cannot be handled.
+ if (originalVar && originalVar.getType() == memrefTy &&
+ memrefTy.hasRank()) {
+ SmallVector<Value> dynamicSizes;
+ for (int64_t i = 0; i < memrefTy.getRank(); ++i) {
+ if (memrefTy.isDynamicDim(i)) {
+ // Extract the size of dimension i from the original variable
+ auto indexValue = arith::ConstantIndexOp::create(builder, loc, i);
+ auto dimSize =
+ memref::DimOp::create(builder, loc, originalVar, indexValue);
+ dynamicSizes.push_back(dimSize);
+ }
+ // Note: We only add dynamic sizes to the dynamicSizes array
+ // Static dimensions are handled automatically by AllocOp
+ }
+ memref::AllocOp::create(builder, loc, memrefTy, dynamicSizes);
+ return true;
+ }
+
+
+ // TODO: Unranked not yet supported.
+ return false;
+ }
+
+ bool genFree(Type pointer, OpBuilder &builder, Location loc,
+ TypedValue<PointerLikeType> varPtr, Type varType) const {
+ if (auto memrefValue = dyn_cast<TypedValue<MemRefType>>(varPtr)) {
+ // Walk through casts to find the original allocation
+ Value currentValue = memrefValue;
+ Operation *originalAlloc = nullptr;
+
+ // Follow the chain of operations to find the original allocation
+ // even if a casted result is provided.
+ while (currentValue) {
+ if (auto *definingOp = currentValue.getDefiningOp()) {
+ // Check if this is an allocation operation
+ if (isa<memref::AllocOp, memref::AllocaOp>(definingOp)) {
+ originalAlloc = definingOp;
+ break;
+ }
+
+ // Check if this is a cast operation we can look through
+ if (auto castOp = dyn_cast<memref::CastOp>(definingOp)) {
+ currentValue = castOp.getSource();
+ continue;
+ }
+
+ // Check for other cast-like operations
+ if (auto reinterpretCastOp =
+ dyn_cast<memref::ReinterpretCastOp>(definingOp)) {
+ currentValue = reinterpretCastOp.getSource();
+ continue;
+ }
+
+ // If we can't look through this operation, stop
+ break;
+ }
+ // This is a block argument or similar - can't trace further.
+ break;
+ }
+
+ if (originalAlloc) {
+ if (isa<memref::AllocaOp>(originalAlloc)) {
+ // This is an alloca - no dealloc needed, but return true (success)
+ return true;
+ }
+ if (isa<memref::AllocOp>(originalAlloc)) {
+ // This is an alloc - generate dealloc
+ memref::DeallocOp::create(builder, loc, memrefValue);
+ return true;
+ }
+ }
+ }
+
+ return false;
+ }
+
+ bool genCopy(Type pointer, OpBuilder &builder, Location loc,
+ TypedValue<PointerLikeType> destination,
+ TypedValue<PointerLikeType> source, Type varType) const {
+ // Generate a copy operation between two memrefs
+ auto destMemref = dyn_cast_if_present<TypedValue<MemRefType>>(destination);
+ auto srcMemref = dyn_cast_if_present<TypedValue<MemRefType>>(source);
+
+ // As per memref documentation, source and destination must have same
+ // element type and shape in order to be compatible. We do not want to fail
+ // with an IR verification error - thus check that before generating the
+ // copy operation.
+ if (destMemref && srcMemref &&
+ destMemref.getType().getElementType() ==
+ srcMemref.getType().getElementType() &&
+ destMemref.getType().getShape() == srcMemref.getType().getShape()) {
+ memref::CopyOp::create(builder, loc, srcMemref, destMemref);
+ return true;
+ }
+
+ return false;
+ }
};
struct LLVMPointerPointerLikeModel
diff --git a/mlir/test/Dialect/OpenACC/pointer-like-interface-alloc.mlir b/mlir/test/Dialect/OpenACC/pointer-like-interface-alloc.mlir
new file mode 100644
index 0000000000000..603ace85072ac
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/pointer-like-interface-alloc.mlir
@@ -0,0 +1,24 @@
+// RUN: mlir-opt %s --split-input-file --pass-pipeline="builtin.module(func.func(test-acc-pointer-like-interface{test-mode=alloc}))" 2>&1 | FileCheck %s
+
+func.func @test_static_memref_alloc() {
+ %0 = memref.alloca() {test.ptr} : memref<10x20xf32>
+ // CHECK: Successfully generated alloc for operation: %[[ORIG:.*]] = memref.alloca() {test.ptr} : memref<10x20xf32>
+ // CHECK: Generated: %{{.*}} = memref.alloca() : memref<10x20xf32>
+ return
+}
+
+// -----
+
+func.func @test_dynamic_memref_alloc() {
+ %c10 = arith.constant 10 : index
+ %c20 = arith.constant 20 : index
+ %orig = memref.alloc(%c10, %c20) {test.ptr} : memref<?x?xf32>
+
+ // CHECK: Successfully generated alloc for operation: %[[ORIG:.*]] = memref.alloc(%[[C10:.*]], %[[C20:.*]]) {test.ptr} : memref<?x?xf32>
+ // CHECK: Generated: %[[C0:.*]] = arith.constant 0 : index
+ // CHECK: Generated: %[[DIM0:.*]] = memref.dim %[[ORIG]], %[[C0]] : memref<?x?xf32>
+ // CHECK: Generated: %[[C1:.*]] = arith.constant 1 : index
+ // CHECK: Generated: %[[DIM1:.*]] = memref.dim %[[ORIG]], %[[C1]] : memref<?x?xf32>
+ // CHECK: Generated: %{{.*}} = memref.alloc(%[[DIM0]], %[[DIM1]]) : memref<?x?xf32>
+ return
+}
diff --git a/mlir/test/Dialect/OpenACC/pointer-like-interface-copy.mlir b/mlir/test/Dialect/OpenACC/pointer-like-interface-copy.mlir
new file mode 100644
index 0000000000000..9220d84a29f94
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/pointer-like-interface-copy.mlir
@@ -0,0 +1,23 @@
+// RUN: mlir-opt %s --split-input-file --pass-pipeline="builtin.module(func.func(test-acc-pointer-like-interface{test-mode=copy}))" 2>&1 | FileCheck %s
+
+func.func @test_copy_static() {
+ %src = memref.alloca() {test.src_ptr} : memref<10x20xf32>
+ %dest = memref.alloca() {test.dest_ptr} : memref<10x20xf32>
+
+ // CHECK: Successfully generated copy from source: %[[SRC:.*]] = memref.alloca() {test.src_ptr} : memref<10x20xf32> to destination: %[[DEST:.*]] = memref.alloca() {test.dest_ptr} : memref<10x20xf32>
+ // CHECK: Generated: memref.copy %[[SRC]], %[[DEST]] : memref<10x20xf32> to memref<10x20xf32>
+ return
+}
+
+// -----
+
+func.func @test_copy_dynamic() {
+ %c10 = arith.constant 10 : index
+ %c20 = arith.constant 20 : index
+ %src = memref.alloc(%c10, %c20) {test.src_ptr} : memref<?x?xf32>
+ %dest = memref.alloc(%c10, %c20) {test.dest_ptr} : memref<?x?xf32>
+
+ // CHECK: Successfully generated copy from source: %[[SRC:.*]] = memref.alloc(%[[C10:.*]], %[[C20:.*]]) {test.src_ptr} : memref<?x?xf32> to destination: %[[DEST:.*]] = memref.alloc(%[[C10]], %[[C20]]) {test.dest_ptr} : memref<?x?xf32>
+ // CHECK: Generated: memref.copy %[[SRC]], %[[DEST]] : memref<?x?xf32> to memref<?x?xf32>
+ return
+}
diff --git a/mlir/test/Dialect/OpenACC/pointer-like-interface-free.mlir b/mlir/test/Dialect/OpenACC/pointer-like-interface-free.mlir
new file mode 100644
index 0000000000000..ecf4f75fe06d1
--- /dev/null
+++ b/mlir/test/Dialect/OpenACC/pointer-like-interface-free.mlir
@@ -0,0 +1,31 @@
+// RUN: mlir-opt %s --split-input-file --pass-pipeline="builtin.module(func.func(test-acc-pointer-like-interface{test-mode=free}))" 2>&1 | FileCheck %s
+
+func.func @test_static_memref_free() {
+ %0 = memref.alloca() {test.ptr} : memref<10x20xf32>
+ // CHECK: Successfully generated free for operation: %[[ORIG:.*]] = memref.alloca() {test.ptr} : memref<10x20xf32>
+ // CHECK-NOT: Generated
+ return
+}
+
+// -----
+
+func.func @test_dynamic_memref_free() {
+ %c10 = arith.constant 10 : index
+ %c20 = arith.constant 20 : index
+ %orig = memref.alloc(%c10, %c20) {test.ptr} : memref<?x?xf32>
+
+ // CHECK: Successfully generated free for operation: %[[ORIG:.*]] = memref.alloc(%[[C10:.*]], %[[C20:.*]]) {test.ptr} : memref<?x?xf32>
+ // CHECK: Generated: memref.dealloc %[[ORIG]] : memref<?x?xf32>
+ return
+}
+
+// -----
+
+func.func @test_cast_walking_free() {
+ %0 = memref.alloca() : memref<10x20xf32>
+ %1 = memref.cast %0 {test.ptr} : memref<10x20xf32> to memref<?x?xf32>
+
+ // CHECK: Successfully generated free for operation: %[[CAST:.*]] = memref.cast %[[ALLOCA:.*]] {test.ptr} : memref<10x20xf32> to memref<?x?xf32>
+ // CHECK-NOT: Generated
+ return
+}
diff --git a/mlir/test/lib/Dialect/CMakeLists.txt b/mlir/test/lib/Dialect/CMakeLists.txt
index 3b7bd9b9637a8..e31140a6f6ba7 100644
--- a/mlir/test/lib/Dialect/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/CMakeLists.txt
@@ -12,6 +12,7 @@ add_subdirectory(Math)
add_subdirectory(MemRef)
add_subdirectory(Shard)
add_subdirectory(NVGPU)
+add_subdirectory(OpenACC)
add_subdirectory(SCF)
add_subdirectory(Shape)
add_subdirectory(SPIRV)
diff --git a/mlir/test/lib/Dialect/OpenACC/CMakeLists.txt b/mlir/test/lib/Dialect/OpenACC/CMakeLists.txt
new file mode 100644
index 0000000000000..f84055df1b6c6
--- /dev/null
+++ b/mlir/test/lib/Dialect/OpenACC/CMakeLists.txt
@@ -0,0 +1,16 @@
+add_mlir_library(MLIROpenACCTestPasses
+ TestOpenACC.cpp
+ TestPointerLikeTypeInterface.cpp
+
+ EXCLUDE_FROM_LIBMLIR
+)
+mlir_target_link_libraries(MLIROpenACCTestPasses PUBLIC
+ MLIRIR
+ MLIRArithDialect
+ MLIRFuncDialect
+ MLIRMemRefDialect
+ MLIROpenACCDialect
+ MLIRPass
+ MLIRSupport
+)
+
diff --git a/mlir/test/lib/Dialect/OpenACC/TestOpenACC.cpp b/mlir/test/lib/Dialect/OpenACC/TestOpenACC.cpp
new file mode 100644
index 0000000000000..6f7bdde176aa7
--- /dev/null
+++ b/mlir/test/lib/Dialect/OpenACC/TestOpenACC.cpp
@@ -0,0 +1,26 @@
+//===- TestOpenACC.cpp - OpenACC Test Registration ------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains unified registration for all OpenACC test passes.
+//
+//===----------------------------------------------------------------------===//
+
+namespace mlir {
+namespace test {
+
+// Forward declarations of individual test pass registration functions
+void registerTestPointerLikeTypeInterfacePass();
+
+// Unified registration function for all OpenACC tests
+void registerTestOpenACC() {
+ registerTestPointerLikeTypeInterfacePass();
+}
+
+} // namespace test
+} // namespace mlir
+
diff --git a/mlir/test/lib/Dialect/OpenACC/TestPointerLikeTypeInterface.cpp b/mlir/test/lib/Dialect/OpenACC/TestPointerLikeTypeInterface.cpp
new file mode 100644
index 0000000000000..fc573d4318b5b
--- /dev/null
+++ b/mlir/test/lib/Dialect/OpenACC/TestPointerLikeTypeInterface.cpp
@@ -0,0 +1,310 @@
+//===- TestPointerLikeTypeInterface.cpp - Test PointerLikeType interface -===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains test passes for testing the OpenACC PointerLikeType
+// interface methods.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/Support/CommandLine.h"
+
+using namespace mlir;
+using namespace mlir::acc;
+
+namespace {
+
+struct OperationTracker : public OpBuilder::Listener {
+ SmallVector<Operation *> insertedOps;
+
+ void notifyOperationInserted(Operation *op,
+ OpBuilder::InsertPoint previous) override {
+ insertedOps.push_back(op);
+ }
+};
+
+struct TestPointerLikeTypeInterfacePass
+ : public PassWrapper<TestPointerLikeTypeInterfacePass,
+ OperationPass<func::FuncOp>> {
+ MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestPointerLikeTypeInterfacePass)
+
+ TestPointerLikeTypeInterfacePass() = default;
+ TestPointerLikeTypeInterfacePass(const TestPointerLikeTypeInterfacePass &pass)
+ : PassWrapper(pass) {
+ testMode = pass.testMode;
+ }
+
+ Pass::Option<std::string> testMode{
+ *this, "test-mode",
+ llvm::cl::desc("Test mode: walk, alloc, copy, or free"),
+ llvm::cl::init("walk")};
+
+ StringRef getArgument() const override {
+ return "test-acc-pointer-like-interface";
+ }
+
+ StringRef getDescription() const override {
+ return "Test OpenACC PointerLikeType interface methods on any implementing "
+ "type";
+ }
+
+ void runOnOperation() override;
+
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<acc::OpenACCDialect>();
+ registry.insert<arith::ArithDialect>();
+ registry.insert<memref::MemRefDialect>();
+ }
+
+private:
+ void walkAndPrint();
+ void testGenAllocate(Operation *op, Value result, PointerLikeType pointerType,
+ OpBuilder &builder);
+ void testGenFree(Operation *op, Value result, PointerLikeType pointerType,
+ OpBuilder &builder);
+ void testGenCopy(Operation *srcOp, Operation *destOp, Value srcResult, Value destResult,
+ PointerLikeType pointerType, OpBuilder &builder);
+
+ struct PointerCandidate {
+ Operation *op;
+ Value result;
+ PointerLikeType pointerType;
+ };
+};
+
+void TestPointerLikeTypeInterfacePass::runOnOperation() {
+ if (testMode == "walk") {
+ walkAndPrint();
+ return;
+ }
+
+ auto func = getOperation();
+ OpBuilder builder(&getContext());
+
+ if (testMode == "alloc" || testMode == "free") {
+ // Collect all candidates first
+ SmallVector<PointerCandidate> candidates;
+ func.walk([&](Operation *op) {
+ if (op->hasAttr("test.ptr")) {
+ for (auto result : op->getResults()) {
+ if (isa<PointerLikeType>(result.getType())) {
+ candidates.push_back(
+ {op, result, cast<PointerLikeType>(result.getType())});
+ break; // Only take the first PointerLikeType result
+ }
+ }
+ }
+ });
+
+ // Now test all candidates
+ for (const auto &candidate : candidates) {
+ if (testMode == "alloc") {
+ testGenAllocate(candidate.op, candidate.result, candidate.pointerType,
+ builder);
+ } else if (testMode == "free") {
+ testGenFree(candidate.op, candidate.result, candidate.pointerType,
+ builder);
+ }
+ }
+ } else if (testMode == "copy") {
+ // Collect all source and destination candidates
+ SmallVector<PointerCandidate> sources, destinations;
+
+ func.walk([&](Operation *op) {
+ if (op->hasAttr("test.src_ptr")) {
+ for (auto resul...
[truncated]
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
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.
Looks great! Thank you, Razvan!
When generating recipes, the allocated value needs to be obtainable so that it can be threaded through to the copying and freeing regions. Thus instead of returning true/false, return value just for genAllocate.
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.
LGTM
InterfaceMethod< | ||
/*description=*/[{ | ||
Generates allocation operations for the pointer-like type. It will create | ||
an allocate that produces memory space for an instance of the current type. |
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.
an allocate that produces memory space for an instance of the current type. | |
an allocate operation that produces memory space for an instance of the current type. |
Seeing allocate
used as a noun is somehow jarring to me. Is this idiomatic in a way I'm just not familiar with?
Generates allocation operations for the pointer-like type. It will create | ||
an allocate that produces memory space for an instance of the current type. | ||
|
||
The `varName` parameter is optional and can be used to provide a name |
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.
Can this be ignored if it is present? I suppose the answer is yes, since your memref implementation does that. Maybe mention that in the description?
A variable in an acc data clause operation must have a type that implements either PointerLikeType or a MappableType interface. These interfaces provide the contract that allows acc dialect and its transform passes to interact with a source dialect.
One of these requirements is ability to generate code that creates memory for a private copy and ability to initialize that copy from another variable. Thus, update the PointerLikeType API to provide the means to create allocation, deallocation, and copy. This will be used as a way to fill in privatization and firstprivatization recipes.
This new API was implemented for memref along with testing to exercise it via the implementation of a testing pass.