[xls][mlir] Support FifoConfig and FlopKinds on top-level Sprocs/Eprocs

Modifies the boundary_channel_names attribute from an array of strings to
an array of #xls.boundary_channel, which contains the name but also FifoConfig
and FlopKind.

PiperOrigin-RevId: 812789063

Author: James Molloy

xls/contrib/mlir/IR/xls_ops.td

@@ -203,6 +203,24 @@ def Xls_FifoConfig : Xls_Attr<"FifoConfig"> {
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
+def Xls_BoundaryChannelAttr : Xls_Attr<"BoundaryChannelAttr"> {
+  let summary = "Boundary channel attribute";
+  let description = [{
+    Contains attributes associated with boundary channels for SprocOps.
+  }];
+  let parameters = (ins
+    "::mlir::StringAttr":$name,
+    OptionalParameter<"FifoConfig", "">:$fifo_config,
+    OptionalParameter<"FlopKindAttr", "">:$input_flop_kind,
+    OptionalParameter<"FlopKindAttr", "">:$output_flop_kind
+  );
+  let mnemonic = "boundary_channel";
+  let assemblyFormat = "`<` struct(params) `>`";
+}
+
+def Xls_BoundaryChannelArrayAttr :
+  TypedArrayAttrBase<Xls_BoundaryChannelAttr, "array of boundary_channel">;
+
 class GetShapeSplat<string name> :
     StrFunc<"getShapeSplat($" # name # ".getType())">;
 
@@ -1874,9 +1892,8 @@ def Xls_SprocOp : Xls_Op<"sproc", [
     region contains Schan declarations and "spawns" of other Sprocs. The only way
     to create an Sproc is to spawn another Sproc.
 
-    Sprocs may be marked as "top" if they EITHER:
-      a) Have no channels.
-      b) Have the optional `boundary_channel_names` attribute set.
+    Sprocs may be marked as "top" if they populate the `boundary_channels`
+    list per argument.
 
     In the latter case, boundary channels (used by XLS for tests etc) are
     created with the given names.
@@ -1926,7 +1943,7 @@ def Xls_SprocOp : Xls_Op<"sproc", [
   let arguments = (ins
     SymbolNameAttr:$sym_name,
     BoolAttr:$is_top,
-    OptionalAttr<ArrayAttr>:$boundary_channel_names,
+    OptionalAttr<Xls_BoundaryChannelArrayAttr>:$boundary_channels,
     DefaultValuedAttr<I64Attr, "1">:$min_pipeline_stages
   );
   let regions = (region

xls/contrib/mlir/testdata/extract_as_top_level_module.mlir

@@ -80,7 +80,7 @@ module attributes {test.name = "sproc"} {
   // CHECK-NOT: func @nope
   // CHECK: xls.sproc @leaf
   // CHECK-NOT: xls.sproc @unrelated
-  // CHECK: xls.sproc @sproc(%arg0: !xls.schan<i32, in>) top attributes {boundary_channel_names = ["arg0"]}
+  // CHECK: xls.sproc @sproc(%arg0: !xls.schan<i32, in>) top attributes {boundary_channels = [#xls.boundary_channel<name = "arg0">]}
   func.func @nope(%arg0: !xls.token, %arg1: i32, %arg2: i1) ->i32 {
     %0 = arith.constant 1 : i32
     return %0 : i32

xls/contrib/mlir/testdata/ops.mlir

@@ -13,7 +13,7 @@ func.func @not(%arg0: tensor<32xi8>) -> tensor<32xi8> {
   %0 = xls.not %arg0 : tensor<32xi8>
   return %0 : tensor<32xi8>
 }
-
+ 
 // CHECK-LABEL: func @and
 func.func @and(%arg0: tensor<32xi8>, %arg1: tensor<32xi8>) -> tensor<32xi8> {
   // CHECK: xls.and %arg0, %arg1

xls/contrib/mlir/testdata/proc_elaboration.mlir

@@ -38,7 +38,7 @@
 // CHECK-NEXT:  xls.chan @fetch_arg0 : i32
 // CHECK-NEXT:  xls.chan @fetch_arg1 : i32
 // CHECK-NEXT:  xls.instantiate_eproc @fetch (@fetch_arg0 as @req, @fetch_arg1 as @resp)
-// CHECK-NEXT:  xls.chan @boundary1 {send_supported = false} : i32
+// CHECK-NEXT:  xls.chan @boundary1 {fifo_config = #xls.fifo_config<fifo_depth = 1, bypass = true, register_push_outputs = true, register_pop_outputs = false>, input_flop_kind = #xls<flop_kind skid>, send_supported = false} : i32
 // CHECK-NEXT:  xls.chan @boundary2 {recv_supported = false} : i32
 // CHECK-NEXT:  xls.instantiate_eproc @rom (@rom_arg0 as @boundary1, @rom_arg1 as @boundary2)
 
@@ -77,7 +77,10 @@ xls.sproc @proxy(%req: !xls.schan<i32, in>, %resp: !xls.schan<i32, out>) attribu
   }
 }
 
-xls.sproc @rom(%req: !xls.schan<i32, in>, %resp: !xls.schan<i32, out>) top attributes {boundary_channel_names = ["boundary1", "boundary2"]} {
+xls.sproc @rom(%req: !xls.schan<i32, in>, %resp: !xls.schan<i32, out>) top attributes {boundary_channels = [
+  #xls.boundary_channel<name = "boundary1", fifo_config = #xls.fifo_config<fifo_depth = 1, bypass = true, register_push_outputs = true, register_pop_outputs = false>, input_flop_kind = #xls<flop_kind skid>>,
+  #xls.boundary_channel<name = "boundary2">
+]} {
   spawns {
     xls.yield %req, %resp : !xls.schan<i32, in>, !xls.schan<i32, out>
   }
@@ -144,7 +147,9 @@ module {
       xls.yield %0 : index
     }
   }
-  xls.sproc @some_wrapped_machine(%arg0: !xls.schan<tensor<i32>, in>, %arg1: !xls.schan<tensor<i32>, out>, %arg2: !xls.schan<tensor<i32>, out>) top attributes {boundary_channel_names = ["x['y']", "x['y']1", "x['y']2"]} {
+  xls.sproc @some_wrapped_machine(%arg0: !xls.schan<tensor<i32>, in>, %arg1: !xls.schan<tensor<i32>, out>, %arg2: !xls.schan<tensor<i32>, out>) top attributes {
+    boundary_channels = [#xls.boundary_channel<name = "x['y']">, #xls.boundary_channel<name = "x['y']1">, #xls.boundary_channel<name = "x['y']2">]
+  } {
     spawns {
       %out, %in = xls.schan<tensor<i32>>("x['y']")
       %out_0, %in_1 = xls.schan<tensor<i32>>("x['y']")

xls/contrib/mlir/transforms/proc_elaboration.cc

@@ -323,18 +323,19 @@ void ProcElaborationPass::runOnOperation() {
 
     OpBuilder builder(sproc);
     SmallVector<ChanOp> boundaryChannels;
-    if (sproc.getBoundaryChannelNames().has_value()) {
-      for (auto [arg, name] : llvm::zip(sproc.getChannelArguments(),
-                                        *sproc.getBoundaryChannelNames())) {
+    if (sproc.getBoundaryChannels().has_value()) {
+      for (auto [arg, boundaryChannelAttr] : llvm::zip(
+               sproc.getChannelArguments(), *sproc.getBoundaryChannels())) {
         SchanType schan = cast<SchanType>(arg.getType());
-        auto nameAttr = cast<StringAttr>(name);
+        auto boundaryChannel = cast<BoundaryChannelAttr>(boundaryChannelAttr);
+        auto nameAttr = boundaryChannel.getName();
         // TODO(jpienaar): Remove unnecessary default args once they are
         // generated automatically.
-        auto echan = ChanOp::create(builder, sproc.getLoc(), nameAttr.str(),
-                                    schan.getElementType(),
-                                    /*fifo_config=*/nullptr,
-                                    /*input_flop_kind=*/nullptr,
-                                    /*output_flop_kind=*/nullptr);
+        auto echan = ChanOp::create(
+            builder, sproc.getLoc(), boundaryChannel.getName(),
+            schan.getElementType(), boundaryChannel.getFifoConfig(),
+            boundaryChannel.getInputFlopKind(),
+            boundaryChannel.getOutputFlopKind());
         // We insert the channel in the symbol table, since there might be
         // a clash with an *eproc* name later, so we need to know. Alternatively
         // we could have done something similar in the constructor of the

xls/contrib/mlir/util/extraction_utils.cc

@@ -137,16 +137,19 @@ void addBoundaryChannelNames(
     };
   }
 
-  if (op.getBoundaryChannelNames().has_value()) {
+  if (op.getBoundaryChannels().has_value()) {
     return;
   }
   ::mlir::OpBuilder builder(op.getOperation());
-  ::mlir::SmallVector<::mlir::Attribute> boundaryChannelNames;
+  ::mlir::SmallVector<::mlir::Attribute> boundaryChannels;
   for (Value value : op.getChannelArguments()) {
-    boundaryChannelNames.push_back(
-        builder.getStringAttr(boundaryChannelName(cast<BlockArgument>(value))));
+    BoundaryChannelAttr boundaryChannel = BoundaryChannelAttr::get(
+        builder.getContext(),
+        builder.getStringAttr(boundaryChannelName(cast<BlockArgument>(value))),
+        FifoConfig(), FlopKindAttr(), FlopKindAttr());
+    boundaryChannels.push_back(boundaryChannel);
   }
-  op.setBoundaryChannelNamesAttr(builder.getArrayAttr(boundaryChannelNames));
+  op.setBoundaryChannelsAttr(builder.getArrayAttr(boundaryChannels));
 }
 }  // namespace