[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.
google · Feb 25, 2025 · 0f4e000 · 0f4e000
1 parent 49cff5c
commit 0f4e000
Show file tree

Hide file tree

Showing 13 changed files with 187 additions and 29 deletions.
diff --git a/xls/contrib/mlir/IR/xls_ops.cc b/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) {

diff --git a/xls/contrib/mlir/IR/xls_ops.td b/xls/contrib/mlir/IR/xls_ops.td
@@ -135,6 +135,36 @@ 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 = "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 +1656,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 +1686,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 +1795,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 +1809,7 @@ 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);
+      build($_builder, $_state, ::mlir::TypeRange{outChanTy, inChanTy}, name, element_type, /*fifo_config=*/nullptr, /*input_flop_kind=*/nullptr, /*output_flop_kind=*/nullptr);
     }]>
   ];
 }

diff --git a/xls/contrib/mlir/testdata/array_to_bits.mlir b/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 {

diff --git a/xls/contrib/mlir/testdata/index_type_conversion.mlir b/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
 
diff --git a/xls/contrib/mlir/testdata/proc_elaboration.mlir b/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>
   }

diff --git a/xls/contrib/mlir/testdata/translate_chn.mlir b/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
+}
+
diff --git a/xls/contrib/mlir/testdata/translate_from_xls/proc.ir b/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

diff --git a/xls/contrib/mlir/tools/xls_translate/xls_translate_from_mlir.cc b/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"
@@ -1408,6 +1409,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) {
@@ -1452,7 +1466,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();

diff --git a/xls/contrib/mlir/tools/xls_translate/xls_translate_to_mlir.cc b/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()));
 }

diff --git a/xls/contrib/mlir/transforms/array_to_bits.cc b/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();
   }
 };

diff --git a/xls/contrib/mlir/transforms/index_type_conversion.cc b/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();
   }
 };

diff --git a/xls/contrib/mlir/transforms/proc_elaboration.cc b/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 {

diff --git a/xls/contrib/mlir/transforms/scalarize.cc b/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();
   }
 };