Skip to content
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

[xls][mlir] Add LiteralOp to XLS MLIR Dialect. #1866

Merged
merged 1 commit into from
Feb 4, 2025
+141 −7
Merged

[xls][mlir] Add LiteralOp to XLS MLIR Dialect. #1866

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
18 changes: 18 additions & 0 deletions xls/contrib/mlir/IR/interfaces.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,5 +78,23 @@ def PredicatableOpInterface : OpInterface<"PredicatableOpInterface"> {
];
}

// Declares that an op is permissible inside a literal region.
def LiteralMemberOpInterface : OpInterface<"LiteralMemberOpInterface"> {
Copy link
Member

Choose a reason for hiding this comment

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

Why not make this a trait and only add when true?

Copy link
Contributor Author

@schilkp schilkp Jan 21, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

I considered this, but as it stands right now the decision for if an op is permissible inside a literal is dynamic:

XLS IR allows for token literals:

literal.1: token = literal(value=token, id=1)

According to Eric this is equivalent to generating a new token with an empty join() in DSLX, resp. after_all in XLS IR without operands. I am assuming the reason this is supported is because it allows the literal IR op to construct complex types with tokens nested inside.

James and I agreed that for the MLIR side we would not introduce an extra op for this, and simply allow an xls.after_all with empty operands inside literal regions. This, in turn, means that for AfterAllOp, this can't really be a trait and needs some (simple) logic:

def Xls_AfterAllOp : Xls_Op<"after_all", [
  Pure,
  DeclareOpInterfaceMethods<LiteralMemberOpInterface>
]> {
  // ....
  let extraClassDeclaration = [{
    bool isPermissibleInsideLiteral() {
      // xls.after_all is only permissible inside literal regions if it is
      // used to create a new token with no prior dependencies.
      return getOperands().empty();
    }
  }];
}

Hence I went with the interface because it still allowed me to have this operation property included in a (halfway) declarative fashion in the ODS.

If you prefer, I can of course refactor this into a helper func which checks op IDs + checks after_all for an empty operand list.

Let me know :)

Copy link
Collaborator

Choose a reason for hiding this comment

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

That rationale makes sense to me.

let description = [{
Declares that an op is permissible inside a literal region.
}];
let cppNamespace = "::mlir::xls";

let methods = [
InterfaceMethod<
/*desc=*/"Returns true if the operation, in its current form, is permissible inside a literal region.",
/*retTy=*/"bool",
/*methodName=*/"isPermissibleInsideLiteral",
/*args=*/(ins),
/*methodBody=*/"",
/*defaultImplementation=*/"return true;"
>,
];
}

#endif
27 changes: 27 additions & 0 deletions xls/contrib/mlir/IR/xls_ops.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -484,6 +484,33 @@ Type getI1TypeOf(Type type) {

namespace mlir::xls {

LogicalResult LiteralOp::verifyRegions() {
Block& b = getInitializerBlock();

// Verify return op & typing:
auto ret = dyn_cast<YieldOp>(b.getTerminator());
if (!ret) {
schilkp marked this conversation as resolved.
Show resolved Hide resolved
return emitOpError(
"initializer region must be terminated by an xls.yield op");
}
if (ret->getNumOperands() != 1) {
return emitOpError("initializer region must return exactly one value");
}
if (ret->getOperand(0).getType() != getType()) {
return emitOpError("initializer region type ")
<< ret->getOperand(0).getType() << " does not match literal type "
<< getType();
}
// Verify that operations inside are allowed inside a literal region:
for (Operation& op : b) {
Copy link
Member

Choose a reason for hiding this comment

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

If you make it a trait, you should be able to use llvm all_of kind of check here with trait checking.

auto iface = dyn_cast<LiteralMemberOpInterface>(op);
if (!iface || !iface.isPermissibleInsideLiteral()) {
return op.emitError() << "op not permissible inside literal region";
}
}
return success();
}

LogicalResult TupleOp::inferReturnTypes(
MLIRContext* context, std::optional<Location> location, ValueRange operands,
DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions,
Expand Down
85 changes: 78 additions & 7 deletions xls/contrib/mlir/IR/xls_ops.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -562,7 +562,12 @@ def Xls_SendOp : Xls_Op<"send", [
// Array operations
//===----------------------------------------------------------------------===//

def Xls_ArrayOp : Xls_NonScalarizableNaryOp<"array", [Pure, SameTypeOperands, TensorArrayTypeFungible]> {
def Xls_ArrayOp : Xls_NonScalarizableNaryOp<"array", [
Pure,
SameTypeOperands,
TensorArrayTypeFungible,
DeclareOpInterfaceMethods<LiteralMemberOpInterface>
]> {
let summary = "Array creation operation";
let description = [{
Constructs an array of its operands.
Expand Down Expand Up @@ -721,8 +726,12 @@ def Xls_ArrayConcatOp : Xls_Op<"array_concat", [Pure, TensorArrayTypeFungible]>
// Tuple operations
//===----------------------------------------------------------------------===//

def Xls_TupleOp : Xls_Op<"tuple", [Pure, DeclareOpInterfaceMethods<InferTypeOpInterface>,
TensorArrayTypeFungible]> {
def Xls_TupleOp : Xls_Op<"tuple", [
Pure,
DeclareOpInterfaceMethods<InferTypeOpInterface>,
DeclareOpInterfaceMethods<LiteralMemberOpInterface>,
TensorArrayTypeFungible,
]> {
let summary = "Tuple operation";
let description = [{
Constructs a tuple of its operands.
Expand Down Expand Up @@ -1048,10 +1057,16 @@ def Xls_MapOp : Xls_Op<"map", []> {
// TODO(jmolloy): DynamicCountedFor (but the XLS code generator cannot handle
// it).

def Xls_YieldOp : Xls_Op<"yield", [Pure, ReturnLike, Terminator]> {
def Xls_YieldOp : Xls_Op<"yield", [
Pure,
ReturnLike,
Terminator,
DeclareOpInterfaceMethods<LiteralMemberOpInterface>
]> {
let summary = "for return op";
let description = [{
Yields an SSA value from the enclosing ForOp, EprocOp, or SprocOp region.
Yields an SSA value from the enclosing LiteralOp, ForOp, EprocOp, or
SprocOp region.
}];

let arguments = (ins Variadic<AnyType>:$results);
Expand Down Expand Up @@ -1159,7 +1174,10 @@ def Xls_CountedForOp : Xls_Op<"counted_for",
// Sequencing operations
//===----------------------------------------------------------------------===//

def Xls_AfterAllOp : Xls_Op<"after_all", [Pure]> {
def Xls_AfterAllOp : Xls_Op<"after_all", [
Pure,
DeclareOpInterfaceMethods<LiteralMemberOpInterface>
]> {
let summary = "Constructs partial orderings among channel operations";
let description = [{
Used to construct partial orderings among channel operations.
Expand All @@ -1177,6 +1195,13 @@ def Xls_AfterAllOp : Xls_Op<"after_all", [Pure]> {
build($_builder, $_state, TokenType::get($_builder.getContext()), ::ValueRange{});
}]>
];
let extraClassDeclaration = [{
bool isPermissibleInsideLiteral() {
// xls.after_all is only permissible inside literal regions if it is
// used to create a new token with no prior dependencies.
return getOperands().empty();
}
}];
let assemblyFormat = [{
$inputs attr-dict `:` custom<VariadicSameOperandsAndResultType>(ref($inputs), type($inputs), type($result))
}];
Expand All @@ -1200,7 +1225,11 @@ def Xls_ConstantTensorOp : Xls_Op<"constant_tensor", [ConstantLike, Pure]> {
let hasFolder = 1;
}

def Xls_ConstantScalarOp : Xls_Op<"constant_scalar", [ConstantLike, Pure]> {
def Xls_ConstantScalarOp : Xls_Op<"constant_scalar", [
ConstantLike,
Pure,
DeclareOpInterfaceMethods<LiteralMemberOpInterface>
]> {
let summary = "Constant scalar operation";
let description = [{
Produces an output from a constant value.
Expand All @@ -1225,6 +1254,48 @@ def Xls_ConstantScalarOp : Xls_Op<"constant_scalar", [ConstantLike, Pure]> {
let hasFolder = 1;
}

def Xls_LiteralOp : Xls_Op<"literal", [
ConstantLike,
Pure,
IsolatedFromAbove,
SingleBlockImplicitTerminator<"YieldOp">
]> {
let summary = "Complex Literal";
let description = [{
Produces a complex (non-bits) constant value from a region containing
simple constants or tuple/array construction operations.

The initializer region must be terminated by a `xls.yield` op, which returns
the constructed value.

Besides this yield, only the following operations are permissible inside
Copy link
Member

Choose a reason for hiding this comment

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

I think it would be computationally cheaper to just have a helper function with the TypeIDs of these ops rather than using traits (you'd be doing pointer comparison based on name rather than doing a linear scan per op through trait list). If these are this well known and we don't need to support any other dialects/unknown supported op types.

Copy link
Collaborator

Choose a reason for hiding this comment

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

Not an unreasonable suggestion, but personally I like that this is all having tablegen as the source of truth, and I'm not majorly worried about linear scans of literal regions in the verifier (they'll be small).

the initializer region:

- `xls.constant_scalar`
- `xls.after_all` (with no operands)
- `xls.tuple`
- `xls.array`

The `LiteralMemberOpInterface` interface is used to determine if a given
incarnation of an operation is permissible inside a literal region.
}];
let regions = (region
SizedRegion<1>:$initializer
);
let results = (outs
Xls_BitsOrTuple:$result
);
let hasRegionVerifier = 1;
let extraClassDeclaration = [{
Block &getInitializerBlock() {
return getInitializer().front();
}
}];
let assemblyFormat = [{
attr-dict `:` type($result) $initializer
}];
}

//===----------------------------------------------------------------------===//
// DSLX operations
//===----------------------------------------------------------------------===//
Expand Down
18 changes: 18 additions & 0 deletions xls/contrib/mlir/testdata/ops.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -376,6 +376,24 @@ func.func @constant_scalar() -> i7 {
return %0 : i7
}

// CHECK-LABEL: literal
func.func @complex_literal() -> tuple<tuple<i1, i2, tuple<i32, i32>>, !xls.array<2 x i3>> {
%lit = xls.literal : tuple<tuple<i1, i2, tuple<i32, i32>>, !xls.array<2 x i3>> {
%0 = "xls.constant_scalar"() <{value = true}> : () -> i1
%1 = "xls.constant_scalar"() <{value = -2 : i2}> : () -> i2
%2 = "xls.constant_scalar"() <{value = 10 : i32}> : () -> i32
%3 = "xls.constant_scalar"() <{value = 0 : i32}> : () -> i32
%4 = "xls.tuple"(%2, %3) : (i32, i32) -> tuple<i32, i32>
%5 = "xls.tuple"(%0, %1, %4) : (i1, i2, tuple<i32, i32>) -> tuple<i1, i2, tuple<i32, i32>>
%6 = "xls.constant_scalar"() <{value = -4 : i3}> : () -> i3
%7 = "xls.constant_scalar"() <{value = -3 : i3}> : () -> i3
%8 = xls.array %6, %7 : (i3, i3) -> !xls.array<2 x i3>
%final = "xls.tuple"(%5, %8) : (tuple<i1, i2, tuple<i32, i32>>, !xls.array<2 x i3>) -> tuple<tuple<i1, i2, tuple<i32, i32>>, !xls.array<2 x i3>>
xls.yield %final : tuple<tuple<i1, i2, tuple<i32, i32>>, !xls.array<2 x i3>>
}
return %lit : tuple<tuple<i1, i2, tuple<i32, i32>>, !xls.array<2 x i3>>
}

// CHECK-LABEL: for
func.func @for(%arg0: i32, %arg1: i8, %arg2: i9) -> i32 {
// CHECK: xls.for
Expand Down