Skip to content

Commit aa2da60

Browse files
krzysz00giacs-epic
krzysz00
authored and
giacs-epic
committed
[Codegen][LLVM] Annotate pointers an noundef + nonnull, consants as noundef (iree-org#19309)
We know, as a matter of the API, that pointers passed into GPU functions and the HAL VM entry point won't be null and will have a defined value. Annotating these facts may unblock LLVM optimizations. In particuar, when future commits will annotate ranges on the push constants, those annotations will need a `noundef` to be used. The existing `noalias` may also sometimes be turned off by the possibility of undefined arguments. Signed-off-by: Giacomo Serafini <[email protected]>
1 parent d34bf98 commit aa2da60

File tree

5 files changed

+39
-27
lines changed

5 files changed

+39
-27
lines changed

‎compiler/src/iree/compiler/Codegen/LLVMCPU/ConvertToLLVM.cpp‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -161,8 +161,10 @@ struct ConvertHALEntryPointFuncOp
161161
// can use the attributes.
162162
// (%arg0: environment, %arg1: dispatch_state, %arg2: workgroup_state)
163163
for (unsigned i = 0; i <= 2; ++i) {
164-
llvmFuncOp.setArgAttr(i, LLVM::LLVMDialect::getNoAliasAttrName(),
165-
rewriter.getUnitAttr());
164+
Attribute unit = rewriter.getUnitAttr();
165+
llvmFuncOp.setArgAttr(i, LLVM::LLVMDialect::getNoAliasAttrName(), unit);
166+
llvmFuncOp.setArgAttr(i, LLVM::LLVMDialect::getNonNullAttrName(), unit);
167+
llvmFuncOp.setArgAttr(i, LLVM::LLVMDialect::getNoUndefAttrName(), unit);
166168
llvmFuncOp.setArgAttr(i, LLVM::LLVMDialect::getAlignAttrName(),
167169
rewriter.getI64IntegerAttr(16));
168170
}

‎compiler/src/iree/compiler/Codegen/LLVMCPU/test/convert_to_llvm.mlir‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,9 @@ builtin.module {
88
}
99
// CHECK: llvm.func @extern_public()
1010
// CHECK: llvm.func @entry_point(
11-
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: !llvm.ptr {llvm.align = 16 : i64, llvm.noalias},
12-
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: !llvm.ptr {llvm.align = 16 : i64, llvm.noalias},
13-
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: !llvm.ptr {llvm.align = 16 : i64, llvm.noalias}) -> i32
11+
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: !llvm.ptr {llvm.align = 16 : i64, llvm.noalias, llvm.nonnull, llvm.noundef},
12+
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: !llvm.ptr {llvm.align = 16 : i64, llvm.noalias, llvm.nonnull, llvm.noundef},
13+
// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: !llvm.ptr {llvm.align = 16 : i64, llvm.noalias, llvm.nonnull, llvm.noundef}) -> i32
1414
// CHECK: llvm.return %{{.+}} : i32
1515

1616
// -----

‎compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.cpp‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -354,15 +354,18 @@ class ConvertIREEBindingSubspanOp : public ConvertToLLVMPattern {
354354
rewriter.getI32IntegerAttr(16));
355355
// It is safe to set the noalias attribute as it is guaranteed that the
356356
// ranges within bindings won't alias.
357+
Attribute unit = rewriter.getUnitAttr();
357358
llvmFuncOp.setArgAttr(llvmBufferArg.getArgNumber(),
358-
LLVM::LLVMDialect::getNoAliasAttrName(),
359-
rewriter.getUnitAttr());
359+
LLVM::LLVMDialect::getNoAliasAttrName(), unit);
360+
llvmFuncOp.setArgAttr(llvmBufferArg.getArgNumber(),
361+
LLVM::LLVMDialect::getNonNullAttrName(), unit);
362+
llvmFuncOp.setArgAttr(llvmBufferArg.getArgNumber(),
363+
LLVM::LLVMDialect::getNoUndefAttrName(), unit);
360364
if (checkAllSubspansReadonly(llvmFuncOp, subspanOp.getBinding())) {
361365
// Setting the readonly attribute here will generate non-coherent cache
362366
// loads.
363367
llvmFuncOp.setArgAttr(llvmBufferArg.getArgNumber(),
364-
LLVM::LLVMDialect::getReadonlyAttrName(),
365-
rewriter.getUnitAttr());
368+
LLVM::LLVMDialect::getReadonlyAttrName(), unit);
366369
}
367370
// Add the byte offset.
368371
Value llvmBufferBasePtr = llvmBufferArg;
@@ -471,6 +474,13 @@ class ConvertIREEConstantOp : public ConvertToLLVMPattern {
471474
mlir::BlockArgument llvmBufferArg = llvmFuncOp.getArgument(
472475
argMapping.size() + ireeConstantOp.getOrdinal().getZExtValue());
473476
assert(llvmBufferArg.getType().isInteger(32));
477+
478+
// Push constants are never `undef`, annotate that here, just as with
479+
// bindings.
480+
llvmFuncOp.setArgAttr(llvmBufferArg.getArgNumber(),
481+
LLVM::LLVMDialect::getNoUndefAttrName(),
482+
rewriter.getUnitAttr());
483+
474484
Type dstType = getTypeConverter()->convertType(ireeConstantOp.getType());
475485
// llvm.zext requires that the result type has a larger bitwidth.
476486
if (dstType == llvmBufferArg.getType()) {

‎compiler/src/iree/compiler/Codegen/LLVMGPU/test/convert_to_nvvm.mlir‎

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -32,9 +32,9 @@ hal.executable @abs_ex_dispatch_0 {
3232
}
3333
}
3434
// CHECK-LABEL: llvm.func @abs_ex_dispatch_0
35-
// CHECK-SAME: (%[[ARG0:.+]]: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias},
36-
// CHECK-SAME: %[[ARG1:.+]]: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias, llvm.readonly},
37-
// CHECK-SAME: %[[ARG2:.+]]: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias})
35+
// CHECK-SAME: (%[[ARG0:.+]]: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias, llvm.nonnull, llvm.noundef},
36+
// CHECK-SAME: %[[ARG1:.+]]: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias, llvm.nonnull, llvm.noundef, llvm.readonly},
37+
// CHECK-SAME: %[[ARG2:.+]]: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias, llvm.nonnull, llvm.noundef})
3838
// CHECK: %[[FADD:.+]] = llvm.fadd %{{.*}}, %{{.*}} : f32
3939
// CHECK: %[[ADDR:.+]] = llvm.getelementptr %[[ARG2]][%{{.*}}] : (!llvm.ptr, i64) -> !llvm.ptr, f32
4040
// CHECK: llvm.store %[[FADD]], %[[ADDR]] : f32, !llvm.ptr
@@ -72,13 +72,13 @@ hal.executable @abs_dynamic {
7272
}
7373
}
7474
// CHECK-LABEL: llvm.func @abs_dynamic
75-
// CHECK-SAME: (%[[ARG0:[a-zA-Z0-9]+]]: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias},
76-
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias},
77-
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias},
78-
// CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]: i32,
79-
// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: i32,
80-
// CHECK-SAME: %[[ARG5:[a-zA-Z0-9]+]]: i32,
81-
// CHECK-SAME: %[[ARG6:[a-zA-Z0-9]+]]: i32)
75+
// CHECK-SAME: (%[[ARG0:[a-zA-Z0-9]+]]: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias, llvm.nonnull, llvm.noundef},
76+
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias, llvm.nonnull, llvm.noundef},
77+
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias, llvm.nonnull, llvm.noundef},
78+
// CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]: i32 {llvm.noundef},
79+
// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: i32 {llvm.noundef},
80+
// CHECK-SAME: %[[ARG5:[a-zA-Z0-9]+]]: i32 {llvm.noundef},
81+
// CHECK-SAME: %[[ARG6:[a-zA-Z0-9]+]]: i32 {llvm.noundef})
8282
// CHECK-DAG: %[[OFFSET:.+]] = llvm.zext %[[ARG3]] : i32 to i64
8383
// CHECK-DAG: %[[D1:.+]] = llvm.zext %[[ARG5]] : i32 to i64
8484
// CHECK-DAG: %[[D2:.+]] = llvm.zext %[[ARG6]] : i32 to i64
@@ -126,8 +126,8 @@ hal.executable @dead_symbol {
126126
}
127127
}
128128
// CHECK-LABEL: llvm.func @dead_symbol
129-
// CHECK-SAME: (%[[ARG0:.+]]: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias},
130-
// CHECK-SAME: %[[ARG1:.+]]: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias})
129+
// CHECK-SAME: (%[[ARG0:.+]]: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias, llvm.nonnull, llvm.noundef},
130+
// CHECK-SAME: %[[ARG1:.+]]: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias, llvm.nonnull, llvm.noundef})
131131
// CHECK: llvm.fadd
132132

133133
// -----
@@ -165,8 +165,8 @@ hal.executable @mixed_type {
165165
}
166166

167167
// CHECK-LABEL: llvm.func @mixed_type
168-
// CHECK-SAME: (%[[ARG0:.+]]: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias},
169-
// CHECK-SAME: %{{.*}}: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias})
168+
// CHECK-SAME: (%[[ARG0:.+]]: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias, llvm.nonnull, llvm.noundef},
169+
// CHECK-SAME: %{{.*}}: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias, llvm.nonnull, llvm.noundef})
170170
// CHECK: nvvm.read.ptx.sreg.tid.x
171171
// CHECK: llvm.getelementptr %[[ARG0]][4] : (!llvm.ptr) -> !llvm.ptr, f32
172172
// CHECK: llvm.fadd
@@ -304,7 +304,7 @@ hal.executable @check_not_readonly {
304304
}
305305
}
306306
// CHECK-LABEL: llvm.func @check_not_readonly
307-
// CHECK-NOT: (%[[ARG0:.+]]: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias, llvm.readonly},
307+
// CHECK-NOT: (%[[ARG0:.+]]: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias, llvm.nonnull, llvm.noundef, llvm.readonly},
308308

309309
// -----
310310

‎compiler/src/iree/compiler/Codegen/LLVMGPU/test/convert_to_rocdl.mlir‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -33,9 +33,9 @@ hal.executable @abs_ex_dispatch_0 {
3333
}
3434
// CHECK-LABEL: llvm.func @abs_ex_dispatch_0
3535
// INDEX32-LABEL: llvm.func @abs_ex_dispatch_0
36-
// CHECK-SAME: (%{{[a-zA-Z0-9]*}}: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias, llvm.readonly},
37-
// CHECK-SAME: %{{[a-zA-Z0-9]*}}: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias},
38-
// CHECK-SAME: %{{[a-zA-Z0-9]*}}: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias})
36+
// CHECK-SAME: (%{{[a-zA-Z0-9]*}}: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias, llvm.nonnull, llvm.noundef, llvm.readonly},
37+
// CHECK-SAME: %{{[a-zA-Z0-9]*}}: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias, llvm.nonnull, llvm.noundef},
38+
// CHECK-SAME: %{{[a-zA-Z0-9]*}}: !llvm.ptr {llvm.align = 16 : i32, llvm.noalias, llvm.nonnull, llvm.noundef})
3939
// CHECK: rocdl.workgroup.dim.x
4040
// CHECK: llvm.getelementptr %{{.*}} : (!llvm.ptr, i64) -> !llvm.ptr, f32
4141
// INDEX32: llvm.getelementptr %{{.*}} : (!llvm.ptr, i32) -> !llvm.ptr, f32

0 commit comments

Comments
 (0)