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

This operation constructs a literal of arbitrarily complex type using a
body region that construct this value using basic scalar operations.

It more closely resembles the XLS literal IR node for loss-less
conversion between the IR and MLIR, but is not intended to be used
during complex operations and transforms in MLIR due to the lack of
constant folding.

A pass to remove this operation by inlining (and potentially a pass to
reconstruct such literal ops by grouping ops) will be added in the
future.

Only a small subset of XLS operations is permitted inside the body
region. The new LiteralMemberOpInterface is used to identify such
ops.

xls/contrib/mlir/IR/interfaces.td

@@ -78,5 +78,23 @@ def PredicatableOpInterface : OpInterface<"PredicatableOpInterface"> {
   ];
 }
 
+// Declares that an op is permissible inside a literal region.
+def LiteralMemberOpInterface : OpInterface<"LiteralMemberOpInterface"> {
+  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 permissibe inside a literal region.",
+      /*retTy=*/"bool",
+      /*methodName=*/"isPermissibleInsideLiteral",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/"return true;"
+    >,
+  ];
+}
 
 #endif

xls/contrib/mlir/IR/xls_ops.cc

@@ -390,6 +390,32 @@ 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) {
+    return emitOpError(
+        "initializer region must be terminated by an xls.yield op");
+  }
+  if (ret->operand_type_begin() == ret.operand_type_end()) {
+    return emitOpError("initializer region cannot return void");
+  }
+  if (*(ret->operand_type_begin()) != getType()) {
+    return emitOpError("initializer region type ")
+           << *ret.operand_type_begin() << " does not match global type "
+           << getType();
+  }
+  // Verify that operations inside are allowed inside a literal region:
+  for (Operation& op : b) {
+    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,

xls/contrib/mlir/IR/xls_ops.td

@@ -563,7 +563,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.
@@ -722,8 +727,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.
@@ -1049,10 +1058,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);
@@ -1160,7 +1175,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.
@@ -1178,6 +1196,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))
   }];
@@ -1201,7 +1226,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.
@@ -1226,6 +1255,50 @@ def Xls_ConstantScalarOp : Xls_Op<"constant_scalar", [ConstantLike, Pure]> {
   let hasFolder = 1;
 }
 
+def Xls_LiteralOp : Xls_Op<"literal", [ConstantLike, Pure, IsolatedFromAbove]> {
+  let summary = "Literal value of arbitrarily complex type.";
+  let description = [{
+    Produces a complex (non-bits) constant value from a region containing
+    simple constants or tuple/array construction operations.
+
+    The body must be terminated by a `xls.yield` op, which returns the
+    constructed value.
+
+    Besides this yield, only the following operations are permissible inside
+    the body 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>:$body
+  );
+  let results = (outs
+    Xls_BitsOrTuple:$result
+  );
+  let hasRegionVerifier = 1;
+  let extraClassDeclaration = [{
+    /// Return the initializer region. This must always include an initializer
+    /// block terminated by a `xls.yield`.
+    Region &getInitializerRegion() {
+      return getOperation()->getRegion(0);
+    }
+    /// Return the initializer block. It terminates with an `xls.yield` which
+    /// returns the literal value.
+    Block &getInitializerBlock() {
+      return getInitializerRegion().front();
+    }
+  }];
+  let assemblyFormat = [{
+    attr-dict `:` type($result) $body
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // DSLX operations
 //===----------------------------------------------------------------------===//

xls/contrib/mlir/testdata/ops.mlir

@@ -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