[Reprogram] Add a global flag + e2e test case for reprogramming Dmas (#1334)

-- This commit adds a global flag `--iree-amdaie-reprogram-dmas` to be
       used when we want to reprogram the DMAs.
-- Also adds an e2e test case.
-- This is in accordance to the [reprogramming DMA
work](#1287) and is the
final PR to start off supporting it.
    
Signed-off-by: Abhishek Varma <[email protected]>

---------

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

Author: Abhishek-Varma

build_tools/ci/cpu_comparison/matmul_test_config.py

@@ -16,6 +16,20 @@
             "--iree-amdaie-num-cols=1",
         ],
     },
+    {
+        "M": 32,
+        "N": 32,
+        "K": 32,
+        "input_type": "i32",
+        "acc_type": "i32",
+        "name_suffix": "reprogram_dma",
+        "additional_labels": ["ReprogramDmas"],
+        "aie_compilation_flags": [
+            "--iree-amdaie-num-rows=1",
+            "--iree-amdaie-num-cols=1",
+            "--iree-amdaie-reprogram-dmas",
+        ],
+    },
     # 2x2 core tests.
     {
         "M": 32,

compiler/plugins/target/AMD-AIE/iree-amd-aie/Target/AIETarget.cpp

@@ -241,7 +241,7 @@ class AIETargetBackend final : public IREE::HAL::TargetBackend {
         options.enableCoalescingLoops, options.enableCollapsingUnitDims,
         options.enableFunctionOutlining, options.callReplication,
         options.insertLoopAroundCoreBlock, options.enableCtrlPkt,
-        options.coreStackSize);
+        options.coreStackSize, options.reprogramDmas);
   }
 
   void buildLinkingPassPipeline(OpPassManager &passManager) override {

compiler/plugins/target/AMD-AIE/iree-amd-aie/Target/AIETarget.h

@@ -64,6 +64,9 @@ struct AMDAIEOptions {
   bool matmulElementwiseFusion{false};
   AMDAIEDevice AMDAIETargetDevice{AMDAIEDevice::npu1_4col};
 
+  // Enable/Disable reprogramming of DMAs.
+  bool reprogramDmas{false};
+
   // The number of rows for the compiler to target. '0' denotes 'all'.
  private:
   unsigned AMDAIENumRows{0};
@@ -257,6 +260,13 @@ struct AMDAIEOptions {
             "for matmul-elementwise fusion. It is currently added for "
             "development purpose and should be removed in the future."));
 
+    binder.opt<bool>(
+        "iree-amdaie-reprogram-dmas", reprogramDmas, llvm::cl::cat(category),
+        llvm::cl::desc(
+            "Flag used to enable/disable reprogramming of DMAs. "
+            "By default it'll be disabled, so we would have circular "
+            "DMAs for L2/L1 cache."));
+
     /// Command line option for selecting the target AIE device.
     binder.opt<AMDAIEDevice>(
         "iree-amdaie-target-device", AMDAIETargetDevice,

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

@@ -107,7 +107,7 @@ static void addAMDAIEBufferizePasses(OpPassManager &pm,
 }
 
 void addAMDAIEToAIEPasses(OpPassManager &passManager,
-                          bool insertLoopAroundCoreBlock) {
+                          bool insertLoopAroundCoreBlock, bool reprogramDmas) {
   // The infinite loop insertion transformation needs to be called before the
   // `AcquireReleaseToUseLock` pass as the latter will perform loop unrolling
   // based on the objFifo depths.
@@ -123,9 +123,7 @@ void addAMDAIEToAIEPasses(OpPassManager &passManager,
   passManager.addPass(createAMDAIEAddNoAliasFunctionArgumentsPass());
   {
     AMDAIELowerToAIEOptions options;
-    // TODO(avarma): In follow-up PRs this will be replaced by a global flag.
-    // Currently setting as `false`.
-    options.reprogramDmas = /*reprogramDmas=*/false;
+    options.reprogramDmas = reprogramDmas;
     passManager.addPass(createAMDAIELowerToAIEPass(options));
   }
   passManager.addPass(createAMDAIERemoveMemorySpacePass());
@@ -672,7 +670,7 @@ void buildAMDAIETransformPassPipeline(
     PacketFlowStrategy packetFlowStrategy, bool enableCoalescingLoops,
     bool enableCollapsingUnitDims, OutliningStrategy enableFunctionOutlining,
     int callReplication, bool insertLoopAroundCoreBlock, bool enableCtrlPkt,
-    uint32_t coreStackSize) {
+    uint32_t coreStackSize, bool reprogramDmas) {
   OpPassManager &modulePassManager = variantPassManager.nest<ModuleOp>();
   {
     FunctionLikeNest funcPassManager(modulePassManager);
@@ -707,7 +705,8 @@ void buildAMDAIETransformPassPipeline(
         modulePassManager, packetFlowStrategy, useTilePipeline,
         enableVectorizationPasses, enableCoalescingLoops,
         enableCollapsingUnitDims, enableFunctionOutlining, callReplication,
-        insertLoopAroundCoreBlock, numCols, enableCtrlPkt, coreStackSize);
+        insertLoopAroundCoreBlock, numCols, enableCtrlPkt, coreStackSize,
+        reprogramDmas);
   } else if (useLowerToAIEPipeline == LowerToAIEPassPipeline::AIR) {
     addMLIRAIRLoweringPasses(modulePassManager, device, useTilePipeline,
                              matmulElementwiseFusion,
@@ -733,7 +732,7 @@ void addAMDAIEObjectFifoLoweringPasses(
     bool enableCoalescingLoops, bool enableCollapsingUnitDims,
     OutliningStrategy enableFunctionOutlining, int callReplication,
     bool insertLoopAroundCoreBlock, uint32_t numCols, bool enableCtrlPkt,
-    uint32_t coreStackSize) {
+    uint32_t coreStackSize, bool reprogramDmas) {
   passManager.addPass(createEraseHALDescriptorTypeFromMemRefPass());
   passManager.addPass(memref::createFoldMemRefAliasOpsPass());
 
@@ -796,18 +795,28 @@ void addAMDAIEObjectFifoLoweringPasses(
 
   passManager.addPass(createCSEPass());
   passManager.addPass(createCanonicalizerPass());
-  passManager.addPass(createAMDAIEAssignLogicalObjectFifoDepthPass());
+  {
+    AMDAIEAssignLogicalObjectFifoDepthOptions options;
+    // TODO(avarma): In case reprogramming Dmas, we currently disable double
+    // buffering. Relax the constraint later after modifying
+    // controlcode-lowering and controlcode-to-transaction-binary pass to work
+    // with double buffering.
+    if (reprogramDmas) {
+      options.l2BufferDepth = 1;
+      options.l1BufferDepth = 1;
+    }
+    passManager.addPass(createAMDAIEAssignLogicalObjectFifoDepthPass(options));
+  }
 
   passManager.addPass(createAMDAIEAssignTilesPass());
   passManager.addPass(createCSEPass());
   passManager.addPass(createCanonicalizerPass());
 
-  passManager.addPass(createAMDAIEDmaToCircularDmaPass());
+  if (!reprogramDmas) passManager.addPass(createAMDAIEDmaToCircularDmaPass());
+
   {
     AMDAIECreateAIEWorkgroupOptions options;
-    // TODO(avarma): In follow-up PRs this will be replaced by a global flag.
-    // Currently setting as `false`.
-    options.reprogramDmas = /*reprogramDmas=*/false;
+    options.reprogramDmas = reprogramDmas;
     passManager.addNestedPass<func::FuncOp>(
         createAMDAIECreateAIEWorkgroupPass(options));
   }
@@ -874,11 +883,30 @@ void addAMDAIEObjectFifoLoweringPasses(
 
   passManager.addPass(createAMDAIENpuDmaToHalfDmaCpyNdPass());
   passManager.addPass(createAMDAIEInsertDmaBdChainPass());
-  passManager.addPass(createAMDAIEFoldDmaWaitsPass());
-  passManager.addPass(createAMDAIEControlCodeLoweringPass());
+  // TODO(avarma): Currently with fold dma wait pass, in case of DMA
+  // reprogramming we get ALL zeroes. To be triaged/fixed later in order to
+  // relax this constraint and optimize the wait ops.
+  if (!reprogramDmas) passManager.addPass(createAMDAIEFoldDmaWaitsPass());
+
+  {
+    AMDAIEControlCodeLoweringOptions options;
+    options.reprogramDmas = reprogramDmas;
+    passManager.addPass(createAMDAIEControlCodeLoweringPass(options));
+  }
+  if (reprogramDmas) {
+    passManager.addPass(createAMDAIEAssignBDIDsPass());
+    {
+      // For Conv ops use basic sequential scheme to avoid numerical error.
+      // TODO: Find a better working scheme for Conv ops
+      AMDAIEAssignBufferAddressOptions options;
+      if (useTilePipeline == TilePassPipeline::ConvDecomposePipeline)
+        options.allocScheme = AllocScheme::Sequential;
+      passManager.addPass(createAMDAIEAssignBufferAddressPass(options));
+    }
+  }
   passManager.addPass(createAMDAIEControlCodeToTransactionPass());
 
-  addAMDAIEToAIEPasses(passManager, insertLoopAroundCoreBlock);
+  addAMDAIEToAIEPasses(passManager, insertLoopAroundCoreBlock, reprogramDmas);
 
   // Now lower using the AIE passes from MLIR-AIE.
   addMLIRAIELoweringPasses(passManager, useTilePipeline);

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

@@ -20,7 +20,7 @@ void addAMDAIEObjectFifoLoweringPasses(
     bool enableCoalescingLoops, bool enableCollapsingUnitDims,
     OutliningStrategy enableFunctionOutlining, int outliningLoopInCallCount,
     bool insertLoopAroundCoreBlock, uint32_t numCols, bool emitCtrlPkt,
-    uint32_t coreStackSize);
+    uint32_t coreStackSize, bool reprogramDmas);
 
 /// Add passes to lower from MLIR-AIR through AIE. This is
 /// currently the default passes used for lowering after IREEs tiling.
@@ -45,7 +45,7 @@ void buildAMDAIETransformPassPipeline(
     PacketFlowStrategy packetFlowStrategy, bool enableCoalescingLoops,
     bool enableCollapsingUnitDims, OutliningStrategy enableFunctionOutlining,
     int outliningLoopInCallCount, bool insertLoopAroundCoreBlock,
-    bool emitCtrlPkt, uint32_t coreStackSize);
+    bool emitCtrlPkt, uint32_t coreStackSize, bool reprogramDmas);
 
 /// Populates passes needed to lower the IR via a Pack-Peel based approach.
 void addPackPeelBasedPassPipeline(OpPassManager &passManager,
@@ -272,7 +272,7 @@ std::unique_ptr<OperationPass<ModuleOp>> createAMDAIELoweringStrategyPass(
 std::unique_ptr<Pass> createAMDAIELowerFuncArgsPass();
 
 /// Create pass to lower from the AMDAIE dialect to the AIE/AIEX dialects.
-void addAMDAIEToAIEPasses(OpPassManager &);
+void addAMDAIEToAIEPasses(OpPassManager &pm, bool reprogramDmas);
 std::unique_ptr<Pass> createAMDAIELowerToAIEPass(
     AMDAIELowerToAIEOptions options = {});