Skip to content

[mlir][arith] Add overflow flags support to arith ops #77211

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 7 commits into from
Jan 9, 2024
Merged

[mlir][arith] Add overflow flags support to arith ops #77211

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
0c1f04d
[mlir][arith] Add overflow flags support to arith ops
yiwu0b11 Jan 5, 2024
82eeda0
fix python test
Hardcode84 Jan 6, 2024
3d09907
Update docs and comments
Hardcode84 Jan 7, 2024
4c6f3ac
Update docs, Arith_IntBinaryOpWithOverflowFlags, cleanup hasFolder/ha…
Hardcode84 Jan 8, 2024
c36ae9d
Canonicalizations comment
Hardcode84 Jan 8, 2024
b69c5b8
better comment
Hardcode84 Jan 8, 2024
750e0f7
formatting
Hardcode84 Jan 9, 2024
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
35 changes: 35 additions & 0 deletions mlir/include/mlir/Conversion/ArithCommon/AttrToLLVMConverter.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,14 @@ convertArithFastMathFlagsToLLVM(arith::FastMathFlags arithFMF);
LLVM::FastmathFlagsAttr
convertArithFastMathAttrToLLVM(arith::FastMathFlagsAttr fmfAttr);

// Map arithmetic overflow enum values to LLVMIR enum values.
LLVM::IntegerOverflowFlags
convertArithOveflowFlagsToLLVM(arith::IntegerOverflowFlags arithFlags);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: typo, Oveflow -> Overflow


// Create an LLVM overflow attribute from a given arithmetic overflow attribute.
LLVM::IntegerOverflowFlagsAttr
convertArithOveflowAttrToLLVM(arith::IntegerOverflowFlagsAttr flagsAttr);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: typo, Oveflow -> Overflow


// Attribute converter that populates a NamedAttrList by removing the fastmath
// attribute from the source operation attributes, and replacing it with an
// equivalent LLVM fastmath attribute.
Expand All @@ -49,6 +57,33 @@ class AttrConvertFastMathToLLVM {

ArrayRef<NamedAttribute> getAttrs() const { return convertedAttr.getAttrs(); }

private:
NamedAttrList convertedAttr;
};

// Attribute converter that populates a NamedAttrList by removing the overflow
// attribute from the source operation attributes, and replacing it with an
// equivalent LLVM overflow attribute.
template <typename SourceOp, typename TargetOp>
class AttrConvertOverflowToLLVM {
public:
AttrConvertOverflowToLLVM(SourceOp srcOp) {
// Copy the source attributes.
convertedAttr = NamedAttrList{srcOp->getAttrs()};
// Get the name of the arith overflow attribute.
llvm::StringRef arithAttrName = SourceOp::getIntegerOverflowAttrName();
// Remove the source overflow attribute.
auto arithAttr = dyn_cast_or_null<arith::IntegerOverflowFlagsAttr>(
convertedAttr.erase(arithAttrName));
if (arithAttr) {
llvm::StringRef targetAttrName = TargetOp::getIntegerOverflowAttrName();
convertedAttr.set(targetAttrName,
convertArithOveflowAttrToLLVM(arithAttr));
}
}

ArrayRef<NamedAttribute> getAttrs() const { return convertedAttr.getAttrs(); }

private:
NamedAttrList convertedAttr;
};
Expand Down
23 changes: 23 additions & 0 deletions mlir/include/mlir/Dialect/Arith/IR/ArithBase.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,4 +133,27 @@ def Arith_FastMathAttr :
let assemblyFormat = "`<` $value `>`";
}

//===----------------------------------------------------------------------===//
// IntegerOverflowFlags
Copy link
Contributor

@superbobry superbobry Jan 11, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This seems to break the Python bindings, since now both the arith and llvm dialects define IntegerOverflowFlags:

Traceback (most recent call last):
  ...
File "[...]/mlir/dialects/_llvm_enum_gen.py", line 681, in <module>
    @register_attribute_builder("IntegerOverflowFlags")
     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "[...]/mlir/ir.py", line 14, in decorator_builder
    AttrBuilder.insert(kind, func, replace=replace)
RuntimeError: Attribute builder for 'IntegerOverflowFlags' is already registered with func: <function _integeroverflowflags at 0x7faf26f8ce00>

The repro is

from mlir.dialects import arith, llvm

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was wondering why this wasn't the case for fastmath flags and they are actually named slightly differently between llvm and arith (FastMathFlags vs FastmathFlags).

Anyways, I'm not familiar with in-tree MLIR python bindings and considering these enums live in different namespaces on C++ level, this sounds like quite a big limitation, which people will continue to hit randomly.

Any ideas how to fix it properly?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just need to namespace (by dialect) the generated enum bindings here and here. Will send a patch soon.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reverted for now

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Independent of Maks upcoming change, it's good to "textually namespace" TableGen side too. That's the common case, but I see unfortunately only documented such for ops and here you were being consistent with how fast math flags were defined. (This reminds me of the ODS linter again).

//===----------------------------------------------------------------------===//

def IOFnone : I32BitEnumAttrCaseNone<"none">;
def IOFnsw : I32BitEnumAttrCaseBit<"nsw", 0>;
def IOFnuw : I32BitEnumAttrCaseBit<"nuw", 1>;

def IntegerOverflowFlags : I32BitEnumAttr<
"IntegerOverflowFlags",
"Integer overflow arith flags",
[IOFnone, IOFnsw, IOFnuw]> {
let separator = ", ";
let cppNamespace = "::mlir::arith";
let genSpecializedAttr = 0;
let printBitEnumPrimaryGroups = 1;
}

def Arith_IntegerOverflowAttr :
EnumAttr<Arith_Dialect, IntegerOverflowFlags, "overflow"> {
let assemblyFormat = "`<` $value `>`";
}

#endif // ARITH_BASE
101 changes: 82 additions & 19 deletions mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,22 @@ class Arith_CompareOpOfAnyRank<string mnemonic, list<Trait> traits = []> :
let results = (outs BoolLikeOfAnyRank:$result);
}

class Arith_IntArithmeticOpWithOverflowFlag<string mnemonic, list<Trait> traits = []> :
Arith_BinaryOp<mnemonic, traits #
[Pure, DeclareOpInterfaceMethods<InferIntRangeInterface>,
DeclareOpInterfaceMethods<ArithIntegerOverflowFlagsInterface>]>,
Arguments<(ins SignlessIntegerLike:$lhs, SignlessIntegerLike:$rhs,
DefaultValuedAttr<
Arith_IntegerOverflowAttr, "::mlir::arith::IntegerOverflowFlags::none">:$overflowFlags)>,
Results<(outs SignlessIntegerLike:$result)> {

let assemblyFormat = [{ $lhs `,` $rhs (`overflow` `` $overflowFlags^)?
attr-dict `:` type($result) }];

let hasFolder = 1;
let hasCanonicalizer = 1;
}

//===----------------------------------------------------------------------===//
// ConstantOp
//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -192,7 +208,7 @@ def Arith_ConstantOp : Op<Arith_Dialect, "constant",
// AddIOp
//===----------------------------------------------------------------------===//

def Arith_AddIOp : Arith_TotalIntBinaryOp<"addi", [Commutative]> {
def Arith_AddIOp : Arith_IntArithmeticOpWithOverflowFlag<"addi", [Commutative]> {
let summary = "integer addition operation";
let description = [{
Performs N-bit addition on the operands. The operands are interpreted as
Expand All @@ -203,15 +219,22 @@ def Arith_AddIOp : Arith_TotalIntBinaryOp<"addi", [Commutative]> {

The `addi` operation takes two operands and returns one result, each of
these is required to be the same type. This type may be an integer scalar type,
a vector whose element type is integer, or a tensor of integers. It has no
standard attributes.
a vector whose element type is integer, or a tensor of integers.

This op supports nuw/nsw flags which stands stand for "No Unsigned Wrap" and
"No Signed Wrap", respectively. If the nuw and/or nsw flags are present, the
result value is undefined if unsigned and/or signed overflow, respectively,
occurs.

Example:

```mlir
// Scalar addition.
%a = arith.addi %b, %c : i64

// Scalar addition with overflow flags.
%a = arith.addi %b, %c overflow<nsw, nuw> : i64

// SIMD vector element-wise addition, e.g. for Intel SSE.
%f = arith.addi %g, %h : vector<4xi32>

Expand Down Expand Up @@ -278,21 +301,41 @@ def Arith_AddUIExtendedOp : Arith_Op<"addui_extended", [Pure, Commutative,
// SubIOp
//===----------------------------------------------------------------------===//

def Arith_SubIOp : Arith_TotalIntBinaryOp<"subi"> {
def Arith_SubIOp : Arith_IntArithmeticOpWithOverflowFlag<"subi"> {
let summary = [{
Integer subtraction operation.
}];
let description = [{
Performs N-bit subtraction on the operands. The operands are interpreted as unsigned
bitvectors. The result is represented by a bitvector containing the mathematical
value of the subtraction modulo 2^n, where `n` is the bitwidth. Because `arith`
integers use a two's complement representation, this operation is applicable on
Performs N-bit subtraction on the operands. The operands are interpreted as unsigned
bitvectors. The result is represented by a bitvector containing the mathematical
value of the subtraction modulo 2^n, where `n` is the bitwidth. Because `arith`
integers use a two's complement representation, this operation is applicable on
both signed and unsigned integer operands.

The `subi` operation takes two operands and returns one result, each of
these is required to be the same type. This type may be an integer scalar type,
a vector whose element type is integer, or a tensor of integers. It has no
standard attributes.
these is required to be the same type. This type may be an integer scalar type,
a vector whose element type is integer, or a tensor of integers.

This op supports nuw/nsw flags which stands stand for "No Unsigned Wrap" and
"No Signed Wrap", respectively. If the nuw and/or nsw flags are present, the
result value is undefined if unsigned and/or signed overflow, respectively,
occurs.

Example:

```mlir
// Scalar subtraction.
%a = arith.subi %b, %c : i64

// Scalar subtraction with overflow flags.
%a = arith.subi %b, %c overflow<nsw, nuw> : i64

// SIMD vector element-wise subtraction, e.g. for Intel SSE.
%f = arith.subi %g, %h : vector<4xi32>

// Tensor element-wise subtraction.
%x = arith.subi %y, %z : tensor<4x?xi8>
```
}];
let hasFolder = 1;
let hasCanonicalizer = 1;
Expand All @@ -302,21 +345,41 @@ def Arith_SubIOp : Arith_TotalIntBinaryOp<"subi"> {
// MulIOp
//===----------------------------------------------------------------------===//

def Arith_MulIOp : Arith_TotalIntBinaryOp<"muli", [Commutative]> {
def Arith_MulIOp : Arith_IntArithmeticOpWithOverflowFlag<"muli", [Commutative]> {
let summary = [{
Integer multiplication operation.
}];
let description = [{
Performs N-bit multiplication on the operands. The operands are interpreted as
unsigned bitvectors. The result is represented by a bitvector containing the
mathematical value of the multiplication modulo 2^n, where `n` is the bitwidth.
Because `arith` integers use a two's complement representation, this operation is
Performs N-bit multiplication on the operands. The operands are interpreted as
unsigned bitvectors. The result is represented by a bitvector containing the
mathematical value of the multiplication modulo 2^n, where `n` is the bitwidth.
Because `arith` integers use a two's complement representation, this operation is
applicable on both signed and unsigned integer operands.

The `muli` operation takes two operands and returns one result, each of
these is required to be the same type. This type may be an integer scalar type,
a vector whose element type is integer, or a tensor of integers. It has no
standard attributes.
these is required to be the same type. This type may be an integer scalar type,
a vector whose element type is integer, or a tensor of integers.

This op supports nuw/nsw flags which stands stand for "No Unsigned Wrap" and
"No Signed Wrap", respectively. If the nuw and/or nsw flags are present, the
result value is undefined if unsigned and/or signed overflow, respectively,
occurs.

Example:

```mlir
// Scalar multiplication.
%a = arith.muli %b, %c : i64

// Scalar multiplication with overflow flags.
%a = arith.muli %b, %c overflow<nsw, nuw> : i64

// SIMD vector element-wise multiplication, e.g. for Intel SSE.
%f = arith.muli %g, %h : vector<4xi32>

// Tensor element-wise multiplication.
%x = arith.muli %y, %z : tensor<4x?xi8>
```
}];
let hasFolder = 1;
let hasCanonicalizer = 1;
Expand Down
57 changes: 57 additions & 0 deletions mlir/include/mlir/Dialect/Arith/IR/ArithOpsInterfaces.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,4 +49,61 @@ def ArithFastMathInterface : OpInterface<"ArithFastMathInterface"> {
];
}

def ArithIntegerOverflowFlagsInterface : OpInterface<"ArithIntegerOverflowFlagsInterface"> {
let description = [{
Access to op integer overflow flags.
}];

let cppNamespace = "::mlir::arith";

let methods = [
InterfaceMethod<
/*desc=*/ "Returns an IntegerOverflowFlagsAttr attribute for the operation",
/*returnType=*/ "IntegerOverflowFlagsAttr",
/*methodName=*/ "getOverflowAttr",
/*args=*/ (ins),
/*methodBody=*/ [{}],
/*defaultImpl=*/ [{
auto op = cast<ConcreteOp>(this->getOperation());
return op.getOverflowFlagsAttr();
}]
>,
InterfaceMethod<
/*desc=*/ "Returns whether the operation has the No Unsigned Wrap keyword",
/*returnType=*/ "bool",
/*methodName=*/ "hasNoUnsignedWrap",
/*args=*/ (ins),
/*methodBody=*/ [{}],
/*defaultImpl=*/ [{
auto op = cast<ConcreteOp>(this->getOperation());
IntegerOverflowFlags flags = op.getOverflowFlagsAttr().getValue();
return bitEnumContainsAll(flags, IntegerOverflowFlags::nuw);
}]
>,
InterfaceMethod<
/*desc=*/ "Returns whether the operation has the No Signed Wrap keyword",
/*returnType=*/ "bool",
/*methodName=*/ "hasNoSignedWrap",
/*args=*/ (ins),
/*methodBody=*/ [{}],
/*defaultImpl=*/ [{
auto op = cast<ConcreteOp>(this->getOperation());
IntegerOverflowFlags flags = op.getOverflowFlagsAttr().getValue();
return bitEnumContainsAll(flags, IntegerOverflowFlags::nsw);
}]
>,
StaticInterfaceMethod<
/*desc=*/ [{Returns the name of the IntegerOveflowFlagsAttr attribute
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: typo, Oveflow -> Overflow

for the operation}],
/*returnType=*/ "StringRef",
/*methodName=*/ "getIntegerOverflowAttrName",
/*args=*/ (ins),
/*methodBody=*/ [{}],
/*defaultImpl=*/ [{
return "overflowFlags";
}]
>
];
}

#endif // ARITH_OPS_INTERFACES
29 changes: 26 additions & 3 deletions mlir/lib/Conversion/ArithCommon/AttrToLLVMConverter.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,9 +22,9 @@ mlir::arith::convertArithFastMathFlagsToLLVM(arith::FastMathFlags arithFMF) {
{arith::FastMathFlags::contract, LLVM::FastmathFlags::contract},
{arith::FastMathFlags::afn, LLVM::FastmathFlags::afn},
{arith::FastMathFlags::reassoc, LLVM::FastmathFlags::reassoc}};
for (auto fmfMap : flags) {
if (bitEnumContainsAny(arithFMF, fmfMap.first))
llvmFMF = llvmFMF | fmfMap.second;
for (auto [arithFlag, llvmFlag] : flags) {
if (bitEnumContainsAny(arithFMF, arithFlag))
llvmFMF = llvmFMF | llvmFlag;
}
return llvmFMF;
}
Expand All @@ -36,3 +36,26 @@ mlir::arith::convertArithFastMathAttrToLLVM(arith::FastMathFlagsAttr fmfAttr) {
return LLVM::FastmathFlagsAttr::get(
fmfAttr.getContext(), convertArithFastMathFlagsToLLVM(arithFMF));
}

// Map arithmetic overflow enum values to LLVMIR enum values.
LLVM::IntegerOverflowFlags mlir::arith::convertArithOveflowFlagsToLLVM(
arith::IntegerOverflowFlags arithFlags) {
LLVM::IntegerOverflowFlags llvmFlags{};
const std::pair<arith::IntegerOverflowFlags, LLVM::IntegerOverflowFlags>
flags[] = {
{arith::IntegerOverflowFlags::nsw, LLVM::IntegerOverflowFlags::nsw},
{arith::IntegerOverflowFlags::nuw, LLVM::IntegerOverflowFlags::nuw}};
for (auto [arithFlag, llvmFlag] : flags) {
if (bitEnumContainsAny(arithFlags, arithFlag))
llvmFlags = llvmFlags | llvmFlag;
}
return llvmFlags;
}

// Create an LLVM overflow attribute from a given arithmetic overflow attribute.
LLVM::IntegerOverflowFlagsAttr mlir::arith::convertArithOveflowAttrToLLVM(
arith::IntegerOverflowFlagsAttr flagsAttr) {
arith::IntegerOverflowFlags arithFlags = flagsAttr.getValue();
return LLVM::IntegerOverflowFlagsAttr::get(
flagsAttr.getContext(), convertArithOveflowFlagsToLLVM(arithFlags));
}
12 changes: 9 additions & 3 deletions mlir/lib/Conversion/ArithToLLVM/ArithToLLVM.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,9 @@ namespace {
using AddFOpLowering =
VectorConvertToLLVMPattern<arith::AddFOp, LLVM::FAddOp,
arith::AttrConvertFastMathToLLVM>;
using AddIOpLowering = VectorConvertToLLVMPattern<arith::AddIOp, LLVM::AddOp>;
using AddIOpLowering =
VectorConvertToLLVMPattern<arith::AddIOp, LLVM::AddOp,
arith::AttrConvertOverflowToLLVM>;
using AndIOpLowering = VectorConvertToLLVMPattern<arith::AndIOp, LLVM::AndOp>;
using BitcastOpLowering =
VectorConvertToLLVMPattern<arith::BitcastOp, LLVM::BitcastOp>;
Expand Down Expand Up @@ -78,7 +80,9 @@ using MinUIOpLowering =
using MulFOpLowering =
VectorConvertToLLVMPattern<arith::MulFOp, LLVM::FMulOp,
arith::AttrConvertFastMathToLLVM>;
using MulIOpLowering = VectorConvertToLLVMPattern<arith::MulIOp, LLVM::MulOp>;
using MulIOpLowering =
VectorConvertToLLVMPattern<arith::MulIOp, LLVM::MulOp,
arith::AttrConvertOverflowToLLVM>;
using NegFOpLowering =
VectorConvertToLLVMPattern<arith::NegFOp, LLVM::FNegOp,
arith::AttrConvertFastMathToLLVM>;
Expand All @@ -102,7 +106,9 @@ using SIToFPOpLowering =
using SubFOpLowering =
VectorConvertToLLVMPattern<arith::SubFOp, LLVM::FSubOp,
arith::AttrConvertFastMathToLLVM>;
using SubIOpLowering = VectorConvertToLLVMPattern<arith::SubIOp, LLVM::SubOp>;
using SubIOpLowering =
VectorConvertToLLVMPattern<arith::SubIOp, LLVM::SubOp,
arith::AttrConvertOverflowToLLVM>;
using TruncFOpLowering =
VectorConvertToLLVMPattern<arith::TruncFOp, LLVM::FPTruncOp>;
using TruncIOpLowering =
Expand Down
Loading