[Mlir-commits] [mlir] [mlir] fix overflow in LLVM dialect inlining (PR #72878)
Oleksandr Alex Zinenko
llvmlistbot at llvm.org
Tue Nov 21 07:23:13 PST 2023
https://github.com/ftynse updated https://github.com/llvm/llvm-project/pull/72878
>From 8f8b079d785ebb2546b2463aea3b69ed636af304 Mon Sep 17 00:00:00 2001
From: Alex Zinenko <zinenko at google.com>
Date: Mon, 20 Nov 2023 12:36:00 +0000
Subject: [PATCH 1/2] [mlir] fix overflow in LLVM dialect inlining
Don't use unsigned for sizes as they may be larger than that type can
hold.
Fixes #72678.
---
mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp | 14 +++++++-------
.../Dialect/LLVMIR/inline-byval-huge.mlir | 19 ++++++++++++++++++-
2 files changed, 25 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
index 2abb9b0e3986872..bfa8bc172969756 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
@@ -557,8 +557,8 @@ static unsigned tryToEnforceAlignment(Value value, unsigned requestedAlignment,
/// the address of the new alloca, then returns the value of the new alloca.
static Value handleByValArgumentInit(OpBuilder &builder, Location loc,
Value argument, Type elementType,
- unsigned elementTypeSize,
- unsigned targetAlignment) {
+ uint64_t elementTypeSize,
+ uint64_t targetAlignment) {
// Allocate the new value on the stack.
Value allocaOp;
{
@@ -587,7 +587,7 @@ static Value handleByValArgumentInit(OpBuilder &builder, Location loc,
/// attribute (or 1 if no align attribute was set).
static Value handleByValArgument(OpBuilder &builder, Operation *callable,
Value argument, Type elementType,
- unsigned requestedAlignment) {
+ uint64_t requestedAlignment) {
auto func = cast<LLVM::LLVMFuncOp>(callable);
LLVM::MemoryEffectsAttr memoryEffects = func.getMemoryAttr();
// If there is no memory effects attribute, assume that the function is
@@ -597,16 +597,16 @@ static Value handleByValArgument(OpBuilder &builder, Operation *callable,
memoryEffects.getArgMem() != LLVM::ModRefInfo::Mod;
// Check if there's an alignment mismatch requiring us to copy.
DataLayout dataLayout = DataLayout::closest(callable);
- unsigned minimumAlignment = dataLayout.getTypeABIAlignment(elementType);
+ uint64_t minimumAlignment = dataLayout.getTypeABIAlignment(elementType);
if (isReadOnly) {
if (requestedAlignment <= minimumAlignment)
return argument;
- unsigned currentAlignment =
+ uint64_t currentAlignment =
tryToEnforceAlignment(argument, requestedAlignment, dataLayout);
if (currentAlignment >= requestedAlignment)
return argument;
}
- unsigned targetAlignment = std::max(requestedAlignment, minimumAlignment);
+ uint64_t targetAlignment = std::max(requestedAlignment, minimumAlignment);
return handleByValArgumentInit(builder, func.getLoc(), argument, elementType,
dataLayout.getTypeSize(elementType),
targetAlignment);
@@ -753,7 +753,7 @@ struct LLVMInlinerInterface : public DialectInlinerInterface {
if (std::optional<NamedAttribute> attr =
argumentAttrs.getNamed(LLVM::LLVMDialect::getByValAttrName())) {
Type elementType = cast<TypeAttr>(attr->getValue()).getValue();
- unsigned requestedAlignment = 1;
+ uint64_t requestedAlignment = 1;
if (std::optional<NamedAttribute> alignAttr =
argumentAttrs.getNamed(LLVM::LLVMDialect::getAlignAttrName())) {
requestedAlignment = cast<IntegerAttr>(alignAttr->getValue())
diff --git a/mlir/test/Dialect/LLVMIR/inline-byval-huge.mlir b/mlir/test/Dialect/LLVMIR/inline-byval-huge.mlir
index 3e246d7f21d5f88..ef39b96b2736617 100644
--- a/mlir/test/Dialect/LLVMIR/inline-byval-huge.mlir
+++ b/mlir/test/Dialect/LLVMIR/inline-byval-huge.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt -inline %s | FileCheck %s
+// RUN: mlir-opt -inline --split-input-file %s | FileCheck %s
// CHECK-LABEL: @byval_2_000_000_000
llvm.func @byval_2_000_000_000(%ptr : !llvm.ptr) {
@@ -12,3 +12,20 @@ llvm.func @byval_2_000_000_000(%ptr : !llvm.ptr) {
llvm.func @with_byval_arg(%ptr : !llvm.ptr { llvm.byval = !llvm.array<1000 x array<1000 x array <500 x i32>>> }) {
llvm.return
}
+
+// -----
+
+// This exceeds representational capacity of 32-bit unsigned value.
+
+// CHECK-LABEL: @byval_8_000_000_000
+llvm.func @byval_8_000_000_000(%ptr : !llvm.ptr) {
+ // CHECK: %[[SIZE:.+]] = llvm.mlir.constant(8000000000 : i64)
+ // CHECK: "llvm.intr.memcpy"(%{{.*}}, %{{.*}}, %[[SIZE]]
+
+ llvm.call @with_byval_arg(%ptr) : (!llvm.ptr) -> ()
+ llvm.return
+}
+
+llvm.func @with_byval_arg(%ptr : !llvm.ptr { llvm.byval = !llvm.array<2000 x array<2000 x array <500 x i32>>> }) {
+ llvm.return
+}
>From 86817e90e65389de90529254c4a37881aa5512c4 Mon Sep 17 00:00:00 2001
From: Alex Zinenko <zinenko at google.com>
Date: Tue, 21 Nov 2023 13:11:23 +0000
Subject: [PATCH 2/2] more uint64_t
---
mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp | 10 +++++-----
1 file changed, 5 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
index bfa8bc172969756..d9f52d7692938f0 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMInlining.cpp
@@ -496,14 +496,14 @@ static void handleAccessGroups(Operation *call,
/// If `requestedAlignment` is higher than the alignment specified on `alloca`,
/// realigns `alloca` if this does not exceed the natural stack alignment.
/// Returns the post-alignment of `alloca`, whether it was realigned or not.
-static unsigned tryToEnforceAllocaAlignment(LLVM::AllocaOp alloca,
- unsigned requestedAlignment,
+static uint64_t tryToEnforceAllocaAlignment(LLVM::AllocaOp alloca,
+ uint64_t requestedAlignment,
DataLayout const &dataLayout) {
- unsigned allocaAlignment = alloca.getAlignment().value_or(1);
+ uint64_t allocaAlignment = alloca.getAlignment().value_or(1);
if (requestedAlignment <= allocaAlignment)
// No realignment necessary.
return allocaAlignment;
- unsigned naturalStackAlignmentBits = dataLayout.getStackAlignment();
+ uint64_t naturalStackAlignmentBits = dataLayout.getStackAlignment();
// If the natural stack alignment is not specified, the data layout returns
// zero. Optimistically allow realignment in this case.
if (naturalStackAlignmentBits == 0 ||
@@ -525,7 +525,7 @@ static unsigned tryToEnforceAllocaAlignment(LLVM::AllocaOp alloca,
/// the pointer, then returns the resulting post-alignment, regardless of
/// whether it was realigned or not. If no existing alignment attribute is
/// found, returns 1 (i.e., assume that no alignment is guaranteed).
-static unsigned tryToEnforceAlignment(Value value, unsigned requestedAlignment,
+static uint64_t tryToEnforceAlignment(Value value, uint64_t requestedAlignment,
DataLayout const &dataLayout) {
if (Operation *definingOp = value.getDefiningOp()) {
if (auto alloca = dyn_cast<LLVM::AllocaOp>(definingOp))
More information about the Mlir-commits
mailing list