Fix overflow in number of bytes for long DMA BD lengths (#3034)

Author: andrej

cmake/modulesXilinx

@@ -1025,15 +1025,20 @@ def AIE_DMABDOp: AIE_Op<"dma_bd", []> {
       auto bufferType = llvm::cast<mlir::BaseMemRefType>(getBuffer().getType());
       return dataLayout.getTypeSize(bufferType.getElementType());
     }
-    uint32_t getLenInBytes() {
-      if (std::optional<uint32_t> len = getLen(); len.has_value())
-        return len.value() * getBufferElementTypeWidthInBytes();
+    uint64_t getLenInBytes() {
+      if (std::optional<int32_t> len = getLen(); len.has_value())
+        return static_cast<uint64_t>(static_cast<uint32_t>(len.value())) *
+               static_cast<uint64_t>(getBufferElementTypeWidthInBytes());
       auto memrefType = llvm::dyn_cast<mlir::MemRefType>(getBuffer().getType());
       if (!memrefType)
         return 0; // Unranked memref without explicit len
-      return memrefType.getNumElements() * getBufferElementTypeWidthInBytes();
+      return static_cast<uint64_t>(memrefType.getNumElements()) *
+             static_cast<uint64_t>(getBufferElementTypeWidthInBytes());
+    }
+    int64_t getOffsetInBytes() {
+      return static_cast<int64_t>(getOffset()) *
+             static_cast<int64_t>(getBufferElementTypeWidthInBytes());
     }
-    int32_t getOffsetInBytes() { return getOffset() * getBufferElementTypeWidthInBytes(); }
   }];
 
   let hasVerifier = 1;

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

@@ -34,7 +34,7 @@ std::string packetStr(int id, int type);
 void generateXAieDmaSetMultiDimAddr(llvm::raw_ostream &output, int ndims,
                                     llvm::ArrayRef<BDDimLayoutAttr> dims,
                                     int col, int row, int bdNum, int baseAddrA,
-                                    int offsetA, int lenA,
+                                    int64_t offsetA, uint32_t lenA,
                                     int elementWidthInBytes,
                                     const char *errorRet);
 

include/aie/Targets/AIETargetShared.h

@@ -2294,7 +2294,7 @@ LogicalResult DMABDOp::verify() {
   if (!isUnrankedMemRef &&
       (parentTile.isMemTile() || parentTile.isCoreTile())) {
     if (auto baseAddr = getBufferOp().getAddress(); baseAddr.has_value()) {
-      int offsetInBytes = *baseAddr + getOffsetInBytes();
+      int64_t offsetInBytes = *baseAddr + getOffsetInBytes();
       if (offsetInBytes % 4)
         return emitOpError("bd address must be 4 byte (32b) aligned; got "
                            "base+offset: ")

lib/Dialect/AIE/IR/AIEDialect.cpp

@@ -86,7 +86,7 @@ static std::string tileDMATensorStr(int col, int row, int bdNum) {
 void generateXAieDmaSetMultiDimAddr(raw_ostream &output, int ndims,
                                     ArrayRef<BDDimLayoutAttr> dims, int col,
                                     int row, int bdNum, int baseAddrA,
-                                    int offsetA, int lenA,
+                                    int64_t offsetA, uint32_t lenA,
                                     int elementWidthInBytes,
                                     const char *errorRetval) {
   // libxaie requires stride in multiples of 32b

lib/Targets/AIETargetShared.cpp

@@ -22,6 +22,8 @@
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/IR/Module.h"
 
+#include <limits>
+
 using namespace mlir;
 using namespace xilinx;
 using namespace xilinx::AIE;
@@ -103,8 +105,8 @@ static mlir::LogicalResult generateDMAConfig(OpType memOp, raw_ostream &output,
     int packetType = 0;
     int packetID = 0;
     bool foundBd = false;
-    int lenA = 0;
-    int offsetA = 0;
+    uint64_t lenA = 0;
+    int64_t offsetA = 0;
     int BaseAddrA = 0;
     int elementWidthInBytes = 0;
     int ndims = 0;
@@ -142,6 +144,12 @@ static mlir::LogicalResult generateDMAConfig(OpType memOp, raw_ostream &output,
         return memOp.emitOpError("DMA contains at least one multi-dimensional "
                                  "buffer descriptor. This is currently only "
                                  "supported for AIE-ML and later devices.");
+    // Ensure BD length in bytes fits in the u32 `Len` argument of the libxaie
+    // DMA APIs (XAie_DmaSetAddrLen / XAie_DmaSetMultiDimAddr).
+    if (foundBd && lenA > std::numeric_limits<uint32_t>::max()) {
+      return memOp.emitOpError("buffer descriptor length in bytes (")
+             << lenA << ") does not fit in 32 bits";
+    }
 
     int acqValue = 0, relValue = 0;
     bool hasAcq = false, hasRel = false;

lib/Targets/AIETargetXAIEV2.cpp

@@ -0,0 +1,35 @@
+//===- test_error_dma_len_overflow.mlir ------------------------*- MLIR -*-===//
+//
+// 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 Advanced Micro Devices, Inc. or its affiliates
+//
+//===----------------------------------------------------------------------===//
+
+// RUN: aie-translate --aie-generate-xaie --verify-diagnostics %s
+
+// `XAie_DmaSetAddrLen` / `XAie_DmaSetMultiDimAddr` take the byte length as a
+// `u32`, so any length that does not fit in 32 bits would be silently
+// truncated by the emitted C call. Verify the translator rejects it with an
+// explicit op error instead.
+//
+// memref<536870912xi64>: 536870912 * 8 = 4294967296 bytes = 0x100000000 which
+// does not fit in 32 bits.
+
+module {
+  aie.device(npu1) {
+    %tile_0_0 = aie.tile(0, 0)
+    // expected-error @+1 {{does not fit in 32 bits}}
+    aie.shim_dma(%tile_0_0) {
+      %srcDma = aie.dma_start(MM2S, 0, ^bd0, ^end)
+    ^bd0:
+      %buf = aie.external_buffer { sym_name = "myBuffer_0_0_0" } : memref<536870912xi64>
+      aie.dma_bd(%buf : memref<536870912xi64>, 0, 536870912)
+      aie.next_bd ^end
+    ^end:
+      aie.end
+    }
+  }
+}

test/Targets/AIEGenerateXAIE/test_error_dma_len_overflow.mlir

@@ -0,0 +1,34 @@
+//
+// 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 AMD Inc.
+
+// RUN: aie-opt --verify-diagnostics --aie-dma-tasks-to-npu %s
+
+// This test verifies that when a large-length dma_bd has mismatched dimensions,
+// the error message correctly reports the byte count without integer overflow.
+//
+// memref<536870912xi64>: 536870912 elements * 8 bytes = 4294967296 bytes.
+// BD length in 32-bit words = 4294967296 / 4 = 1073741824.
+// The lowest three dims product: sizes[0] = 2*64/32 = 4 words, sizes[1] = sizes[2] = 1.
+// Total = 4 words = 16 bytes, which does not match 4294967296 bytes.
+// The error message must report "4294967296 bytes" (not a truncated 32-bit value
+// like "0 bytes" that the overflow bug would produce) and "16 bytes" for the dims.
+
+module {
+  aie.device(npu1) {
+    %tile_0_0 = aie.tile(0, 0)
+
+    aie.runtime_sequence(%arg0: memref<536870912xi64>) {
+      %t1 = aiex.dma_configure_task(%tile_0_0, MM2S, 0) {
+          // expected-error@+2 {{Buffer descriptor length does not match length of transfer expressed by lowest three dimensions of data layout transformation strides/wraps. BD length is 4294967296 bytes. Lowest three dimensions of data layout transformation would result in transfer of 16 bytes.}}
+          // expected-note@+1 {{}}
+          aie.dma_bd(%arg0 : memref<536870912xi64>, 0, 536870912,
+                     [<size=1, stride=1>, <size=1, stride=1>, <size=2, stride=1>]) {bd_id = 0 : i32}
+          aie.end
+      }
+    }
+  }
+}

test/bd-chains-and-dma-tasks/dma-tasks-to-npu/overflow-len-bad.mlir

@@ -0,0 +1,36 @@
+//
+// 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 AMD Inc.
+
+// RUN: aie-opt --aie-dma-tasks-to-npu %s | FileCheck %s
+
+// This test verifies that a dma_bd with a transfer length that would overflow a
+// 32-bit integer when converted to bytes is handled correctly. Without the fix,
+// `len * element_size_bytes` would overflow to 0 or a wrong value, causing a
+// false "Buffer descriptor length does not match" error.
+//
+// memref<536870912xi64>: 536870912 elements * 8 bytes = 4294967296 bytes.
+// In 32-bit words: 4294967296 / 4 = 1073741824 (= 0x40000000).
+// Without overflow fix, the byte count truncates to 0, so len_addr_granularity
+// would be 0 instead of 1073741824, triggering a spurious validation error.
+
+// CHECK-LABEL: aie.runtime_sequence
+// CHECK: aiex.npu.writebd
+// CHECK-SAME: buffer_length = 1073741824 : i32
+
+module {
+  aie.device(npu1) {
+    %tile_0_0 = aie.tile(0, 0)
+
+    aie.runtime_sequence(%arg0: memref<536870912xi64>) {
+      %t1 = aiex.dma_configure_task(%tile_0_0, MM2S, 0) {
+          aie.dma_bd(%arg0 : memref<536870912xi64>, 0, 536870912,
+                     [<size=1, stride=1>, <size=1, stride=1>, <size=536870912, stride=1>]) {bd_id = 0 : i32}
+          aie.end
+      }
+    }
+  }
+}