Skip to content

Commit fdbd75d

Browse files
razvanlupusorurazvanlupusoru
authored
[mlir][acc] Ensure genAllocate uses provided variable name (#163861)
The genAllocate API was documented to have the `varName` argument as optional. However, when it is provided, it becomes unexpected if the implementation does not use it. Since not all dialects have a way to store variable names, add one in the acc dialect and use it to store names of memref variables. This updates the API documentation, implementation of genAllocate for memref, and IR testing.
1 parent 7466915 commit fdbd75d

File tree

7 files changed

+50
-19
lines changed

7 files changed

+50
-19
lines changed

‎mlir/include/mlir/Dialect/OpenACC/OpenACC.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -183,6 +183,10 @@ static constexpr StringLiteral getRoutineInfoAttrName() {
183183
return StringLiteral("acc.routine_info");
184184
}
185185

186+
static constexpr StringLiteral getVarNameAttrName() {
187+
return VarNameAttr::name;
188+
}
189+
186190
static constexpr StringLiteral getCombinedConstructsAttrName() {
187191
return CombinedConstructsTypeAttr::name;
188192
}

‎mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -415,6 +415,13 @@ def OpenACC_ConstructResource : Resource<"::mlir::acc::ConstructResource">;
415415
// Define a resource for the OpenACC current device setting.
416416
def OpenACC_CurrentDeviceIdResource : Resource<"::mlir::acc::CurrentDeviceIdResource">;
417417

418+
// Attribute for saving variable names - this can be attached to non-acc-dialect
419+
// operations in order to ensure the name is preserved.
420+
def OpenACC_VarNameAttr : OpenACC_Attr<"VarName", "var_name"> {
421+
let parameters = (ins StringRefParameter<"">:$name);
422+
let assemblyFormat = "`<` $name `>`";
423+
}
424+
418425
// Used for data specification in data clauses (2.7.1).
419426
// Either (or both) extent and upperbound must be specified.
420427
def OpenACC_DataBoundsOp : OpenACC_Op<"bounds",

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

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -73,12 +73,18 @@ def OpenACC_PointerLikeTypeInterface : TypeInterface<"PointerLikeType"> {
7373
InterfaceMethod<
7474
/*description=*/[{
7575
Generates allocation operations for the pointer-like type. It will create
76-
an allocate that produces memory space for an instance of the current type.
76+
an allocate operation that produces memory space for an instance of the
77+
current type.
7778

7879
The `varName` parameter is optional and can be used to provide a name
79-
for the allocated variable. If the current type is represented
80-
in a way that it does not capture the pointee type, `varType` must be
81-
passed in to provide the necessary type information.
80+
for the allocated variable. When provided, it must be used by the
81+
implementation; and if the implementing dialect does not have its own
82+
way to save it, the discardable `acc.var_name` attribute from the acc
83+
dialect will be used.
84+
85+
If the current type is represented in a way that it does not capture
86+
the pointee type, `varType` must be passed in to provide the necessary
87+
type information.
8288

8389
The `originalVar` parameter is optional but enables support for dynamic
8490
types (e.g., dynamic memrefs). When provided, implementations can extract

‎mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,16 @@ static bool isScalarLikeType(Type type) {
4040
return type.isIntOrIndexOrFloat() || isa<ComplexType>(type);
4141
}
4242

43+
/// Helper function to attach the `VarName` attribute to an operation
44+
/// if a variable name is provided.
45+
static void attachVarNameAttr(Operation *op, OpBuilder &builder,
46+
StringRef varName) {
47+
if (!varName.empty()) {
48+
auto varNameAttr = acc::VarNameAttr::get(builder.getContext(), varName);
49+
op->setAttr(acc::getVarNameAttrName(), varNameAttr);
50+
}
51+
}
52+
4353
struct MemRefPointerLikeModel
4454
: public PointerLikeType::ExternalModel<MemRefPointerLikeModel,
4555
MemRefType> {
@@ -83,7 +93,9 @@ struct MemRefPointerLikeModel
8393
// then we can generate an alloca operation.
8494
if (memrefTy.hasStaticShape()) {
8595
needsFree = false; // alloca doesn't need deallocation
86-
return memref::AllocaOp::create(builder, loc, memrefTy).getResult();
96+
auto allocaOp = memref::AllocaOp::create(builder, loc, memrefTy);
97+
attachVarNameAttr(allocaOp, builder, varName);
98+
return allocaOp.getResult();
8799
}
88100

89101
// For dynamic memrefs, extract sizes from the original variable if
@@ -103,8 +115,10 @@ struct MemRefPointerLikeModel
103115
// Static dimensions are handled automatically by AllocOp
104116
}
105117
needsFree = true; // alloc needs deallocation
106-
return memref::AllocOp::create(builder, loc, memrefTy, dynamicSizes)
107-
.getResult();
118+
auto allocOp =
119+
memref::AllocOp::create(builder, loc, memrefTy, dynamicSizes);
120+
attachVarNameAttr(allocOp, builder, varName);
121+
return allocOp.getResult();
108122
}
109123

110124
// TODO: Unranked not yet supported.

‎mlir/test/Dialect/OpenACC/pointer-like-interface-alloc.mlir‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
func.func @test_static_memref_alloc() {
44
%0 = memref.alloca() {test.ptr} : memref<10x20xf32>
55
// CHECK: Successfully generated alloc for operation: %[[ORIG:.*]] = memref.alloca() {test.ptr} : memref<10x20xf32>
6-
// CHECK: Generated: %{{.*}} = memref.alloca() : memref<10x20xf32>
6+
// CHECK: Generated: %{{.*}} = memref.alloca() {acc.var_name = #acc.var_name<"test_alloc">} : memref<10x20xf32>
77
return
88
}
99

@@ -19,6 +19,6 @@ func.func @test_dynamic_memref_alloc() {
1919
// CHECK: Generated: %[[DIM0:.*]] = memref.dim %[[ORIG]], %[[C0]] : memref<?x?xf32>
2020
// CHECK: Generated: %[[C1:.*]] = arith.constant 1 : index
2121
// CHECK: Generated: %[[DIM1:.*]] = memref.dim %[[ORIG]], %[[C1]] : memref<?x?xf32>
22-
// CHECK: Generated: %{{.*}} = memref.alloc(%[[DIM0]], %[[DIM1]]) : memref<?x?xf32>
22+
// CHECK: Generated: %{{.*}} = memref.alloc(%[[DIM0]], %[[DIM1]]) {acc.var_name = #acc.var_name<"test_alloc">} : memref<?x?xf32>
2323
return
2424
}

‎mlir/test/Dialect/OpenACC/recipe-populate-firstprivate.mlir‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
// CHECK: acc.firstprivate.recipe @firstprivate_scalar : memref<f32> init {
44
// CHECK: ^bb0(%{{.*}}: memref<f32>):
5-
// CHECK: %[[ALLOC:.*]] = memref.alloca() : memref<f32>
5+
// CHECK: %[[ALLOC:.*]] = memref.alloca() {acc.var_name = #acc.var_name<"scalar">} : memref<f32>
66
// CHECK: acc.yield %[[ALLOC]] : memref<f32>
77
// CHECK: } copy {
88
// CHECK: ^bb0(%[[SRC:.*]]: memref<f32>, %[[DST:.*]]: memref<f32>):
@@ -20,7 +20,7 @@ func.func @test_scalar() {
2020

2121
// CHECK: acc.firstprivate.recipe @firstprivate_static_2d : memref<10x20xf32> init {
2222
// CHECK: ^bb0(%{{.*}}: memref<10x20xf32>):
23-
// CHECK: %[[ALLOC:.*]] = memref.alloca() : memref<10x20xf32>
23+
// CHECK: %[[ALLOC:.*]] = memref.alloca() {acc.var_name = #acc.var_name<"static_2d">} : memref<10x20xf32>
2424
// CHECK: acc.yield %[[ALLOC]] : memref<10x20xf32>
2525
// CHECK: } copy {
2626
// CHECK: ^bb0(%[[SRC:.*]]: memref<10x20xf32>, %[[DST:.*]]: memref<10x20xf32>):
@@ -42,7 +42,7 @@ func.func @test_static_2d() {
4242
// CHECK: %[[DIM0:.*]] = memref.dim %[[ARG]], %[[C0]] : memref<?x?xf32>
4343
// CHECK: %[[C1:.*]] = arith.constant 1 : index
4444
// CHECK: %[[DIM1:.*]] = memref.dim %[[ARG]], %[[C1]] : memref<?x?xf32>
45-
// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[DIM0]], %[[DIM1]]) : memref<?x?xf32>
45+
// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[DIM0]], %[[DIM1]]) {acc.var_name = #acc.var_name<"dynamic_2d">} : memref<?x?xf32>
4646
// CHECK: acc.yield %[[ALLOC]] : memref<?x?xf32>
4747
// CHECK: } copy {
4848
// CHECK: ^bb0(%[[SRC:.*]]: memref<?x?xf32>, %[[DST:.*]]: memref<?x?xf32>):
@@ -65,7 +65,7 @@ func.func @test_dynamic_2d(%arg0: index, %arg1: index) {
6565
// CHECK: ^bb0(%[[ARG:.*]]: memref<10x?xf32>):
6666
// CHECK: %[[C1:.*]] = arith.constant 1 : index
6767
// CHECK: %[[DIM1:.*]] = memref.dim %[[ARG]], %[[C1]] : memref<10x?xf32>
68-
// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[DIM1]]) : memref<10x?xf32>
68+
// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[DIM1]]) {acc.var_name = #acc.var_name<"mixed_dims">} : memref<10x?xf32>
6969
// CHECK: acc.yield %[[ALLOC]] : memref<10x?xf32>
7070
// CHECK: } copy {
7171
// CHECK: ^bb0(%[[SRC:.*]]: memref<10x?xf32>, %[[DST:.*]]: memref<10x?xf32>):
@@ -86,7 +86,7 @@ func.func @test_mixed_dims(%arg0: index) {
8686

8787
// CHECK: acc.firstprivate.recipe @firstprivate_scalar_int : memref<i32> init {
8888
// CHECK: ^bb0(%{{.*}}: memref<i32>):
89-
// CHECK: %[[ALLOC:.*]] = memref.alloca() : memref<i32>
89+
// CHECK: %[[ALLOC:.*]] = memref.alloca() {acc.var_name = #acc.var_name<"scalar_int">} : memref<i32>
9090
// CHECK: acc.yield %[[ALLOC]] : memref<i32>
9191
// CHECK: } copy {
9292
// CHECK: ^bb0(%[[SRC:.*]]: memref<i32>, %[[DST:.*]]: memref<i32>):

‎mlir/test/Dialect/OpenACC/recipe-populate-private.mlir‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
// CHECK: acc.private.recipe @private_scalar : memref<f32> init {
44
// CHECK: ^bb0(%{{.*}}: memref<f32>):
5-
// CHECK: %[[ALLOC:.*]] = memref.alloca() : memref<f32>
5+
// CHECK: %[[ALLOC:.*]] = memref.alloca() {acc.var_name = #acc.var_name<"scalar">} : memref<f32>
66
// CHECK: acc.yield %[[ALLOC]] : memref<f32>
77
// CHECK: }
88
// CHECK-NOT: destroy
@@ -16,7 +16,7 @@ func.func @test_scalar() {
1616

1717
// CHECK: acc.private.recipe @private_static_2d : memref<10x20xf32> init {
1818
// CHECK: ^bb0(%{{.*}}: memref<10x20xf32>):
19-
// CHECK: %[[ALLOC:.*]] = memref.alloca() : memref<10x20xf32>
19+
// CHECK: %[[ALLOC:.*]] = memref.alloca() {acc.var_name = #acc.var_name<"static_2d">} : memref<10x20xf32>
2020
// CHECK: acc.yield %[[ALLOC]] : memref<10x20xf32>
2121
// CHECK: }
2222
// CHECK-NOT: destroy
@@ -34,7 +34,7 @@ func.func @test_static_2d() {
3434
// CHECK: %[[DIM0:.*]] = memref.dim %[[ARG]], %[[C0]] : memref<?x?xf32>
3535
// CHECK: %[[C1:.*]] = arith.constant 1 : index
3636
// CHECK: %[[DIM1:.*]] = memref.dim %[[ARG]], %[[C1]] : memref<?x?xf32>
37-
// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[DIM0]], %[[DIM1]]) : memref<?x?xf32>
37+
// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[DIM0]], %[[DIM1]]) {acc.var_name = #acc.var_name<"dynamic_2d">} : memref<?x?xf32>
3838
// CHECK: acc.yield %[[ALLOC]] : memref<?x?xf32>
3939
// CHECK: } destroy {
4040
// CHECK: ^bb0(%{{.*}}: memref<?x?xf32>, %[[VAL:.*]]: memref<?x?xf32>):
@@ -53,7 +53,7 @@ func.func @test_dynamic_2d(%arg0: index, %arg1: index) {
5353
// CHECK: ^bb0(%[[ARG:.*]]: memref<10x?xf32>):
5454
// CHECK: %[[C1:.*]] = arith.constant 1 : index
5555
// CHECK: %[[DIM1:.*]] = memref.dim %[[ARG]], %[[C1]] : memref<10x?xf32>
56-
// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[DIM1]]) : memref<10x?xf32>
56+
// CHECK: %[[ALLOC:.*]] = memref.alloc(%[[DIM1]]) {acc.var_name = #acc.var_name<"mixed_dims">} : memref<10x?xf32>
5757
// CHECK: acc.yield %[[ALLOC]] : memref<10x?xf32>
5858
// CHECK: } destroy {
5959
// CHECK: ^bb0(%{{.*}}: memref<10x?xf32>, %[[VAL:.*]]: memref<10x?xf32>):
@@ -70,7 +70,7 @@ func.func @test_mixed_dims(%arg0: index) {
7070

7171
// CHECK: acc.private.recipe @private_scalar_int : memref<i32> init {
7272
// CHECK: ^bb0(%{{.*}}: memref<i32>):
73-
// CHECK: %[[ALLOC:.*]] = memref.alloca() : memref<i32>
73+
// CHECK: %[[ALLOC:.*]] = memref.alloca() {acc.var_name = #acc.var_name<"scalar_int">} : memref<i32>
7474
// CHECK: acc.yield %[[ALLOC]] : memref<i32>
7575
// CHECK: }
7676
// CHECK-NOT: destroy

0 commit comments

Comments
 (0)