[xls][mlir] Support XLS fifo properties/config in MLIR channels.

This also adds support for converting these properties to/from XLS, and
updates different passes (array_to_bits, index_type_conversion) to
ensure these attributes are maintained when channels are modified.

Finally, these properties are also exposed for sproc channels, and the
proc elaboration pass is updated to assign said properties to all
eproc channels that are generated from an sproc channel.

Author: schilkp

xls/contrib/mlir/IR/xls_ops.cc

@@ -620,7 +620,9 @@ ParseResult SchanOp::parse(OpAsmParser& parser, OperationState& result) {
   result.addAttribute("type", TypeAttr::get(type));
   result.types.push_back(SchanType::get(parser.getContext(), type, false));
   result.types.push_back(SchanType::get(parser.getContext(), type, true));
-  return success();
+
+  return failure(
+      parser.parseOptionalAttrDictWithKeyword(result.attributes).failed());
 }
 
 void SchanOp::print(OpAsmPrinter& printer) {
@@ -630,6 +632,9 @@ void SchanOp::print(OpAsmPrinter& printer) {
   printer << '(';
   printer.printString(getName());
   printer << ')';
+  printer << ") ";
+  printer.printOptionalAttrDictWithKeyword(getOperation()->getAttrs(),
+                                           {"type", "name"});
 }
 
 void SprocOp::print(OpAsmPrinter& printer) {

xls/contrib/mlir/IR/xls_ops.td

@@ -135,6 +135,75 @@ def Xls_TranslationLinkage : Xls_Attr<"TranslationLinkage"> {
   let attrName = "xls.translation_linkage";
 }
 
+def Xls_FlopKindNone        : I32EnumAttrCase<"kNone", 0, "none"> {
+  let summary = "No flop.";
+  let description = [{
+    Adds no buffer to the input/output.
+
+    See https://google.github.io/xls/codegen_options/#io-behavior for more information.
+  }];
+}
+def Xls_FlopKindFlop        : I32EnumAttrCase<"kFlop", 1, "flop"> {
+  let summary = "Basic flop.";
+  let description = [{
+    Adds a pipeline stage to the input/output. This is essentially a single-element FIFO
+
+    See https://google.github.io/xls/codegen_options/#io-behavior for more information.
+  }];
+}
+def Xls_FlopKindSkid        : I32EnumAttrCase<"kSkid", 2, "skid"> {
+  let summary = "Skip flop.";
+  let description = [{
+    Adds a skid buffer to the input/output. The skid buffer can hold 2 entries.
+
+    See https://google.github.io/xls/codegen_options/#io-behavior for more information.
+  }];
+}
+def Xls_FlopKindZeroLatency : I32EnumAttrCase<"kZeroLatency", 3, "zero-latency"> {
+  let summary = "Zero-latency flop.";
+  let description = [{
+    Adds a zero-latency buffer to the input/output. This is essentially a
+    single-element FIFO with bypass.
+
+    See https://google.github.io/xls/codegen_options/#io-behavior for more information.
+  }];
+}
+def Xls_FlopKind : I32EnumAttr<"FlopKind", "XLS flop kind",
+  [Xls_FlopKindNone, Xls_FlopKindFlop, Xls_FlopKindSkid, Xls_FlopKindZeroLatency]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::xls";
+}
+def Xls_FlopKindAttr : EnumAttr<Xls_Dialect, Xls_FlopKind, "flop_kind"> {
+  let summary = "Channel flop/buffer kind.";
+  let description = [{
+    The inputs/outputs of procs/blocks can optionally be fitted with one of a
+    number of buffer types. This is specified on the channel connected to this
+    proc input/output, using the `input_flop_kind` and `output_flop_kind`
+    property.
+
+    See https://google.github.io/xls/codegen_options/#io-behavior for more information.
+  }];
+}
+
+def Xls_FifoConfig : Xls_Attr<"FifoConfig"> {
+  let summary = "FIFO Configuration Attribute";
+  let description = [{
+    Definies the properties of the physical FIFO to be instantiated from
+    an (internal) channel.
+  }];
+
+  let parameters = (ins
+    "int64_t":$fifo_depth,
+    "bool":$bypass,
+    "bool":$register_push_outputs,
+    "bool":$register_pop_outputs
+  );
+
+  let mnemonic = "fifo_config";
+
+  let assemblyFormat = "`<` struct(params) `>`";
+}
+
 class GetShapeSplat<string name> :
     StrFunc<"getShapeSplat($" # name # ".getType())">;
 
@@ -1626,6 +1695,9 @@ def Xls_ChanOp : Xls_Op<"chan", [
   let arguments = (ins
     SymbolNameAttr:$sym_name,
     TypeAttr:$type,
+    OptionalAttr<Xls_FifoConfig>: $fifo_config,
+    OptionalAttr<Xls_FlopKindAttr>: $input_flop_kind,
+    OptionalAttr<Xls_FlopKindAttr>: $output_flop_kind,
     DefaultValuedAttr<BoolAttr, "true">:$send_supported,
     DefaultValuedAttr<BoolAttr, "true">:$recv_supported
   );
@@ -1653,7 +1725,6 @@ def Xls_ChanOp : Xls_Op<"chan", [
     void setArgAttrsAttr(ArrayAttr) { return; }
     void setResAttrsAttr(ArrayAttr) { return; }
   }];
-
 }
 
 def Xls_InstantiateEprocOp : Xls_Op<"instantiate_eproc", [DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
@@ -1763,7 +1834,10 @@ def Xls_SchanOp : Xls_Op<"schan", [Pure, HasParent<"SprocOp">]> {
   }];
   let arguments = (ins
     StrAttr:$name,
-    TypeAttr:$type
+    TypeAttr:$type,
+    OptionalAttr<Xls_FifoConfig>: $fifo_config,
+    OptionalAttr<Xls_FlopKindAttr>: $input_flop_kind,
+    OptionalAttr<Xls_FlopKindAttr>: $output_flop_kind
   );
   let results = (outs
     Xls_SchanType:$out,
@@ -1774,7 +1848,8 @@ def Xls_SchanOp : Xls_Op<"schan", [Pure, HasParent<"SprocOp">]> {
     OpBuilder<(ins "::mlir::StringRef":$name, "::mlir::Type":$element_type), [{
       auto inChanTy = SchanType::get($_builder.getContext(), element_type, /*isInput=*/true);
       auto outChanTy = SchanType::get($_builder.getContext(), element_type, /*isInput=*/false);
-      build($_builder, $_state, ::mlir::TypeRange{outChanTy, inChanTy}, name, element_type);
+      // TODO(jpienaar): Remove unecessary default args once they are generated automatically.
+      build($_builder, $_state, ::mlir::TypeRange{outChanTy, inChanTy}, name, element_type, /*fifo_config=*/nullptr, /*input_flop_kind=*/nullptr, /*output_flop_kind=*/nullptr);
     }]>
   ];
 }

xls/contrib/mlir/testdata/array_to_bits.mlir

@@ -215,8 +215,8 @@ func.func @tensor_empty() -> !xls.array<4 x i32> attributes {xls = true} {
   return %0 : !xls.array<4 x i32>
 }
 
-// CHECK-LABEL: xls.chan @mychan : i24
-xls.chan @mychan : !xls.array<3 x i8>
+// CHECK-LABEL: xls.chan @mychan {fifo_config = #xls.fifo_config<fifo_depth = 1, bypass = true, register_push_outputs = true, register_pop_outputs = false>} : i24
+xls.chan @mychan {fifo_config = #xls.fifo_config<fifo_depth = 1, bypass = true, register_push_outputs = true, register_pop_outputs = false>} : !xls.array<3 x i8>
 
 // CHECK-LABEL:   xls.eproc @eproc(
 // CHECK-SAME:                     %[[VAL_0:.*]]: i32) zeroinitializer {

xls/contrib/mlir/testdata/index_type_conversion.mlir

@@ -164,7 +164,7 @@ func.func @forloop(%arg0: i32, %arg1: i8, %arg2: i9) -> i32 attributes {xls = tr
   return %0 : i32
 }
 
-// INDEX32-LABEL:   xls.chan @mychan : i32
-// INDEX64-LABEL:   xls.chan @mychan : i64
-xls.chan @mychan : index
+// INDEX32-LABEL:   xls.chan @mychan {fifo_config = #xls.fifo_config<fifo_depth = 1, bypass = true, register_push_outputs = true, register_pop_outputs = false>} : i32
+// INDEX64-LABEL:   xls.chan @mychan {fifo_config = #xls.fifo_config<fifo_depth = 1, bypass = true, register_push_outputs = true, register_pop_outputs = false>} : i64
+xls.chan @mychan {fifo_config = #xls.fifo_config<fifo_depth = 1, bypass = true, register_push_outputs = true, register_pop_outputs = false>} : index
 

xls/contrib/mlir/testdata/proc_elaboration.mlir

@@ -1,6 +1,6 @@
 // RUN: xls_opt -elaborate-procs -split-input-file %s 2>&1 | FileCheck %s
-// CHECK-LABEL:       xls.chan @req : i32
-// CHECK-NEXT:  xls.chan @resp : i32
+// CHECK-LABEL: xls.chan @req : i32
+// CHECK-NEXT:  xls.chan @resp {fifo_config = #xls.fifo_config<fifo_depth = 1, bypass = true, register_push_outputs = true, register_pop_outputs = false>} : i32
 // CHECK-NEXT:  xls.chan @rom1_req : i32
 // CHECK-NEXT:  xls.chan @rom1_resp : i32
 // CHECK-NEXT:  xls.eproc @rom(%arg0: i32) zeroinitializer discardable {
@@ -45,7 +45,7 @@
 xls.sproc @fetch() top {
   spawns {
     %req_out, %req_in = xls.schan<i32>("req")
-    %resp_out, %resp_in = xls.schan<i32>("resp")
+    %resp_out, %resp_in = xls.schan<i32>("resp") attributes { fifo_config = #xls.fifo_config<fifo_depth = 1, bypass = true, register_push_outputs = true, register_pop_outputs = false> }
     xls.spawn @proxy(%req_in, %resp_out) : !xls.schan<i32, in>, !xls.schan<i32, out>
     xls.yield %req_out, %resp_in : !xls.schan<i32, out>, !xls.schan<i32, in>
   }

xls/contrib/mlir/testdata/translate_chn.mlir

@@ -0,0 +1,27 @@
+// RUN: xls_translate --mlir-xls-to-xls %s -- 2>&1 | FileCheck %s
+
+// CHECK-LABEL: chan ch_inp
+// CHECK-SAME:    kind=streaming
+// CHECK-SAME:    ops=receive_only
+xls.chan @ch_inp {send_supported = false} : i32
+
+// CHECK-LABEL: chan ch_out
+// CHECK-SAME:    kind=streaming
+// CHECK-SAME:    ops=send_only
+// CHECK-SAME:    fifo_depth=1
+// CHECK-SAME:    bypass=true
+// CHECK-SAME:    register_push_outputs=true
+// CHECK-SAME:    register_pop_outputs=false
+xls.chan @ch_out {
+  fifo_config = #xls.fifo_config<fifo_depth = 1, bypass = true, register_push_outputs = true, register_pop_outputs = false>,
+  recv_supported = false
+} : i32
+
+// CHECK: top proc eproc
+xls.eproc @eproc() zeroinitializer {
+  %tok0 = xls.after_all : !xls.token
+  %tok1, %val = xls.blocking_receive %tok0, @ch_inp : i32
+  xls.send %tok1, %val, @ch_out : i32
+  xls.yield
+}
+

xls/contrib/mlir/testdata/translate_from_xls/proc.ir

@@ -7,8 +7,13 @@ file_number 0 "./simple_proc.x"
 // CHECK-LABEL: xls.chan @ch_inp {send_supported = false} : i32
 chan ch_inp(bits[32], id=0, kind=streaming, ops=receive_only)
 
-// CHECK-LABEL: xls.chan @ch_out {recv_supported = false} : i32
-chan ch_out(bits[32], id=1, kind=streaming, ops=send_only)
+// CHECK-LABEL: xls.chan @ch_out {
+// CHECK-SAME:    fifo_config = #xls.fifo_config<fifo_depth = 1, bypass = true, register_push_outputs = true, register_pop_outputs = false>,
+// CHECK-SAME:    input_flop_kind = #xls<flop_kind zero-latency>,
+// CHECK-SAME:    output_flop_kind = #xls<flop_kind skid>
+// CHECK-SAME:    recv_supported = false
+// CHECK-SAME:  } : i32
+chan ch_out(bits[32], id=1, kind=streaming, ops=send_only, flow_control=ready_valid, strictness=proven_mutually_exclusive, fifo_depth=1, bypass=true, register_push_outputs=true, register_pop_outputs=false, input_flop_kind=zero_latency, output_flop_kind=skid)
 
 // CHECK-LABEL: xls.eproc @ident() zeroinitializer {
 // CHECK:         xls.yield

xls/contrib/mlir/tools/xls_translate/xls_translate_from_mlir.cc

@@ -68,6 +68,7 @@
 #include "xls/common/file/get_runfile_path.h"
 #include "xls/contrib/mlir/IR/xls_ops.h"
 #include "xls/ir/bits.h"
+#include "xls/ir/channel.h"
 #include "xls/ir/foreign_function.h"
 #include "xls/ir/foreign_function_data.pb.h"
 #include "xls/ir/nodes.h"
@@ -1409,6 +1410,19 @@ absl::StatusOr<::xls::Function*> wrapDslxFunctionIfNeeded(
   return xlsFunc;
 }
 
+::xls::FlopKind convertFlopKind(FlopKind kind) {
+  switch (kind) {
+    case FlopKind::kNone:
+      return ::xls::FlopKind::kNone;
+    case FlopKind::kFlop:
+      return ::xls::FlopKind::kFlop;
+    case FlopKind::kSkid:
+      return ::xls::FlopKind::kSkid;
+    case FlopKind::kZeroLatency:
+      return ::xls::FlopKind::kZeroLatency;
+  }
+}
+
 FailureOr<std::unique_ptr<Package>> mlirXlsToXls(
     Operation* op, llvm::StringRef dslx_search_path,
     DslxPackageCache& dslx_cache) {
@@ -1453,7 +1467,30 @@ FailureOr<std::unique_ptr<Package>> mlirXlsToXls(
     } else if (!chan_op.getRecvSupported()) {
       kind = ::xls::ChannelOps::kSendOnly;  // NOLINT
     }
-    auto channel = package->CreateStreamingChannel(name, kind, xls_type);
+
+    std::optional<::xls::FifoConfig> fifo_config = std::nullopt;
+    if (auto mlir_fifo_config = chan_op.getFifoConfig()) {
+      fifo_config = ::xls::FifoConfig(
+          mlir_fifo_config->getFifoDepth(), mlir_fifo_config->getBypass(),
+          mlir_fifo_config->getRegisterPushOutputs(),
+          mlir_fifo_config->getRegisterPopOutputs());
+    }
+
+    std::optional<::xls::FlopKind> input_flop = std::nullopt;
+    std::optional<::xls::FlopKind> output_flop = std::nullopt;
+    if (auto attr = chan_op.getInputFlopKind()) {
+      input_flop = convertFlopKind(*attr);
+    }
+    if (auto attr = chan_op.getOutputFlopKind()) {
+      output_flop = convertFlopKind(*attr);
+    }
+
+    auto channel_config =
+        ::xls::ChannelConfig(fifo_config, input_flop, output_flop);
+
+    auto channel = package->CreateStreamingChannel(name, kind, xls_type, {},
+                                                   channel_config);
+
     if (!channel.ok()) {
       chan_op.emitOpError("failed to create streaming channel: ")
           << channel.status().message();

xls/contrib/mlir/tools/xls_translate/xls_translate_to_mlir.cc

@@ -45,6 +45,7 @@
 #include "xls/ir/bits.h"
 #include "xls/ir/channel.h"
 #include "xls/ir/fileno.h"
+#include "xls/ir/channel.h"
 #include "xls/ir/function.h"
 #include "xls/ir/function_base.h"
 #include "xls/ir/lsb_or_msb.h"
@@ -1600,15 +1601,50 @@ absl::StatusOr<Operation*> translateProc(::xls::Proc& xls_proc,
 // Channel Translation
 //===----------------------------------------------------------------------===//
 
+FlopKind convertFlopKind(::xls::FlopKind kind) {
+  switch (kind) {
+    case ::xls::FlopKind::kNone:
+      return FlopKind::kNone;
+    case ::xls::FlopKind::kFlop:
+      return FlopKind::kFlop;
+    case ::xls::FlopKind::kSkid:
+      return FlopKind::kSkid;
+    case ::xls::FlopKind::kZeroLatency:
+      return FlopKind::kZeroLatency;
+  }
+}
+
 absl::Status translateChannel(::xls::Channel& xls_chn, OpBuilder& builder,
                               MLIRContext* ctx, TranslationState& state) {
   auto chn = builder.create<xls::ChanOp>(
       builder.getUnknownLoc(),
       /*name=*/builder.getStringAttr(xls_chn.name()),
       /*type=*/TypeAttr::get(translateType(xls_chn.type(), builder, ctx)),
+      /*fifo_config=*/nullptr,
+      /*input_flop_kind=*/nullptr,
+      /*output_flop_kind=*/nullptr,
       /*send_supported=*/builder.getBoolAttr(xls_chn.CanSend()),
       /*recv_supported=*/builder.getBoolAttr(xls_chn.CanReceive()));
 
+  if (auto xls_streaming_chn =
+          dynamic_cast<::xls::StreamingChannel*>(&xls_chn)) {
+    auto xls_ch_config = xls_streaming_chn->channel_config();
+    if (auto xls_fifo_config = xls_ch_config.fifo_config()) {
+      chn.setFifoConfigAttr(FifoConfig::get(
+          ctx, xls_fifo_config->depth(), xls_fifo_config->bypass(),
+          xls_fifo_config->register_push_outputs(),
+          xls_fifo_config->register_pop_outputs()));
+    }
+
+    if (auto xls_input_flop = xls_ch_config.input_flop_kind()) {
+      chn.setInputFlopKind(convertFlopKind(*xls_input_flop));
+    }
+
+    if (auto xls_output_flop = xls_ch_config.output_flop_kind()) {
+      chn.setOutputFlopKind(convertFlopKind(*xls_output_flop));
+    }
+  }
+
   return state.setChannel(std::string(xls_chn.name()),
                           SymbolRefAttr::get(chn.getNameAttr()));
 }

xls/contrib/mlir/transforms/array_to_bits.cc

@@ -147,10 +147,10 @@ class LegalizeChanOpPattern : public OpConversionPattern<ChanOp> {
       ChanOp op, OpAdaptor adaptor,
       ConversionPatternRewriter& rewriter) const override {
     (void)adaptor;
-    auto newOp = rewriter.replaceOpWithNewOp<ChanOp>(
-        op, op.getSymName(), typeConverter->convertType(op.getType()));
-    newOp.setSendSupported(op.getSendSupported());
-    newOp.setRecvSupported(op.getRecvSupported());
+    rewriter.replaceOpWithNewOp<ChanOp>(
+        op, op.getSymName(), typeConverter->convertType(op.getType()),
+        op.getFifoConfigAttr(), op.getInputFlopKindAttr(), op.getOutputFlopKindAttr(),
+        op.getSendSupported(), op.getRecvSupported());
     return success();
   }
 };