Add pass to assign BD IDs to NPU DMA ops and refactor `AMDAIEAssignBufferDescriptorIDsPass` #551

jtuyls · 2024-07-15T17:34:09Z

Adds a pass to generate BD IDs for NPU DMA operations that avoids race conditions with BDs being reconfigured during execution by some DMA.

This only assigns BD IDs on the L3 side of NPU DMA operations as the backend can currently only handle Shim DMA instructions and MemTile/core DMA configurations are assumed to be part of the xclbin. Therefore, you will see the checks like npuDmaOp.getSourceMemorySpaceAsUInt() == 0 and npuDmaOp.getTargetMemorySpaceAsUInt() == 0 being used (memory space 0 is L3).

BD IDs are represented with an amdaie.bd_id operation, which takes a tile and an integer value as operands as every tile has its own set of BDs. This helps with guaranteeing/verifying correct reuse of the same id across DMA operations. This usage was derived from Xilinx/mlir-aie#1610

makslevental · 2024-07-16T03:07:21Z

runtime/src/iree-amd-aie/aie_runtime/iree_aie_runtime.h

+  /// Return a map from channels to valid BD ids for the requested tile type.
+  /// TODO(jornt): find these ranges in the device model.


it's here https://github.com/Xilinx/aie-rt/blob/c28d57972957d8032839f4da5d2e3b9dd01805ea/driver/src/global/xaie2ipugbl_reginit.c#L389-L391 but they don't distinguish between even/odd

Yeah, they have the real numbers harcoded inside a verifier here: https://github.com/Xilinx/aie-rt/blob/1ad203de0b7f282b1c0659fd2ae3f218652c7274/driver/src/dma/xaie_dma_aieml.c#L1216

We can send them a PR to factor that out (but also that's a good util to use somewhere around here too)

compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/AMDAIEOps.td

compiler/plugins/target/AMD-AIE/iree-amd-aie/Utils/ChannelBdIdGenerator.cpp

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/lower_to_aie.mlir

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

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

makslevental

mostly fine - a few nittish type things to debate - but can you put a description somewhere (eg this PR itself) for the rules for assigning bdids re npuDmaOp.getSourceMemorySpaceAsUInt() and npuDmaOp.getTargetMemorySpaceAsUInt().

makslevental · 2024-07-16T03:40:47Z

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

+            npuDmaOp.getTargetMixedSizes(), npuDmaOp.getTargetMixedStrides(),
+            npuDmaOp.getSourceMixedOffsets(), npuDmaOp.getSourceMixedSizes(),
+            npuDmaOp.getSourceMixedStrides(), nullptr, bdIdOp);
+      } else if (npuDmaOp.getTargetMemorySpaceAsUInt() == 0) {


these conditionals aren't mutually exclusive (I don't think? maybe the first two are but the last one definitely isn't, at least not according to first principles) so I would prefer if each conditional were separate and each just did a return WalkResult::advance(); at the end.

these conditionals aren't mutually exclusive

In principal you could create a NPU DMA operation with both source and target on L3/DDR, but that doesn't make much sense in the context of AIE. However, maybe this is not a good place to bake in any of those 'assumptions' and I can make it two separate if statements, but with a single return WalkResult::advance(); at the end?

I generally like early returns in ifs because then you don't have to hunt for the exit. And also it prevents multiple ifs firing.

discussed offline

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

makslevental · 2024-07-16T03:48:28Z

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

+        // towards other DMAs and future hardware generations, channel
+        // assignment should happen before BD assignemnt. This requires more
+        // refactoring.
+        std::optional<uint32_t> bdId = generator.getAndAssignBdId(0);


i guess the next step is to put channel on NpuDmaCpyNdOp...

I think channel belongs on the CircularDmaCpyNd operations which are referenced by the NpuDmaCpyNd ops.

Hm really? I guess I haven't been paying attention to these ops. That's different from mlir-aie (obv) so that's why I was expecting a channel on this op.

The connections (routing, channel/port selection) is set up based on the CircularDmaCpyNd ops after lowering to AIE dialect. As it's set up now, the NPU instructions shouldn't change the channels/DMA ports

weird to me (again only because coming from the mlir-aie perspective) but we can discuss it later/elsewhere.

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

makslevental

LGTM

makslevental · 2024-07-16T13:33:38Z

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

+        // towards other DMAs and future hardware generations, channel
+        // assignment should happen before BD assignemnt. This requires more
+        // refactoring.
+        std::optional<uint32_t> bdId = generator.getAndAssignBdId(0);


weird to me (again only because coming from the mlir-aie perspective) but we can discuss it later/elsewhere.

makslevental · 2024-07-16T13:35:39Z

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

+            npuDmaOp.getTargetMixedSizes(), npuDmaOp.getTargetMixedStrides(),
+            npuDmaOp.getSourceMixedOffsets(), npuDmaOp.getSourceMixedSizes(),
+            npuDmaOp.getSourceMixedStrides(), nullptr, bdIdOp);
+      } else if (npuDmaOp.getTargetMemorySpaceAsUInt() == 0) {


discussed offline

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

jtuyls requested review from makslevental, nirvedhmeshram, MaheshRavishankar, yzhang93 and Abhishek-Varma as code owners July 15, 2024 17:34

makslevental reviewed Jul 16, 2024

View reviewed changes

compiler/plugins/target/AMD-AIE/iree-amd-aie/IR/AMDAIEOps.td Show resolved Hide resolved

makslevental reviewed Jul 16, 2024

View reviewed changes

compiler/plugins/target/AMD-AIE/iree-amd-aie/Utils/ChannelBdIdGenerator.cpp Outdated Show resolved Hide resolved

makslevental reviewed Jul 16, 2024

View reviewed changes

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/test/lower_to_aie.mlir Outdated Show resolved Hide resolved

makslevental reviewed Jul 16, 2024

View reviewed changes

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

makslevental reviewed Jul 16, 2024

View reviewed changes

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEAssignNpuDmaBdIds.cpp Outdated Show resolved Hide resolved

makslevental reviewed Jul 16, 2024

View reviewed changes

compiler/plugins/target/AMD-AIE/iree-amd-aie/Transforms/AMDAIEAssignNpuDmaBdIds.cpp Outdated Show resolved Hide resolved

makslevental requested changes Jul 16, 2024

View reviewed changes

Add pass to assign BD IDs to NPU DMA ops and refactor

7e88dea

jtuyls force-pushed the add-assign-npu-bd-ids branch from e6a55a6 to 7e88dea Compare July 16, 2024 10:11

makslevental approved these changes Jul 16, 2024

View reviewed changes

Merge branch 'main' into add-assign-npu-bd-ids

05a353f

jtuyls merged commit c8b8112 into nod-ai:main Jul 16, 2024
2 checks passed

jtuyls deleted the add-assign-npu-bd-ids branch July 16, 2024 14:28

jtuyls mentioned this pull request Jul 16, 2024

Add DMA loop iteration subsumption #495

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add pass to assign BD IDs to NPU DMA ops and refactor `AMDAIEAssignBufferDescriptorIDsPass` #551

Add pass to assign BD IDs to NPU DMA ops and refactor `AMDAIEAssignBufferDescriptorIDsPass` #551

jtuyls commented Jul 15, 2024 •

edited

Loading

makslevental Jul 16, 2024

jtuyls Jul 16, 2024

makslevental Jul 16, 2024

makslevental left a comment

makslevental Jul 16, 2024 •

edited

Loading

jtuyls Jul 16, 2024 •

edited

Loading

makslevental Jul 16, 2024

makslevental Jul 16, 2024

makslevental Jul 16, 2024

jtuyls Jul 16, 2024

makslevental Jul 16, 2024 •

edited

Loading

jtuyls Jul 16, 2024

makslevental Jul 16, 2024

makslevental left a comment

makslevental Jul 16, 2024

makslevental Jul 16, 2024

		/// Return a map from channels to valid BD ids for the requested tile type.
		/// TODO(jornt): find these ranges in the device model.

Add pass to assign BD IDs to NPU DMA ops and refactor AMDAIEAssignBufferDescriptorIDsPass #551

Add pass to assign BD IDs to NPU DMA ops and refactor AMDAIEAssignBufferDescriptorIDsPass #551

Conversation

jtuyls commented Jul 15, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

makslevental left a comment

Choose a reason for hiding this comment

makslevental Jul 16, 2024 • edited Loading

Choose a reason for hiding this comment

jtuyls Jul 16, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

makslevental Jul 16, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

makslevental left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Add pass to assign BD IDs to NPU DMA ops and refactor `AMDAIEAssignBufferDescriptorIDsPass` #551

Add pass to assign BD IDs to NPU DMA ops and refactor `AMDAIEAssignBufferDescriptorIDsPass` #551

jtuyls commented Jul 15, 2024 •

edited

Loading

makslevental Jul 16, 2024 •

edited

Loading

jtuyls Jul 16, 2024 •

edited

Loading

makslevental Jul 16, 2024 •

edited

Loading