CascadeFlowOp and ConfigureCascadeOp

include/aie/Dialect/AIE/IR/AIEAttrs.td

@@ -39,19 +39,29 @@ def AIEI64Attr : AIETypedSignlessIntegerAttrBase<
 // AIE attributes.
 //===----------------------------------------------------------------------===//
 
+def CoreWire:   I32EnumAttrCase<"Core", 0>;
+def DMAWire:    I32EnumAttrCase<"DMA", 1>;
+def FIFOWire:   I32EnumAttrCase<"FIFO", 2>;
+def SouthWire:  I32EnumAttrCase<"South", 3>;
+def WestWire:   I32EnumAttrCase<"West", 4>;
+def NorthWire:  I32EnumAttrCase<"North", 5>;
+def EastWire:   I32EnumAttrCase<"East", 6>;
+def PLIOWire:   I32EnumAttrCase<"PLIO", 7>;
+def NOCWire:    I32EnumAttrCase<"NOC", 8>;
+def TraceWire:  I32EnumAttrCase<"Trace", 9>;
 
 def WireBundle: I32EnumAttr<"WireBundle", "Bundle of wires",
   [
-    I32EnumAttrCase<"Core", 0>,
-    I32EnumAttrCase<"DMA", 1>,
-    I32EnumAttrCase<"FIFO", 2>,
-    I32EnumAttrCase<"South", 3>,
-    I32EnumAttrCase<"West", 4>,
-    I32EnumAttrCase<"North", 5>,
-    I32EnumAttrCase<"East", 6>,
-    I32EnumAttrCase<"PLIO", 7>,
-    I32EnumAttrCase<"NOC", 8>,
-    I32EnumAttrCase<"Trace", 9>
+    CoreWire, DMAWire, FIFOWire, SouthWire, WestWire, NorthWire,
+    EastWire, PLIOWire, NOCWire, TraceWire
+  ]> {
+
+  let cppNamespace = "xilinx::AIE";
+}
+
+def CascadeDir: I32EnumAttr<"CascadeDir", "Directions for cascade",
+  [
+    SouthWire, WestWire, NorthWire, EastWire
   ]> {
 
   let cppNamespace = "xilinx::AIE";

include/aie/Dialect/AIE/IR/AIEOps.td

@@ -1345,6 +1345,57 @@ def AIE_PutStreamOp: AIE_Op<"put_stream", [HasParent<"CoreOp">]> {
   }];
 }
 
+def AIE_CascadeFlowOp: AIE_Op<"cascade_flow", []> {
+  let arguments = (
+    ins Index:$source_tile,
+        Index:$dest_tile
+  );
+  let summary = "A cascade connection between tiles";
+  let description = [{
+    The `aie.cascade_flow` operation represents a cascade connection between two `aie.tile` operations.  
+    During lowering, this is replaced by `aie.configure_cascade` operations for each `aie.tile` based on 
+    their relative placement to one another.
+
+    Example:
+    ```
+      %tile03 = aie.tile(0, 3)
+      %tile13 = aie.tile(1, 3)
+      aie.cascade_flow(%tile03, %tile13)
+    ```
+  }];
+  let hasVerifier = 1;
+  let assemblyFormat = [{
+    `(` $source_tile `,` $dest_tile `)` attr-dict
+  }];
+  let extraClassDeclaration = [{
+    TileOp getSourceTileOp();
+    TileOp getDestTileOp();
+  }];
+}
+
+def AIE_ConfigureCascadeOp: AIE_Op<"configure_cascade", [HasParent<"DeviceOp">]> {
+  let summary = "An op to configure the input and output directions of the cascade for a single AIE tile";
+  let description = [{
+    An operation to configure the cascade on a single tile in both the input and the output 
+    directions.
+
+    Example:
+    ```
+      %tile00 = aie.tile(1, 3)
+      aie.configure_cascade(%tile00, West, East)
+    ```
+    Configures the input cascade port of %tile00 to the West direction, and the output port to the East direction.
+  }];
+  let arguments = (
+    ins Index:$tile,
+        CascadeDir:$inputDir,
+        CascadeDir:$outputDir
+  );
+  let results = (outs);
+  let hasVerifier = 1;
+  let assemblyFormat = [{ `(` $tile `,` $inputDir `,` $outputDir `)` attr-dict }];
+}
+
 def AIE_GetCascadeOp: AIE_Op<"get_cascade", [HasParent<"CoreOp">]>, Results<(outs AnyType:$cascade_value)> {
   let summary = "An op to read from a cascading stream from a neighboring core";
   let description = [{

include/aie/Dialect/AIE/Transforms/AIEPasses.h

@@ -41,6 +41,7 @@ std::unique_ptr<mlir::OperationPass<DeviceOp>>
 createAIEObjectFifoStatefulTransformPass();
 std::unique_ptr<mlir::OperationPass<DeviceOp>>
 createAIEObjectFifoRegisterProcessPass();
+std::unique_ptr<mlir::OperationPass<DeviceOp>> createAIELowerCascadeFlowsPass();
 
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION

include/aie/Dialect/AIE/Transforms/AIEPasses.td

@@ -211,4 +211,17 @@ def AIEObjectFifoRegisterProcess : Pass<"aie-register-objectFifos", "DeviceOp">
   ];
 }
 
+def AIELowerCascadeFlows : Pass<"aie-lower-cascade-flows", "DeviceOp"> {
+  let summary = "Lower aie.cascade_flow operations through `aie.configure_cascade` operations";
+  let description = [{
+    Replace each aie.cascade_flow operation with an equivalent set of `aie.configure_cascade` 
+    operations. 
+  }];
+
+  let constructor = "xilinx::AIE::createAIELowerCascadeFlowsPass()";
+  let dependentDialects = [
+    "xilinx::AIE::AIEDialect",
+  ];
+}
+
 #endif

lib/Dialect/AIE/IR/AIEDialect.cpp

@@ -864,6 +864,70 @@ ObjectFifoCreateOp ObjectFifoRegisterProcessOp::getObjectFifo() {
   return {};
 }
 
+//===----------------------------------------------------------------------===//
+// CascadeFlowOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult CascadeFlowOp::verify() {
+  TileOp src = getSourceTileOp();
+  TileOp dst = getDestTileOp();
+  const auto &t = getTargetModel(src);
+
+  if (src.isShimTile() || dst.isShimTile())
+    return emitOpError("shimTile row has no cascade stream interface");
+  if (t.isMemTile(src.colIndex(), src.rowIndex()) ||
+      t.isMemTile(dst.colIndex(), dst.rowIndex()))
+    return emitOpError("memTile row has no cascade stream interface");
+
+  if (!t.isSouth(src.getCol(), src.getRow(), dst.getCol(), dst.getRow()) &&
+      !t.isWest(src.getCol(), src.getRow(), dst.getCol(), dst.getRow()) &&
+      !t.isNorth(src.getCol(), src.getRow(), dst.getCol(), dst.getRow()) &&
+      !t.isEast(src.getCol(), src.getRow(), dst.getCol(), dst.getRow())) {
+    return emitOpError("tiles must be adjacent");
+  }
+  return success();
+}
+
+TileOp CascadeFlowOp::getSourceTileOp() {
+  return cast<TileOp>(getSourceTile().getDefiningOp());
+}
+
+TileOp CascadeFlowOp::getDestTileOp() {
+  return cast<TileOp>(getDestTile().getDefiningOp());
+}
+
+//===----------------------------------------------------------------------===//
+// ConfigureCascadeOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult ConfigureCascadeOp::verify() {
+  const auto &t = getTargetModel(*this);
+  TileOp tile = cast<TileOp>(getTile().getDefiningOp());
+  CascadeDir inputDir = getInputDir();
+  CascadeDir outputDir = getOutputDir();
+
+  if (tile.isShimTile())
+    return emitOpError("shimTile row has no cascade stream interface");
+  if (t.isMemTile(tile.colIndex(), tile.rowIndex()))
+    return emitOpError("memTile row has no cascade stream interface");
+
+  if (t.getTargetArch() == AIEArch::AIE2) {
+    if (inputDir == CascadeDir::South || inputDir == CascadeDir::East) {
+      return emitOpError("input direction of cascade must be North or West on ")
+             << stringifyAIEArch(t.getTargetArch());
+    }
+    if (outputDir == CascadeDir::North || outputDir == CascadeDir::West) {
+      return emitOpError(
+                 "output direction of cascade must be South or East on ")
+             << stringifyAIEArch(t.getTargetArch());
+    }
+  } else {
+    return emitOpError("cascade not supported in ")
+           << stringifyAIEArch(t.getTargetArch());
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // PutCascadeOp
 //===----------------------------------------------------------------------===//

lib/Dialect/AIE/Transforms/AIECoreToStandard.cpp

@@ -570,7 +570,8 @@ struct AIECoreToStandardPass : AIECoreToStandardBase<AIECoreToStandardPass> {
         AIEOpRemoval<DeviceOp>, AIEOpRemoval<TileOp>, AIEOpRemoval<FlowOp>,
         AIEOpRemoval<MemOp>, AIEOpRemoval<ShimDMAOp>, AIEOpRemoval<ShimMuxOp>,
         AIEOpRemoval<SwitchboxOp>, AIEOpRemoval<LockOp>, AIEOpRemoval<BufferOp>,
-        AIEOpRemoval<ExternalBufferOp>, AIEOpRemoval<ShimDMAAllocationOp>>(
+        AIEOpRemoval<ExternalBufferOp>, AIEOpRemoval<ShimDMAAllocationOp>,
+        AIEOpRemoval<CascadeFlowOp>, AIEOpRemoval<ConfigureCascadeOp>>(
         m.getContext(), m);
 
     if (failed(applyPartialConversion(m, target, std::move(removepatterns))))

lib/Dialect/AIE/Transforms/AIELowerCascadeFlows.cpp

@@ -0,0 +1,96 @@
+//===- AIELowerCascadeFlows.cpp ---------------------------------*- C++ -*-===//
+//
+// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// (c) Copyright 2024 Xilinx Inc.
+//
+//===----------------------------------------------------------------------===//
+
+#include "aie/Dialect/AIE/IR/AIEDialect.h"
+#include "aie/Dialect/AIE/Transforms/AIEPasses.h"
+
+#include "mlir/IR/Attributes.h"
+#include "mlir/Pass/Pass.h"
+
+#include "llvm/ADT/Twine.h"
+
+#define DEBUG_TYPE "aie-lower-cascade-flows"
+
+using namespace mlir;
+using namespace xilinx;
+using namespace xilinx::AIE;
+
+struct AIELowerCascadeFlowsPass
+    : AIELowerCascadeFlowsBase<AIELowerCascadeFlowsPass> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<AIEDialect>();
+  }
+  void runOnOperation() override {
+    DeviceOp device = getOperation();
+    const auto &targetModel = device.getTargetModel();
+    OpBuilder builder = OpBuilder::atBlockEnd(device.getBody());
+
+    std::set<TileOp> tilesWithCascadeFlow;
+    DenseMap<TileOp, WireBundle> cascadeInputsPerTile;
+    DenseMap<TileOp, WireBundle> cascadeOutputsPerTile;
+
+    // identify cascade_flows and what ports they use on each tile
+    for (auto cascadeFlow : device.getOps<CascadeFlowOp>()) {
+      // for each cascade flow
+      TileOp src = cascadeFlow.getSourceTileOp();
+      TileOp dst = cascadeFlow.getDestTileOp();
+      tilesWithCascadeFlow.insert(src);
+      tilesWithCascadeFlow.insert(dst);
+
+      if (targetModel.isSouth(src.getCol(), src.getRow(), dst.getCol(),
+                              dst.getRow())) {
+        cascadeInputsPerTile[dst] = WireBundle::North;
+        cascadeOutputsPerTile[src] = WireBundle::South;
+      } else if (targetModel.isEast(src.getCol(), src.getRow(), dst.getCol(),
+                                    dst.getRow())) {
+        cascadeInputsPerTile[dst] = WireBundle::West;
+        cascadeOutputsPerTile[src] = WireBundle::East;
+      } else {
+        // TODO: remove when this pass supports routing
+        cascadeFlow.emitOpError(
+            "source tile must be to the North or West of the destination tile");
+        return;
+      }
+    }
+
+    // generate configure cascade ops
+    for (TileOp tile : tilesWithCascadeFlow) {
+      WireBundle inputDir;
+      if (cascadeInputsPerTile.find(tile) != cascadeInputsPerTile.end()) {
+        inputDir = cascadeInputsPerTile[tile];
+      } else {
+        inputDir = WireBundle::North;
+      }
+      WireBundle outputDir;
+      if (cascadeOutputsPerTile.find(tile) != cascadeOutputsPerTile.end()) {
+        outputDir = cascadeOutputsPerTile[tile];
+      } else {
+        outputDir = WireBundle::South;
+      }
+      builder.create<ConfigureCascadeOp>(builder.getUnknownLoc(), tile,
+                                         static_cast<CascadeDir>(inputDir),
+                                         static_cast<CascadeDir>(outputDir));
+    }
+
+    // erase CascadeFlowOps
+    SetVector<Operation *> opsToErase;
+    device.walk([&](Operation *op) {
+      if (isa<CascadeFlowOp>(op))
+        opsToErase.insert(op);
+    });
+    IRRewriter rewriter(&getContext());
+    for (auto it = opsToErase.rbegin(); it != opsToErase.rend(); ++it)
+      (*it)->erase();
+  }
+};
+
+std::unique_ptr<OperationPass<DeviceOp>> AIE::createAIELowerCascadeFlowsPass() {
+  return std::make_unique<AIELowerCascadeFlowsPass>();
+}

lib/Dialect/AIE/Transforms/CMakeLists.txt

@@ -20,6 +20,7 @@ add_mlir_dialect_library(
   AIEVectorOpt.cpp
   AIEObjectFifoStatefulTransform.cpp
   AIEObjectFifoRegisterProcess.cpp
+  AIELowerCascadeFlows.cpp
   ADDITIONAL_HEADER_DIRS
   ${AIE_BINARY_DIR}/include
 

lib/Targets/AIETargetAirbin.cpp

@@ -1061,6 +1061,33 @@ static void configureSwitchBoxes(DeviceOp &targetOp) {
   */
 }
 
+static void configureCascade(DeviceOp &targetOp) {
+  const auto &target_model = xilinx::AIE::getTargetModel(targetOp);
+  if (target_model.getTargetArch() == AIEArch::AIE2) {
+    for (auto configOp : targetOp.getOps<ConfigureCascadeOp>()) {
+      TileOp tile = cast<TileOp>(configOp.getTile().getDefiningOp());
+      auto inputDir = stringifyCascadeDir(configOp.getInputDir()).upper();
+      auto outputDir = stringifyCascadeDir(configOp.getOutputDir()).upper();
+
+      Address address{tile, 0x36060u};
+
+      /*
+       *  * Register value for output BIT 1: 0 == SOUTH, 1 == EAST
+       *  * Register value for input BIT 0: 0 == NORTH, 1 == WEST
+       */
+      uint8_t outputValue = (outputDir == "SOUTH") ? 0 : 1;
+      uint8_t inputValue = (inputDir == "NORTH") ? 0 : 1;
+
+      constexpr Field<1> Output;
+      constexpr Field<0> Input;
+
+      auto regValue = Output(outputValue) | Input(inputValue);
+
+      write32(address, regValue);
+    }
+  }
+}
+
 /*
         Convert memory address to index
 
@@ -1184,6 +1211,7 @@ mlir::LogicalResult AIETranslateToAirbin(mlir::ModuleOp module,
   }
 
   configureSwitchBoxes(targetOp);
+  configureCascade(targetOp);
   configureDMAs(targetOp);
   groupSections(sections);
 

lib/Targets/AIETargetCDODirect.cpp

@@ -628,7 +628,7 @@ struct AIEControl {
         int arbiter = -1;
 
         for (auto val : connectOp.getAmsels()) {
-          AMSelOp amsel = dyn_cast<AMSelOp>(val.getDefiningOp());
+          AMSelOp amsel = cast<AMSelOp>(val.getDefiningOp());
           arbiter = amsel.arbiterIndex();
           int msel = amsel.getMselValue();
           mask |= (1 << msel);
@@ -653,8 +653,7 @@ struct AIEControl {
         int slot = 0;
         Block &block = connectOp.getRules().front();
         for (auto slotOp : block.getOps<PacketRuleOp>()) {
-          AMSelOp amselOp =
-              dyn_cast<AMSelOp>(slotOp.getAmsel().getDefiningOp());
+          AMSelOp amselOp = cast<AMSelOp>(slotOp.getAmsel().getDefiningOp());
           int arbiter = amselOp.arbiterIndex();
           int msel = amselOp.getMselValue();
           TRY_XAIE_API_EMIT_ERROR(
@@ -702,6 +701,22 @@ struct AIEControl {
             WIRE_BUNDLE_TO_STRM_SW_PORT_TYPE.at(connectOp.getDestBundle()),
             connectOp.destIndex());
     }
+
+    // Cascade configuration
+    const auto &target_model = xilinx::AIE::getTargetModel(targetOp);
+    if (target_model.getTargetArch() == AIEArch::AIE2) {
+      for (auto configOp : targetOp.getOps<ConfigureCascadeOp>()) {
+        TileOp tile = cast<TileOp>(configOp.getTile().getDefiningOp());
+        auto tileLoc = XAie_TileLoc(tile.getCol(), tile.getRow());
+        TRY_XAIE_API_EMIT_ERROR(
+            targetOp, XAie_CoreConfigAccumulatorControl, &devInst, tileLoc,
+            WIRE_BUNDLE_TO_STRM_SW_PORT_TYPE.at(
+                static_cast<WireBundle>(configOp.getInputDir())),
+            WIRE_BUNDLE_TO_STRM_SW_PORT_TYPE.at(
+                static_cast<WireBundle>(configOp.getOutputDir())));
+      }
+    }
+
     return success();
   }