NOT YET MERGED: Expose FifoProps in MLIR

google#1916

xls/contrib/mlir/IR/xls_ops.cc

@@ -620,6 +620,11 @@ 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));
+
+  if (parser.parseOptionalAttrDictWithKeyword(result.attributes)) {
+    return failure();
+  }
+
   return success();
 }
 
@@ -630,6 +635,8 @@ 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,33 @@ def Xls_TranslationLinkage : Xls_Attr<"TranslationLinkage"> {
   let attrName = "xls.translation_linkage";
 }
 
+def Xls_FlopKindNone        : I32EnumAttrCase<"kNone", 0, "none">;
+def Xls_FlopKindFlop        : I32EnumAttrCase<"kFlop", 1, "flop">;
+def Xls_FlopKindSkid        : I32EnumAttrCase<"kSkid", 2, "skid">;
+def Xls_FlopKindZeroLatency : I32EnumAttrCase<"kZeroLatency", 3, "zero-latency">;
+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">;
+
+def Xls_FifoConfig : Xls_Attr<"FifoConfig"> {
+  let summary = "TODO";
+  let description = [{TODO}];
+
+  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())">;
 

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();
   }
 };

xls/contrib/mlir/transforms/index_type_conversion.cc

@@ -246,9 +246,10 @@ class LegalizeChanOp : public OpConversionPattern<ChanOp> {
       ChanOp op, OpAdaptor /*adaptor*/,
       ConversionPatternRewriter &rewriter) const override {
     Type resultType = getTypeConverter()->convertType(op.getType());
-    rewriter.replaceOpWithNewOp<xls::ChanOp>(op, op.getSymName(), resultType,
-                                             op.getSendSupported(),
-                                             op.getRecvSupported());
+    rewriter.replaceOpWithNewOp<xls::ChanOp>(
+        op, op.getSymName(), resultType, op.getFifoConfigAttr(),
+        op.getInputFlopKindAttr(), op.getOutputFlopKindAttr(), op.getSendSupported(),
+        op.getRecvSupported());
     return success();
   }
 };

xls/contrib/mlir/transforms/proc_elaboration.cc

@@ -141,7 +141,10 @@ class ElaborationContext
     for (auto [i, arg] : llvm::enumerate(sproc.getNextChannels())) {
       auto chan = builder.create<ChanOp>(
           sproc.getLoc(), absl::StrFormat("%s_arg%d", originalName.str(), i),
-          cast<SchanType>(arg.getType()).getElementType());
+          cast<SchanType>(arg.getType()).getElementType(),
+          /*fifo_config=*/nullptr,
+          /*input_flop_kind=*/nullptr,
+          /*output_flop_kind=*/nullptr);
       eprocChannels.push_back(chan);
       chanMap[mapping.lookup(arg)] = SymbolRefAttr::get(chan.getSymNameAttr());
     }
@@ -224,8 +227,10 @@ class ElaborationInterpreter
   absl::Status Interpret(SchanOp op, ElaborationContext& ctx) {
     std::string name = op.getName().str();
     auto uniqueName = ctx.Uniquify(op.getNameAttr());
-    ChanOp chan =
-        ctx.getBuilder().create<ChanOp>(op.getLoc(), uniqueName, op.getType());
+    ChanOp chan = ctx.getBuilder().create<ChanOp>(
+        op.getLoc(), uniqueName, op.getType(), op.getFifoConfigAttr(),
+        op.getInputFlopKindAttr(), op.getOutputFlopKindAttr());
+
     ctx.getSymbolTable().insert(chan);
     ctx.Set(op.getResult(0), chan);
     ctx.Set(op.getResult(1), chan);
@@ -318,7 +323,11 @@ void ProcElaborationPass::runOnOperation() {
         SchanType schan = cast<SchanType>(arg.getType());
         auto nameAttr = cast<StringAttr>(name);
         auto echan = builder.create<ChanOp>(sproc.getLoc(), nameAttr.str(),
-                                            schan.getElementType());
+                                            schan.getElementType(),
+                                            /*fifo_config=*/nullptr,
+                                            /*input_flop_kind=*/nullptr,
+                                            /*output_flop_kind=*/nullptr);
+
         if (schan.getIsInput()) {
           echan.setSendSupported(false);
         } else {

xls/contrib/mlir/transforms/scalarize.cc

@@ -546,10 +546,11 @@ 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();
   }
 };