[ObjectFifo] Combine Logical ObjectFifos for reuse

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIECombineLogicalObjFifosForConnectionReuse.cpp

@@ -0,0 +1,50 @@
+// Copyright 2024 The IREE Authors
+//
+// 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
+#include "iree-amd-aie/IR/AMDAIEOps.h"
+#include "iree-amd-aie/Transforms/AMDAIELogicalObjFifoSplittingUtils.h"
+#include "iree-amd-aie/Transforms/Passes.h"
+#include "mlir/IR/Iterators.h"
+#include "mlir/Pass/Pass.h"
+
+#define DEBUG_TYPE "iree-amdaie-combine-logical-objectfifos-for-connection-reuse"
+
+namespace mlir::iree_compiler::AMDAIE {
+
+namespace {
+
+class AMDAIECombineLogicalObjFifosForConnectionReusePass
+    : public impl::AMDAIECombineLogicalObjFifosForConnectionReuseBase<
+          AMDAIECombineLogicalObjFifosForConnectionReusePass> {
+ public:
+  using AMDAIECombineLogicalObjFifosForConnectionReuseBase::
+      AMDAIECombineLogicalObjFifosForConnectionReuseBase;
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<AMDAIEDialect>();
+  }
+  void runOnOperation() override;
+};
+
+void AMDAIECombineLogicalObjFifosForConnectionReusePass::runOnOperation() {
+  ModuleOp moduleOp = getOperation();
+  MLIRContext *context = &getContext();
+  IRRewriter rewriter(context);
+
+  SmallVector<AMDAIE::DmaCpyNdOp> l2ToL1DmaOps =
+      fetchDmaCpyNdOpsToSplitOrCombine(moduleOp);
+
+  if (failed(combineLogicalObjectFifos(rewriter, l2ToL1DmaOps, context))) {
+    return signalPassFailure();
+  }
+}
+
+}  // namespace
+
+std::unique_ptr<Pass> createAMDAIECombineLogicalObjFifosForConnectionReusePass() {
+  return std::make_unique<AMDAIECombineLogicalObjFifosForConnectionReusePass>();
+}
+
+}  // namespace mlir::iree_compiler::AMDAIE

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEDmaLoopSubsumption.cpp

@@ -52,60 +52,6 @@ int64_t calculateNbIterations(int64_t lowerBound, int64_t upperBound,
 
 namespace {
 
-/// Utility affine expression visitor to retrieve the scale and optional bias
-/// from the expression.
-struct RetrieveScaleAndBias
-    : public AffineExprVisitor<RetrieveScaleAndBias, LogicalResult> {
-  std::optional<int64_t> scale;
-  std::optional<int64_t> bias;
-  LogicalResult visitAffineBinaryOpExpr(AffineBinaryOpExpr /*expr*/) {
-    return failure();
-  }
-  LogicalResult visitConstantExpr(AffineConstantExpr /*expr*/) {
-    return failure();
-  }
-  LogicalResult visitDimExpr(AffineDimExpr /*expr*/) { return failure(); }
-  LogicalResult visitSymbolExpr(AffineSymbolExpr /*expr*/) { return failure(); }
-  LogicalResult visitMulExpr(AffineBinaryOpExpr expr) {
-    if (auto rhsSize = dyn_cast<AffineConstantExpr>(expr.getRHS());
-        isa<AffineDimExpr>(expr.getLHS())) {
-      scale = rhsSize.getValue();
-    } else if (auto lhsSize = dyn_cast<AffineConstantExpr>(expr.getLHS());
-               isa<AffineDimExpr>(expr.getRHS())) {
-      scale = lhsSize.getValue();
-    }
-    return success();
-  }
-  LogicalResult visitAddExpr(AffineBinaryOpExpr expr) {
-    if (bias) return failure();
-    if (auto rhsSize = dyn_cast<AffineConstantExpr>(expr.getRHS())) {
-      bias = rhsSize.getValue();
-      if (bias.value() < 0) return failure();
-      if (isa<AffineBinaryOpExpr>(expr.getLHS())) {
-        return visit(expr.getLHS());
-      } else if (isa<AffineDimExpr>(expr.getLHS())) {
-        scale = 1;
-        return success();
-      } else {
-        return failure();
-      }
-    } else if (auto lhsSize = dyn_cast<AffineConstantExpr>(expr.getLHS())) {
-      bias = lhsSize.getValue();
-      if (bias.value() < 0) return failure();
-      if (isa<AffineBinaryOpExpr>(expr.getRHS())) {
-        return visit(expr.getRHS());
-      } else if (isa<AffineDimExpr>(expr.getRHS())) {
-        scale = 1;
-        return success();
-      } else {
-        return failure();
-      }
-    } else {
-      return failure();
-    }
-  }
-};
-
 struct SubsumeLoopIntoDMA
     : public OpInterfaceRewritePattern<AMDAIE::DoublyStridedOpInterface> {
   using OpInterfaceRewritePattern::OpInterfaceRewritePattern;

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEDmaUtils.h

@@ -12,12 +12,71 @@
 #include "iree-amd-aie/IR/AMDAIEOps.h"
 #include "iree-amd-aie/aie_runtime/iree_aie_runtime.h"
 #include "llvm/ADT/SmallVector.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/PatternMatch.h"
 
 namespace mlir::iree_compiler::AMDAIE {
 
+/// Utility to retrieve a constant index from an OpFoldResult.
+int64_t getConstantIndexOrAssert(OpFoldResult dim);
+
+/// Utility affine expression visitor to retrieve the scale and optional bias
+/// from the expression.
+struct RetrieveScaleAndBias
+    : public AffineExprVisitor<RetrieveScaleAndBias, LogicalResult> {
+  std::optional<int64_t> scale;
+  std::optional<int64_t> bias;
+  LogicalResult visitAffineBinaryOpExpr(AffineBinaryOpExpr /*expr*/) {
+    return failure();
+  }
+  LogicalResult visitConstantExpr(AffineConstantExpr /*expr*/) {
+    return failure();
+  }
+  LogicalResult visitDimExpr(AffineDimExpr /*expr*/) { return failure(); }
+  LogicalResult visitSymbolExpr(AffineSymbolExpr /*expr*/) { return failure(); }
+  LogicalResult visitMulExpr(AffineBinaryOpExpr expr) {
+    if (auto rhsSize = dyn_cast<AffineConstantExpr>(expr.getRHS());
+        isa<AffineDimExpr>(expr.getLHS())) {
+      scale = rhsSize.getValue();
+    } else if (auto lhsSize = dyn_cast<AffineConstantExpr>(expr.getLHS());
+               isa<AffineDimExpr>(expr.getRHS())) {
+      scale = lhsSize.getValue();
+    }
+    return success();
+  }
+  LogicalResult visitAddExpr(AffineBinaryOpExpr expr) {
+    if (bias) return failure();
+    if (auto rhsSize = dyn_cast<AffineConstantExpr>(expr.getRHS())) {
+      bias = rhsSize.getValue();
+      if (bias.value() < 0) return failure();
+      if (isa<AffineBinaryOpExpr>(expr.getLHS())) {
+        return visit(expr.getLHS());
+      } else if (isa<AffineDimExpr>(expr.getLHS())) {
+        scale = 1;
+        return success();
+      } else {
+        return failure();
+      }
+    } else if (auto lhsSize = dyn_cast<AffineConstantExpr>(expr.getLHS())) {
+      bias = lhsSize.getValue();
+      if (bias.value() < 0) return failure();
+      if (isa<AffineBinaryOpExpr>(expr.getRHS())) {
+        return visit(expr.getRHS());
+      } else if (isa<AffineDimExpr>(expr.getRHS())) {
+        scale = 1;
+        return success();
+      } else {
+        return failure();
+      }
+    } else {
+      return failure();
+    }
+  }
+};
+
 // Constant specifying the number of inter-iteration dimension for DMA
 // operations.
 //
@@ -194,9 +253,9 @@ struct DmaDimConfig {
   AMDAIE::AMDAIETileType sourceTileType;
   AMDAIE::AMDAIETileType targetTileType;
   /// The maximum number of addressing dimensions on the source side of the DMA.
-  uint8_t sourceMaxNbDims{0};
+  int64_t sourceMaxNbDims{0};
   /// The maximum number of addressing dimensions on the target side of the DMA.
-  uint8_t targetMaxNbDims{0};
+  int64_t targetMaxNbDims{0};
 
   DmaDimConfig(const AMDAIE::AMDAIEDeviceModel &deviceModel,
                uint8_t sourceMemspaceInt, uint8_t targetMemspaceInt)