[Reprogram] Modify controlcode-lowering and lower-to-aie for DMA reprogramming (#1330)

-- This commit includes modifications to `controlcode-lowering` and
   `lower-to-aie` passes for DMA reprogramming.

-- This is being added to AMDAIE dialect to make [DMA
reprogramming](#1287) work.

Signed-off-by: Abhishek Varma <[email protected]>

Author: Abhishek-Varma

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

@@ -35,82 +35,6 @@
 using namespace xilinx;
 
 namespace mlir::iree_compiler::AMDAIE {
-
-/// Compute the 'global' repetition count: the product over all dimensions with
-/// zero stride of the size of the dimension.
-///
-/// The case where sizes and strides are empty is a special case, and '0' is
-/// returned.
-static int64_t getRepetitionCount(ArrayRef<OpFoldResult> sizes,
-                                  ArrayRef<OpFoldResult> strides) {
-  assert(sizes.size() == strides.size() &&
-         "expected stride and size vectors of same size");
-  if (strides.empty()) return 0;
-  size_t repetitionCount{1};
-  for (uint32_t i = 0; i < strides.size(); ++i) {
-    if (!isConstantIntValue(strides[i], 0)) continue;
-    std::optional<int64_t> maybeSize = getConstantIntValue(sizes[i]);
-    assert(maybeSize.has_value() &&
-           "expected constant size in this zero stride dimension");
-    assert(maybeSize.value() >= 0 && "expected a non-negative size");
-    repetitionCount *= maybeSize.value();
-  }
-  return repetitionCount;
-}
-
-/// Utility to retrieve the common repetition count from all producers and
-/// consumers of a logical objectFifo.
-static FailureOr<size_t> getRepetitionCount(LogicalObjFifoOpInterface op) {
-  SmallVector<int64_t> repetitionCounts;
-  auto appendRepetitionCount = [&](ArrayRef<OpFoldResult> sizes,
-                                   ArrayRef<OpFoldResult> strides) {
-    size_t repetitionCount = getRepetitionCount(sizes, strides);
-    if (repetitionCount != 0) repetitionCounts.push_back(repetitionCount);
-  };
-
-  for (Operation *userOp : op->getUsers()) {
-    if (auto connectionOp = dyn_cast<AMDAIE::ConnectionOp>(userOp)) {
-      FailureOr<AMDAIE::NpuCircularDmaCpyNdOp> maybeNpuDmaUserOp =
-          connectionOp.getNpuCircularDmaCpyNdUser();
-
-      if (failed(maybeNpuDmaUserOp)) continue;
-
-      AMDAIE::NpuCircularDmaCpyNdOp npuDma = maybeNpuDmaUserOp.value();
-
-      if (connectionOp.getTarget() &&
-          dyn_cast_if_present<LogicalObjFifoOpInterface>(
-              connectionOp.getTarget().getDefiningOp()) == op) {
-        appendRepetitionCount(npuDma.getTargetMixedSizes(),
-                              npuDma.getTargetMixedStrides());
-      }
-
-      if (connectionOp.getSource() &&
-          dyn_cast_if_present<LogicalObjFifoOpInterface>(
-              connectionOp.getSource().getDefiningOp()) == op) {
-        appendRepetitionCount(npuDma.getSourceMixedSizes(),
-                              npuDma.getSourceMixedStrides());
-      }
-    }
-  }
-
-  // merge the repetition counts:
-  if (repetitionCounts.empty()) return 1;
-  int64_t combinedRepetitionCount =
-      *std::min_element(repetitionCounts.begin(), repetitionCounts.end());
-
-  // if any of the repetition counts are not divisible by the combined
-  // repetition count, that's a problem:
-  if (!std::all_of(
-          repetitionCounts.begin(), repetitionCounts.end(),
-          [&](size_t c) { return c % combinedRepetitionCount == 0; })) {
-    return op.emitOpError()
-           << " could not resolved a common repetition count based on the "
-              "individual repetition counts: "
-           << getArrayString<int64_t>(repetitionCounts);
-  }
-  return combinedRepetitionCount;
-}
-
 //===----------------------------------------------------------------------===//
 // AIEDeviceBuilder utilities
 //===----------------------------------------------------------------------===//
@@ -334,44 +258,6 @@ void AIEDeviceBuilder::eraseOp(Operation *op) {
   rewriter.eraseOp(op);
 }
 
-LogicalResult AIEDeviceBuilder::foldDimsAndReturnAsStatic(
-    SmallVector<OpFoldResult> sizes, SmallVector<OpFoldResult> strides,
-    SmallVector<int64_t> &newSizes, SmallVector<int64_t> &newStrides,
-    size_t repetitionCount, uint8_t memSpace,
-    function_ref<InFlightDiagnostic()> emitError) {
-  if (failed(foldRepetitionCount(rewriter.getContext(), sizes, strides,
-                                 repetitionCount))) {
-    return emitError() << "could not fold repetition counts from sizes: "
-                       << getConstantIntValuesString(sizes)
-                       << " strides: " << getConstantIntValuesString(strides)
-                       << " repetitionCount: " << repetitionCount << ".";
-  }
-  SmallVector<OpFoldResult> offsets(
-      strides.size(), getAsIndexOpFoldResult(rewriter.getContext(), 0));
-  (void)foldUnitDims(rewriter.getContext(), offsets, sizes, strides);
-
-  DmaDimConfig dmaDimConfig(deviceModel, memSpace);
-  SmallVector<int64_t> maxSizes = dmaDimConfig.getMaxSizes(offsets.size());
-  SmallVector<OpFoldResult> linearOffsets, linearSizes, linearStrides;
-  (void)foldLinearDims(
-      rewriter.getContext(), offsets, sizes, strides, linearOffsets,
-      linearSizes, linearStrides, [&](size_t idxFromEnd, int64_t size) {
-        return idxFromEnd < maxSizes.size() &&
-               size <= maxSizes[maxSizes.size() - idxFromEnd - 1];
-      });
-  std::optional<SmallVector<int64_t>> maybeStaticSizes =
-      getConstantIntValues(linearSizes);
-  std::optional<SmallVector<int64_t>> maybeStaticStrides =
-      getConstantIntValues(linearStrides);
-  if (!maybeStaticSizes || !maybeStaticStrides) {
-    return emitError()
-           << "found dynamic sizes or strides which is not supported";
-  }
-  newSizes = std::move(maybeStaticSizes.value());
-  newStrides = std::move(maybeStaticStrides.value());
-  return success();
-}
-
 void AIEDeviceBuilder::remapOperands(Operation *op) {
   for (int i = 0; i < op->getNumOperands(); ++i) {
     Value operand = op->getOperand(i);
@@ -577,6 +463,7 @@ LogicalResult AIEDeviceBuilder::bufferToAIE(AMDAIE::BufferOp bufferOp,
 LogicalResult AIEDeviceBuilder::connectionToAIE(
     AMDAIE::ConnectionOp connectionOp, Block *deviceBlock,
     int &connectionIndex) {
+  if (reprogramDmas) return success();
   LLVM_DEBUG(llvm::dbgs() << "Convert [AMDAIE::ConnectionOp]\n");
   rewriter.setInsertionPoint(deviceBlock->getTerminator());
   SmallVector<AMDAIE::ChannelOp> producerChannels;
@@ -705,7 +592,7 @@ LogicalResult AIEDeviceBuilder::connectionToAIE(
           std::make_pair(consumerLocks[0], producerLocks[0]);
       SmallVector<int64_t> canonicalizedSizes, canonicalizedStrides;
       if (failed(foldDimsAndReturnAsStatic(
-              maybeNpuDmaUserOp->getSourceMixedSizes(),
+              rewriter, deviceModel, maybeNpuDmaUserOp->getSourceMixedSizes(),
               maybeNpuDmaUserOp->getSourceMixedStrides(), canonicalizedSizes,
               canonicalizedStrides, repetitionCount.value(),
               maybeSourceMemSpace.value(),
@@ -810,7 +697,7 @@ LogicalResult AIEDeviceBuilder::connectionToAIE(
           std::make_pair(producerLocks[0], consumerLocks[0]);
       SmallVector<int64_t> canonicalizedSizes, canonicalizedStrides;
       if (failed(foldDimsAndReturnAsStatic(
-              maybeNpuDmaUserOp->getTargetMixedSizes(),
+              rewriter, deviceModel, maybeNpuDmaUserOp->getTargetMixedSizes(),
               maybeNpuDmaUserOp->getTargetMixedStrides(), canonicalizedSizes,
               canonicalizedStrides, repetitionCount.value(),
               maybeTargetMemSpace.value(),
@@ -993,7 +880,7 @@ LogicalResult AIEDeviceBuilder::workgroupToAIE(AMDAIE::WorkgroupOp workgroupOp,
           if (failed(connectionToAIE(dmaOp, deviceBlock, connectionIndex))) {
             return WalkResult::interrupt();
           }
-          return WalkResult::advance();
+          return WalkResult::skip();
         })
         .Case<AMDAIE::ControlCodeOp>([&](auto controlCodeOp) {
           // Skip control code as it should already be translated into firmware
@@ -1178,6 +1065,9 @@ LogicalResult AIEDeviceBuilder::lowerToAIE(ModuleOp moduleOp) {
 class AMDAIELowerToAIEPass
     : public impl::AMDAIELowerToAIEBase<AMDAIELowerToAIEPass> {
  public:
+  AMDAIELowerToAIEPass(const AMDAIELowerToAIEOptions &options)
+      : AMDAIELowerToAIEBase(options) {}
+
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<mlir::memref::MemRefDialect, AMDAIEDialect,
                     xilinx::AIE::AIEDialect, xilinx::AIEX::AIEXDialect>();
@@ -1196,13 +1086,15 @@ class AMDAIELowerToAIEPass
       return signalPassFailure();
     }
     AMDAIEDeviceModel deviceModel = getDeviceModel(maybeDevice.value());
-    AIEDeviceBuilder builder(moduleOp.getContext(), std::move(deviceModel));
+    AIEDeviceBuilder builder(moduleOp.getContext(), std::move(deviceModel),
+                             reprogramDmas);
     if (failed(builder.lowerToAIE(moduleOp))) return signalPassFailure();
   }
 };
 
-std::unique_ptr<Pass> createAMDAIELowerToAIEPass() {
-  return std::make_unique<AMDAIELowerToAIEPass>();
+std::unique_ptr<Pass> createAMDAIELowerToAIEPass(
+    AMDAIELowerToAIEOptions options) {
+  return std::make_unique<AMDAIELowerToAIEPass>(options);
 }
 
 }  // namespace mlir::iree_compiler::AMDAIE

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

@@ -32,8 +32,11 @@ using BDDimLayoutAndLength = std::pair<AIE::BDDimLayoutArrayAttr, int64_t>;
 /// `amdaie.workgroup`.
 class AIEDeviceBuilder {
  public:
-  AIEDeviceBuilder(MLIRContext *ctx, AMDAIEDeviceModel deviceModel)
-      : rewriter(ctx), deviceModel(std::move(deviceModel)) {}
+  AIEDeviceBuilder(MLIRContext *ctx, AMDAIEDeviceModel deviceModel,
+                   bool reprogramDmas)
+      : rewriter(ctx),
+        deviceModel(std::move(deviceModel)),
+        reprogramDmas(reprogramDmas) {}
 
   LogicalResult lowerToAIE(ModuleOp moduleOp);
 
@@ -98,14 +101,6 @@ class AIEDeviceBuilder {
   /// might be used after `op` is erased.
   void eraseOp(Operation *op);
 
-  /// Utility to fold the provided repetition count, unit dims, linear dims and
-  /// to convert the sizes and strides into static versions and return them.
-  LogicalResult foldDimsAndReturnAsStatic(
-      SmallVector<OpFoldResult> sizes, SmallVector<OpFoldResult> strides,
-      SmallVector<int64_t> &newSizes, SmallVector<int64_t> &newStrides,
-      size_t repetitionCount, uint8_t memSpace,
-      function_ref<InFlightDiagnostic()> emitError);
-
   /// Utility to remap the provided operation's operands.
   void remapOperands(Operation *op);
 
@@ -127,6 +122,9 @@ class AIEDeviceBuilder {
       connectionToSourceTargetMemOps;
   /// Map from connection ops to the flow ops they have been converted into.
   DenseMap<AMDAIE::ConnectionOp, SmallVector<Operation *>> connectionToFlowOps;
+  /// Set using the pass' flag `reprogram-dmas` and is used to enable/disable
+  /// reprogramming of DMAs.
+  bool reprogramDmas;
 };
 
 }  // namespace mlir::iree_compiler::AMDAIE

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

@@ -121,7 +121,13 @@ void addAMDAIEToAIEPasses(OpPassManager &passManager,
   passManager.addPass(createAMDAIESinkIntoCorePass());
   passManager.addPass(createCanonicalizerPass());
   passManager.addPass(createAMDAIEAddNoAliasFunctionArgumentsPass());
-  passManager.addPass(createAMDAIELowerToAIEPass());
+  {
+    AMDAIELowerToAIEOptions options;
+    // TODO(avarma): In follow-up PRs this will be replaced by a global flag.
+    // Currently setting as `false`.
+    options.reprogramDmas = /*reprogramDmas=*/false;
+    passManager.addPass(createAMDAIELowerToAIEPass(options));
+  }
   passManager.addPass(createAMDAIERemoveMemorySpacePass());
   passManager.addPass(createCanonicalizerPass());
 }

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

@@ -273,7 +273,8 @@ std::unique_ptr<Pass> createAMDAIELowerFuncArgsPass();
 
 /// Create pass to lower from the AMDAIE dialect to the AIE/AIEX dialects.
 void addAMDAIEToAIEPasses(OpPassManager &);
-std::unique_ptr<Pass> createAMDAIELowerToAIEPass();
+std::unique_ptr<Pass> createAMDAIELowerToAIEPass(
+    AMDAIELowerToAIEOptions options = {});
 
 /// Create pass to lower a sequence of operation(s) to a iree_codegen.ukernel.*
 /// operation.

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

@@ -215,6 +215,8 @@ def AMDAIEControlCodeLowering :
   let options = [
     Option<"argIdxOffset", "arg-idx-offset", "int32_t", /*default=*/"0",
       "The offset to be added to the argument index.">,
+    Option<"reprogramDmas", "reprogram-dmas", "bool", /*default=*/"false",
+           "Flag to reprogram DMAs. When enabled, no circular DMAs will be produced">,
   ];
 }
 
@@ -639,6 +641,10 @@ def AMDAIELowerToAIE :
     Pass<"iree-amdaie-lower-to-aie", "ModuleOp"> {
   let summary = "Lower from the AMDAIE dialect to the AIE/AIEX dialects";
   let constructor = "mlir::iree_compiler::AMDAIE::createAMDAIELowerToAIEPass()";
+  let options = [
+    Option<"reprogramDmas", "reprogram-dmas", "bool", /*default=*/"false",
+           "Flag to reprogram DMAs. When enabled, no circular DMAs will be produced">,
+  ];
 }
 
 def AMDAIELowerToUKernels :