Skip to content

[MLIR][LLVM] Handle floats in Mem2Reg of memset intrinsics #131621

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.

Sign up for GitHub

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

Merged
merged 1 commit into from
Mar 17, 2025
+90 −44
Merged

[MLIR][LLVM] Handle floats in Mem2Reg of memset intrinsics #131621

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 47 additions & 26 deletions mlir/lib/Dialect/LLVMIR/IR/LLVMMemorySlot.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1051,30 +1051,52 @@ static bool memsetCanRewire(MemsetIntr op, const DestructurableMemorySlot &slot,
template <class MemsetIntr>
static Value memsetGetStored(MemsetIntr op, const MemorySlot &slot,
OpBuilder &builder) {
// TODO: Support non-integer types.
return TypeSwitch<Type, Value>(slot.elemType)
.Case([&](IntegerType intType) -> Value {
if (intType.getWidth() == 8)
return op.getVal();

assert(intType.getWidth() % 8 == 0);

// Build the memset integer by repeatedly shifting the value and
// or-ing it with the previous value.
uint64_t coveredBits = 8;
Value currentValue =
builder.create<LLVM::ZExtOp>(op.getLoc(), intType, op.getVal());
while (coveredBits < intType.getWidth()) {
Value shiftBy = builder.create<LLVM::ConstantOp>(op.getLoc(), intType,
coveredBits);
Value shifted =
builder.create<LLVM::ShlOp>(op.getLoc(), currentValue, shiftBy);
currentValue =
builder.create<LLVM::OrOp>(op.getLoc(), currentValue, shifted);
coveredBits *= 2;
}
/// Returns an integer value that is `width` bits wide representing the value
/// assigned to the slot by memset.
auto buildMemsetValue = [&](unsigned width) -> Value {
assert(width % 8 == 0);
auto intType = IntegerType::get(op.getContext(), width);

// If we know the pattern at compile time, we can compute and assign a
// constant directly.
IntegerAttr constantPattern;
if (matchPattern(op.getVal(), m_Constant(&constantPattern))) {
assert(constantPattern.getValue().getBitWidth() == 8);
APInt memsetVal(/*numBits=*/width, /*val=*/0);
for (unsigned loBit = 0; loBit < width; loBit += 8)
memsetVal.insertBits(constantPattern.getValue(), loBit);
return builder.create<LLVM::ConstantOp>(
op.getLoc(), IntegerAttr::get(intType, memsetVal));
}

// If the output is a single byte, we can return the pattern directly.
if (width == 8)
return op.getVal();

// Otherwise build the memset integer at runtime by repeatedly shifting the
// value and or-ing it with the previous value.
uint64_t coveredBits = 8;
Value currentValue =
builder.create<LLVM::ZExtOp>(op.getLoc(), intType, op.getVal());
while (coveredBits < width) {
Value shiftBy =
builder.create<LLVM::ConstantOp>(op.getLoc(), intType, coveredBits);
Value shifted =
builder.create<LLVM::ShlOp>(op.getLoc(), currentValue, shiftBy);
currentValue =
builder.create<LLVM::OrOp>(op.getLoc(), currentValue, shifted);
coveredBits *= 2;
}

return currentValue;
return currentValue;
};
return TypeSwitch<Type, Value>(slot.elemType)
.Case([&](IntegerType type) -> Value {
return buildMemsetValue(type.getWidth());
})
.Case([&](FloatType type) -> Value {
Value intVal = buildMemsetValue(type.getWidth());
return builder.create<LLVM::BitcastOp>(op.getLoc(), type, intVal);
})
.Default([](Type) -> Value {
llvm_unreachable(
Expand All @@ -1088,11 +1110,10 @@ memsetCanUsesBeRemoved(MemsetIntr op, const MemorySlot &slot,
const SmallPtrSetImpl<OpOperand *> &blockingUses,
SmallVectorImpl<OpOperand *> &newBlockingUses,
const DataLayout &dataLayout) {
// TODO: Support non-integer types.
bool canConvertType =
TypeSwitch<Type, bool>(slot.elemType)
.Case([](IntegerType intType) {
return intType.getWidth() % 8 == 0 && intType.getWidth() > 0;
.Case<IntegerType, FloatType>([](auto type) {
return type.getWidth() % 8 == 0 && type.getWidth() > 0;
})
.Default([](Type) { return false; });
if (!canConvertType)
Expand Down
61 changes: 43 additions & 18 deletions mlir/test/Dialect/LLVMIR/mem2reg-intrinsics.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,30 @@ llvm.func @basic_memset(%memset_value: i8) -> i32 {

// -----

// CHECK-LABEL: llvm.func @memset_float
// CHECK-SAME: (%[[MEMSET_VALUE:.*]]: i8)
llvm.func @memset_float(%memset_value: i8) -> f32 {
%one = llvm.mlir.constant(1 : i32) : i32
%alloca = llvm.alloca %one x i32 {alignment = 4 : i64} : (i32) -> !llvm.ptr
%memset_len = llvm.mlir.constant(4 : i32) : i32
"llvm.intr.memset"(%alloca, %memset_value, %memset_len) <{isVolatile = false}> : (!llvm.ptr, i8, i32) -> ()
// CHECK-NOT: "llvm.intr.memset"
// CHECK: %[[VALUE_8:.*]] = llvm.zext %[[MEMSET_VALUE]] : i8 to i32
// CHECK: %[[C8:.*]] = llvm.mlir.constant(8 : i32) : i32
// CHECK: %[[SHIFTED_8:.*]] = llvm.shl %[[VALUE_8]], %[[C8]]
// CHECK: %[[VALUE_16:.*]] = llvm.or %[[VALUE_8]], %[[SHIFTED_8]]
// CHECK: %[[C16:.*]] = llvm.mlir.constant(16 : i32) : i32
// CHECK: %[[SHIFTED_16:.*]] = llvm.shl %[[VALUE_16]], %[[C16]]
// CHECK: %[[VALUE_32:.*]] = llvm.or %[[VALUE_16]], %[[SHIFTED_16]]
// CHECK: %[[VALUE_FLOAT:.+]] = llvm.bitcast %[[VALUE_32]] : i32 to f32
// CHECK-NOT: "llvm.intr.memset"
%load = llvm.load %alloca {alignment = 4 : i64} : !llvm.ptr -> f32
// CHECK: llvm.return %[[VALUE_FLOAT]] : f32
llvm.return %load : f32
}

// -----

// CHECK-LABEL: llvm.func @basic_memset_inline
// CHECK-SAME: (%[[MEMSET_VALUE:.*]]: i8)
llvm.func @basic_memset_inline(%memset_value: i8) -> i32 {
Expand Down Expand Up @@ -53,36 +77,37 @@ llvm.func @basic_memset_constant() -> i32 {
%memset_len = llvm.mlir.constant(4 : i32) : i32
"llvm.intr.memset"(%1, %memset_value, %memset_len) <{isVolatile = false}> : (!llvm.ptr, i8, i32) -> ()
%2 = llvm.load %1 {alignment = 4 : i64} : !llvm.ptr -> i32
// CHECK: %[[C42:.*]] = llvm.mlir.constant(42 : i8) : i8
// CHECK: %[[VALUE_42:.*]] = llvm.zext %[[C42]] : i8 to i32
// CHECK: %[[C8:.*]] = llvm.mlir.constant(8 : i32) : i32
// CHECK: %[[SHIFTED_42:.*]] = llvm.shl %[[VALUE_42]], %[[C8]] : i32
// CHECK: %[[OR0:.*]] = llvm.or %[[VALUE_42]], %[[SHIFTED_42]] : i32
// CHECK: %[[C16:.*]] = llvm.mlir.constant(16 : i32) : i32
// CHECK: %[[SHIFTED:.*]] = llvm.shl %[[OR0]], %[[C16]] : i32
// CHECK: %[[RES:..*]] = llvm.or %[[OR0]], %[[SHIFTED]] : i32
// CHECK: llvm.return %[[RES]] : i32
// CHECK: %[[CONSTANT_VAL:..*]] = llvm.mlir.constant(707406378 : i32) : i32
// CHECK: llvm.return %[[CONSTANT_VAL]] : i32
llvm.return %2 : i32
}

// -----

// CHECK-LABEL: llvm.func @memset_one_byte_constant
llvm.func @memset_one_byte_constant() -> i8 {
%one = llvm.mlir.constant(1 : i32) : i32
%alloca = llvm.alloca %one x i8 : (i32) -> !llvm.ptr
// CHECK: %{{.+}} = llvm.mlir.constant(42 : i8) : i8
%value = llvm.mlir.constant(42 : i8) : i8
"llvm.intr.memset"(%alloca, %value, %one) <{isVolatile = false}> : (!llvm.ptr, i8, i32) -> ()
%load = llvm.load %alloca : !llvm.ptr -> i8
// CHECK: %[[CONSTANT_VAL:..*]] = llvm.mlir.constant(42 : i8) : i8
// CHECK: llvm.return %[[CONSTANT_VAL]] : i8
llvm.return %load : i8
}

// -----

// CHECK-LABEL: llvm.func @basic_memset_inline_constant
llvm.func @basic_memset_inline_constant() -> i32 {
%0 = llvm.mlir.constant(1 : i32) : i32
%1 = llvm.alloca %0 x i32 {alignment = 4 : i64} : (i32) -> !llvm.ptr
%memset_value = llvm.mlir.constant(42 : i8) : i8
"llvm.intr.memset.inline"(%1, %memset_value) <{isVolatile = false, len = 4}> : (!llvm.ptr, i8) -> ()
%2 = llvm.load %1 {alignment = 4 : i64} : !llvm.ptr -> i32
// CHECK: %[[C42:.*]] = llvm.mlir.constant(42 : i8) : i8
// CHECK: %[[VALUE_42:.*]] = llvm.zext %[[C42]] : i8 to i32
// CHECK: %[[C8:.*]] = llvm.mlir.constant(8 : i32) : i32
// CHECK: %[[SHIFTED_42:.*]] = llvm.shl %[[VALUE_42]], %[[C8]] : i32
// CHECK: %[[OR0:.*]] = llvm.or %[[VALUE_42]], %[[SHIFTED_42]] : i32
// CHECK: %[[C16:.*]] = llvm.mlir.constant(16 : i32) : i32
// CHECK: %[[SHIFTED:.*]] = llvm.shl %[[OR0]], %[[C16]] : i32
// CHECK: %[[RES:..*]] = llvm.or %[[OR0]], %[[SHIFTED]] : i32
// CHECK: llvm.return %[[RES]] : i32
// CHECK: %[[CONSTANT_VAL:..*]] = llvm.mlir.constant(707406378 : i32) : i32
// CHECK: llvm.return %[[CONSTANT_VAL]] : i32
llvm.return %2 : i32
}

Expand Down