[FIRRTL] LowerLayers: Inline enabled layers

Author: rwy7

include/circt/Dialect/FIRRTL/LayerSet.h

@@ -42,6 +42,43 @@ struct LayerSetCompare {
 
 using LayerSet = SmallSet<SymbolRefAttr, 4, LayerSetCompare>;
 
+//===----------------------------------------------------------------------===//
+// Layer Verification Utilities
+//===----------------------------------------------------------------------===//
+
+/// Get the ambient layers active at the given op.
+LayerSet getAmbientLayersAt(Operation *op);
+
+/// Get the ambient layer requirements at the definition site of the value.
+LayerSet getAmbientLayersFor(Value value);
+
+/// Get the effective layer requirements for the given value.
+/// The effective layers for a value is the union of
+///   - the ambient layers for the cannonical storage location.
+///   - any explicit layer annotations in the value's type.
+LayerSet getLayersFor(Value value);
+
+/// Check that the source layer is compatible with the destination layer.
+/// Either the source and destination are identical, or the source-layer
+/// is a parent of the destination. For example `A` is compatible with `A.B.C`,
+/// because any definition valid in `A` is also valid in `A.B.C`.
+bool isLayerCompatibleWith(mlir::SymbolRefAttr srcLayer,
+                           mlir::SymbolRefAttr dstLayer);
+
+/// Check that the source layer is present in the destination layers.
+bool isLayerCompatibleWith(SymbolRefAttr srcLayer, const LayerSet &dstLayers);
+
+/// Check that the source layers are all present in the destination layers.
+/// True if all source layers are present in the destination.
+/// Outputs the set of source layers that are missing in the destination.
+bool isLayerSetCompatibleWith(const LayerSet &src, const LayerSet &dst,
+                              SmallVectorImpl<SymbolRefAttr> &missing);
+
+LogicalResult checkLayerCompatibility(
+    Operation *op, const LayerSet &src, const LayerSet &dst,
+    const Twine &errorMsg,
+    const Twine &noteMsg = Twine("missing layer requirements"));
+
 } // namespace firrtl
 } // namespace circt
 

lib/Dialect/FIRRTL/CMakeLists.txt

@@ -19,6 +19,7 @@ set(CIRCT_FIRRTL_Sources
   FIRRTLOps.cpp
   FIRRTLTypes.cpp
   FIRRTLUtils.cpp
+  LayerSet.cpp
   NLATable.cpp
 )
 

lib/Dialect/FIRRTL/FIRRTLOps.cpp

@@ -313,112 +313,6 @@ void getAsmBlockArgumentNamesImpl(Operation *op, mlir::Region &region,
 static ParseResult parseNameKind(OpAsmParser &parser,
                                  firrtl::NameKindEnumAttr &result);
 
-//===----------------------------------------------------------------------===//
-// Layer Verification Utilities
-//===----------------------------------------------------------------------===//
-
-/// Get the ambient layers active at the given op.
-static LayerSet getAmbientLayersAt(Operation *op) {
-  // Crawl through the parent ops, accumulating all ambient layers at the given
-  // operation.
-  LayerSet result;
-  for (; op != nullptr; op = op->getParentOp()) {
-    if (auto module = dyn_cast<FModuleLike>(op)) {
-      auto layers = module.getLayersAttr().getAsRange<SymbolRefAttr>();
-      result.insert(layers.begin(), layers.end());
-      break;
-    }
-    if (auto layerblock = dyn_cast<LayerBlockOp>(op)) {
-      result.insert(layerblock.getLayerName());
-      continue;
-    }
-  }
-  return result;
-}
-
-/// Get the ambient layer requirements at the definition site of the value.
-static LayerSet getAmbientLayersFor(Value value) {
-  return getAmbientLayersAt(getFieldRefFromValue(value).getDefiningOp());
-}
-
-/// Get the effective layer requirements for the given value.
-/// The effective layers for a value is the union of
-///   - the ambient layers for the cannonical storage location.
-///   - any explicit layer annotations in the value's type.
-static LayerSet getLayersFor(Value value) {
-  auto result = getAmbientLayersFor(value);
-  if (auto type = dyn_cast<RefType>(value.getType()))
-    if (auto layer = type.getLayer())
-      result.insert(type.getLayer());
-  return result;
-}
-
-/// Check that the source layer is compatible with the destination layer.
-/// Either the source and destination are identical, or the source-layer
-/// is a parent of the destination. For example `A` is compatible with `A.B.C`,
-/// because any definition valid in `A` is also valid in `A.B.C`.
-static bool isLayerCompatibleWith(mlir::SymbolRefAttr srcLayer,
-                                  mlir::SymbolRefAttr dstLayer) {
-  // A non-colored probe may be cast to any colored probe.
-  if (!srcLayer)
-    return true;
-
-  // A colored probe cannot be cast to an uncolored probe.
-  if (!dstLayer)
-    return false;
-
-  // Return true if the srcLayer is a prefix of the dstLayer.
-  if (srcLayer.getRootReference() != dstLayer.getRootReference())
-    return false;
-
-  auto srcNames = srcLayer.getNestedReferences();
-  auto dstNames = dstLayer.getNestedReferences();
-  if (dstNames.size() < srcNames.size())
-    return false;
-
-  return llvm::all_of(llvm::zip_first(srcNames, dstNames),
-                      [](auto x) { return std::get<0>(x) == std::get<1>(x); });
-}
-
-/// Check that the source layer is present in the destination layers.
-static bool isLayerCompatibleWith(SymbolRefAttr srcLayer,
-                                  const LayerSet &dstLayers) {
-  // fast path: the required layer is directly listed in the provided layers.
-  if (dstLayers.contains(srcLayer))
-    return true;
-
-  // Slow path: the required layer is not directly listed in the provided
-  // layers, but the layer may still be provided by a nested layer.
-  return any_of(dstLayers, [=](SymbolRefAttr dstLayer) {
-    return isLayerCompatibleWith(srcLayer, dstLayer);
-  });
-}
-
-/// Check that the source layers are all present in the destination layers.
-/// True if all source layers are present in the destination.
-/// Outputs the set of source layers that are missing in the destination.
-static bool isLayerSetCompatibleWith(const LayerSet &src, const LayerSet &dst,
-                                     SmallVectorImpl<SymbolRefAttr> &missing) {
-  for (auto srcLayer : src)
-    if (!isLayerCompatibleWith(srcLayer, dst))
-      missing.push_back(srcLayer);
-
-  llvm::sort(missing, LayerSetCompare());
-  return missing.empty();
-}
-
-static LogicalResult checkLayerCompatibility(
-    Operation *op, const LayerSet &src, const LayerSet &dst,
-    const Twine &errorMsg,
-    const Twine &noteMsg = Twine("missing layer requirements")) {
-  SmallVector<SymbolRefAttr> missing;
-  if (isLayerSetCompatibleWith(src, dst, missing))
-    return success();
-  interleaveComma(missing, op->emitOpError(errorMsg).attachNote()
-                               << noteMsg << ": ");
-  return failure();
-}
-
 //===----------------------------------------------------------------------===//
 // CircuitOp
 //===----------------------------------------------------------------------===//

lib/Dialect/FIRRTL/LayerSet.cpp

@@ -0,0 +1,98 @@
+#include "circt/Dialect/FIRRTL/LayerSet.h"
+#include "circt/Dialect/FIRRTL/FIRRTLOps.h"
+#include "circt/Dialect/FIRRTL/FIRRTLUtils.h"
+#include "llvm/ADT/SmallSet.h"
+
+using namespace mlir;
+using namespace circt;
+using namespace firrtl;
+
+LayerSet circt::firrtl::getAmbientLayersAt(Operation *op) {
+  // Crawl through the parent ops, accumulating all ambient layers at the given
+  // operation.
+  LayerSet result;
+  for (; op != nullptr; op = op->getParentOp()) {
+    if (auto module = dyn_cast<FModuleLike>(op)) {
+      auto layers = module.getLayersAttr().getAsRange<SymbolRefAttr>();
+      result.insert(layers.begin(), layers.end());
+      break;
+    }
+    if (auto layerblock = dyn_cast<LayerBlockOp>(op)) {
+      result.insert(layerblock.getLayerName());
+      continue;
+    }
+  }
+  return result;
+}
+
+LayerSet circt::firrtl::getAmbientLayersFor(Value value) {
+  return getAmbientLayersAt(getFieldRefFromValue(value).getDefiningOp());
+}
+
+LayerSet circt::firrtl::getLayersFor(Value value) {
+  auto result = getAmbientLayersFor(value);
+  if (auto type = dyn_cast<RefType>(value.getType()))
+    if (auto layer = type.getLayer())
+      result.insert(type.getLayer());
+  return result;
+}
+
+bool circt::firrtl::isLayerCompatibleWith(mlir::SymbolRefAttr srcLayer,
+                                          mlir::SymbolRefAttr dstLayer) {
+  // A non-colored probe may be cast to any colored probe.
+  if (!srcLayer)
+    return true;
+
+  // A colored probe cannot be cast to an uncolored probe.
+  if (!dstLayer)
+    return false;
+
+  // Return true if the srcLayer is a prefix of the dstLayer.
+  if (srcLayer.getRootReference() != dstLayer.getRootReference())
+    return false;
+
+  auto srcNames = srcLayer.getNestedReferences();
+  auto dstNames = dstLayer.getNestedReferences();
+  if (dstNames.size() < srcNames.size())
+    return false;
+
+  return llvm::all_of(llvm::zip_first(srcNames, dstNames),
+                      [](auto x) { return std::get<0>(x) == std::get<1>(x); });
+}
+
+bool circt::firrtl::isLayerCompatibleWith(SymbolRefAttr srcLayer,
+                                          const LayerSet &dstLayers) {
+  // fast path: the required layer is directly listed in the provided layers.
+  if (dstLayers.contains(srcLayer))
+    return true;
+
+  // Slow path: the required layer is not directly listed in the provided
+  // layers, but the layer may still be provided by a nested layer.
+  return any_of(dstLayers, [=](SymbolRefAttr dstLayer) {
+    return isLayerCompatibleWith(srcLayer, dstLayer);
+  });
+}
+
+bool circt::firrtl::isLayerSetCompatibleWith(
+    const LayerSet &src, const LayerSet &dst,
+    SmallVectorImpl<SymbolRefAttr> &missing) {
+  for (auto srcLayer : src)
+    if (!isLayerCompatibleWith(srcLayer, dst))
+      missing.push_back(srcLayer);
+
+  llvm::sort(missing, LayerSetCompare());
+  return missing.empty();
+}
+
+LogicalResult circt::firrtl::checkLayerCompatibility(Operation *op,
+                                                     const LayerSet &src,
+                                                     const LayerSet &dst,
+                                                     const Twine &errorMsg,
+                                                     const Twine &noteMsg) {
+  SmallVector<SymbolRefAttr> missing;
+  if (isLayerSetCompatibleWith(src, dst, missing))
+    return success();
+  interleaveComma(missing, op->emitOpError(errorMsg).attachNote()
+                               << noteMsg << ": ");
+  return failure();
+}

lib/Dialect/FIRRTL/Transforms/LowerLayers.cpp

@@ -13,6 +13,7 @@
 #include "circt/Dialect/FIRRTL/FIRRTLInstanceGraph.h"
 #include "circt/Dialect/FIRRTL/FIRRTLOps.h"
 #include "circt/Dialect/FIRRTL/FIRRTLUtils.h"
+#include "circt/Dialect/FIRRTL/LayerSet.h"
 #include "circt/Dialect/FIRRTL/Namespace.h"
 #include "circt/Dialect/FIRRTL/Passes.h"
 #include "circt/Dialect/HW/HierPathCache.h"
@@ -22,6 +23,7 @@
 #include "mlir/Pass/Pass.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Mutex.h"
 
@@ -231,7 +233,8 @@ class LowerLayersPass
   FailureOr<InnerRefMap> runOnModuleLike(FModuleLike moduleLike);
 
   /// Extract layerblocks and strip probe colors from all ops under the module.
-  LogicalResult runOnModuleBody(FModuleOp moduleOp, InnerRefMap &innerRefMap);
+  LogicalResult runOnModuleBody(FModuleOp moduleOp, InnerRefMap &innerRefMap,
+                                const LayerSet &enabledLayers);
 
   /// Update the module's port types to remove any explicit layer requirements
   /// from any probe types.
@@ -243,6 +246,9 @@ class LowerLayersPass
   /// Lower an inline layerblock to an ifdef block.
   void lowerInlineLayerBlock(LayerOp layer, LayerBlockOp layerBlock);
 
+  /// Inline an enabled layerblock directly into its parent.
+  void inlineLayerBlock(LayerBlockOp layerBlock);
+
   /// Build macro declarations and cache information about the layers.
   void preprocessLayers(CircuitNamespace &ns, OpBuilder &b, LayerOp layer,
                         StringRef circuitName,
@@ -364,9 +370,12 @@ LowerLayersPass::runOnModuleLike(FModuleLike moduleLike) {
   auto result =
       TypeSwitch<Operation *, LogicalResult>(moduleLike.getOperation())
           .Case<FModuleOp>([&](auto op) {
+            LayerSet enabledLayers;
+            enabledLayers.insert_range(
+                op.getLayersAttr().template getAsRange<SymbolRefAttr>());
             op.setLayers({});
             removeLayersFromPorts(op);
-            return runOnModuleBody(op, innerRefMap);
+            return runOnModuleBody(op, innerRefMap, enabledLayers);
           })
           .Case<FExtModuleOp>([&](auto op) {
             op.setKnownLayers({});
@@ -400,8 +409,15 @@ void LowerLayersPass::lowerInlineLayerBlock(LayerOp layer,
   layerBlock.erase();
 }
 
+void LowerLayersPass::inlineLayerBlock(LayerBlockOp layerBlock) {
+  layerBlock->getBlock()->getOperations().splice(
+      Block::iterator(layerBlock), layerBlock.getBody()->getOperations());
+  layerBlock.erase();
+}
+
 LogicalResult LowerLayersPass::runOnModuleBody(FModuleOp moduleOp,
-                                               InnerRefMap &innerRefMap) {
+                                               InnerRefMap &innerRefMap,
+                                               const LayerSet &enabledLayers) {
   hw::InnerSymbolNamespace ns(moduleOp);
 
   // Get or create a node op for a value captured by a layer block.
@@ -694,6 +710,11 @@ LogicalResult LowerLayersPass::runOnModuleBody(FModuleOp moduleOp,
     // After this point, we are dealing with a layer block.
     auto layer = symbolToLayer.lookup(layerBlock.getLayerName());
 
+    if (isLayerCompatibleWith(layerBlock.getLayerName(), enabledLayers)) {
+      inlineLayerBlock(layerBlock);
+      return WalkResult::advance();
+    }
+
     if (layer.getConvention() == LayerConvention::Inline) {
       lowerInlineLayerBlock(layer, layerBlock);
       return WalkResult::advance();