[mlir][acc] Extend PointerLikeType to provide alloc, dealloc, copy

mlir/include/mlir/Dialect/OpenACC/OpenACCTypeInterfaces.td

@@ -70,6 +70,86 @@ 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 a Value representing the result of the allocation. If no value
+        is returned, it means the allocation was not successfully generated.
+      }],
+      /*retTy=*/"::mlir::Value",
+      /*methodName=*/"genAllocate",
+      /*args=*/(ins "::mlir::OpBuilder &":$builder,
+                    "::mlir::Location":$loc,
+                    "::llvm::StringRef":$varName,
+                    "::mlir::Type":$varType,
+                    "::mlir::Value":$originalVar),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return {};
+      }]
+    >,
+    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;
+      }]
+    >,
   ];
 }
 

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,115 @@ struct MemRefPointerLikeModel
     assert(memrefTy.getRank() > 0 && "rank expected to be positive");
     return mlir::acc::VariableTypeCategory::array;
   }
+
+  mlir::Value 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())
+      return memref::AllocaOp::create(builder, loc, memrefTy).getResult();
+
+    // 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
+      }
+      return memref::AllocOp::create(builder, loc, memrefTy, dynamicSizes)
+          .getResult();
+    }
+
+    // TODO: Unranked not yet supported.
+    return {};
+  }
+
+  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

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
+}

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
+}

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
+}

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)

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
+)
+

mlir/test/lib/Dialect/OpenACC/TestOpenACC.cpp

@@ -0,0 +1,23 @@
+//===- 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