AIE: dynamic objectFifo lock lowering (SSA lock values)

Make aie.use_lock carry its lock count as an i32 SSA operand instead of
an optional integer attribute, and emit runtime-computed lock counts in
the dynamic (loop-preserving) objectFifo lowering via a per-(fifo,port)
"held" counter. Passes that need a compile-time constant read it back
through UseLockOp::getConstantValue().

The static loop-unrolling path (dynamic-objFifos=false) is left in place
and still emits compile-time constant lock counts, so static and dynamic
lock lowering coexist on this branch.

Compiler code only; the corresponding test updates follow in the next
commit.

Author: andrej

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

@@ -1378,34 +1378,49 @@ def AIE_UseLockOp: AIE_Op<"use_lock", []> {
       Then, the value of the lock is decremented by `value`.
     - Release: In AIE1, set the lock to `value`.
       In AIE2, increment the lock by `value`.
+
+    The lock value is always given as an `i32` SSA `value` operand. Passes that
+    need a compile-time constant obtain it via `getConstantValue()`, which
+    expects the operand to be defined by an `arith.constant` and fails
+    otherwise.
   }];
 
   let arguments = (
     ins Index:$lock,
         LockAction:$action,
-        OptionalAttr<AIEI32Attr>:$value,
+        I32:$value,
         OptionalAttr<LockBlocking>:$blocking,
         DefaultValuedOptionalAttr<BoolAttr, "true">:$acq_en
   );
 
   let assemblyFormat = [{
-    `(` $lock `,` $action (`,` $value^)? (`,` $blocking^)? `)` attr-dict
+    `(` $lock `,` custom<UseLockValue>($action, $value, $blocking) `)` attr-dict
   }];
 
   let hasVerifier = 1;
   let builders = [
     OpBuilder<(ins "mlir::Value":$lock,
                    "xilinx::AIE::LockAction":$action,
                    "int32_t":$value), [{
-      build($_builder, $_state, lock, action, $_builder.getI32IntegerAttr(value), nullptr);
+      auto constant = mlir::arith::ConstantOp::create(
+          $_builder, $_state.location, $_builder.getI32Type(),
+          $_builder.getI32IntegerAttr(value));
+      build($_builder, $_state, lock, action, constant, /*blocking=*/nullptr);
+    }]>,
+    OpBuilder<(ins "mlir::Value":$lock,
+                   "xilinx::AIE::LockAction":$action,
+                   "mlir::Value":$value), [{
+      build($_builder, $_state, lock, action, value, /*blocking=*/nullptr);
     }]>
   ];
 
   let extraClassDeclaration = [{
     bool acquire() { return (getAction() == LockAction::Acquire); }
     bool acquireGE() { return (getAction() == LockAction::AcquireGreaterEqual); }
     bool release() { return (getAction() == LockAction::Release); }
-    int getLockValue() { return getValue().value_or(1); }
+    // Returns the compile-time constant lock value when `value` is defined by
+    // an arith.constant; otherwise emits an op error and returns failure.
+    mlir::FailureOr<int32_t> getConstantValue();
     int getTimeout() {
       // LockBlocking is an EnumAttr.
       if (auto val = getBlocking())

lib/Dialect/AIE/IR/AIEDialect.cpp

@@ -8,9 +8,11 @@
 
 #include "aie/Dialect/AIE/IR/AIEDialect.h"
 
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/FoldInterfaces.h"
@@ -73,7 +75,15 @@ struct AIEDialectFoldInterface : DialectFoldInterface {
     //    them with a core's identical constants, which would leave the core
     //    referencing a device-level value that is erased when the core is
     //    outlined.
-    return isa<CoreOp, RuntimeSequenceOp>(region->getParentOp());
+    //  - aie.mem/aie.memtile_dma/aie.shim_dma bodies carry constant lock values
+    //    (aie.use_lock operands). If these hoisted to the device, CSE would
+    //    merge them with the identical constant kept local to a core, and since
+    //    the device-level constant dominates the core, the core's use would be
+    //    rewritten to reference it -- reintroducing the cross-region reference
+    //    that breaks core outlining. Keeping them local avoids the dominating
+    //    device-level duplicate entirely.
+    return isa<CoreOp, RuntimeSequenceOp, MemOp, MemTileDMAOp, ShimDMAOp>(
+        region->getParentOp());
   }
 };
 
@@ -2679,6 +2689,17 @@ LogicalResult LockOp::verify() {
   return success();
 }
 
+// Centralized lookup for the compile-time constant lock value: returns the
+// integer if `op`'s value operand is defined by an arith.constant, otherwise
+// std::nullopt. This is the single place that knows how a static lock value is
+// recovered from the SSA operand.
+static std::optional<int32_t> getConstantLockValue(UseLockOp op) {
+  if (auto constant = op.getValue().getDefiningOp<arith::ConstantOp>())
+    if (auto intAttr = llvm::dyn_cast<IntegerAttr>(constant.getValue()))
+      return (int32_t)intAttr.getInt();
+  return std::nullopt;
+}
+
 struct UsesOneLockInDMABlock {
   static LogicalResult verifyTrait(Operation *op) {
     auto *block = op->getBlock();
@@ -2700,16 +2721,21 @@ struct AcquireReleaseOneStateInDMABlock {
     auto *block = op->getBlock();
     int acqValue = -1, relValue = -1;
     for (auto op : block->getOps<UseLockOp>()) {
+      // Non-constant lock values cannot be compared here; the passes that
+      // require a constant enforce that separately via getConstantValue().
+      auto value = getConstantLockValue(op);
+      if (!value)
+        continue;
       if (op.acquire() || op.acquireGE()) {
-        if (acqValue != -1 && acqValue != op.getLockValue()) {
+        if (acqValue != -1 && acqValue != *value) {
           return failure();
         }
-        acqValue = op.getLockValue();
+        acqValue = *value;
       } else if (op.release()) {
-        if (relValue != -1 && relValue != op.getLockValue()) {
+        if (relValue != -1 && relValue != *value) {
           return failure();
         }
-        relValue = op.getLockValue();
+        relValue = *value;
       }
     }
     return success();
@@ -2737,6 +2763,15 @@ LogicalResult UseLockOp::verify() {
     return (*this)->emitOpError(
         "AcquireGreaterEqual is not supported in AIE1.");
 
+  // Locks used inside a DMA/BD block are configured via MMIO and therefore
+  // require a compile-time constant value (an arith.constant operand).
+  if (HasSomeParent<MemOp, MemTileDMAOp, ShimDMAOp>::verifyTrait(*this)
+          .succeeded() &&
+      !getConstantLockValue(*this))
+    return (*this)->emitOpError(
+        "lock value in a DMA/BD block must be a compile-time constant "
+        "(defined by an arith.constant).");
+
   // Otherwise, AIE.useLock should be inside MemOp, MemTileDMAOp, or
   // ShimDMAOp,
   if (HasSomeParent<MemOp, MemTileDMAOp, ShimDMAOp>::verifyTrait(*this)
@@ -2821,9 +2856,80 @@ static void printTraceRegValue(OpAsmPrinter &printer, Operation *op,
   xilinx::AIE::printTraceEventEnum(printer, value);
 }
 
+// Helper to parse a LockBlocking enum keyword.
+static ParseResult parseLockBlockingKeyword(OpAsmParser &parser,
+                                            xilinx::AIE::LockBlockingAttr &blocking) {
+  StringRef keyword;
+  if (parser.parseKeyword(&keyword))
+    return failure();
+  auto symbol = xilinx::AIE::symbolizeLockBlocking(keyword);
+  if (!symbol)
+    return parser.emitError(parser.getCurrentLocation())
+           << "invalid lock blocking value: " << keyword;
+  blocking =
+      xilinx::AIE::LockBlockingAttr::get(parser.getContext(), *symbol);
+  return success();
+}
+
+// Custom parser for the value slot of aie.use_lock. The lock value is a
+// required `i32` SSA operand, optionally followed by a `blocking` keyword. The
+// action keyword is parsed here as well so that the printed form has no stray
+// space before the comma.
+static ParseResult
+parseUseLockValue(OpAsmParser &parser, xilinx::AIE::LockActionAttr &action,
+                  OpAsmParser::UnresolvedOperand &value,
+                  xilinx::AIE::LockBlockingAttr &blocking) {
+  // The action is accepted either as a bare keyword (`Acquire`) or, for
+  // backward compatibility with older textual IR, as a quoted string
+  // (`"Acquire"`).
+  StringRef actionKeyword;
+  std::string actionString;
+  if (succeeded(parser.parseOptionalKeyword(&actionKeyword))) {
+    // Parsed a bare keyword.
+  } else if (succeeded(parser.parseOptionalString(&actionString))) {
+    actionKeyword = actionString;
+  } else {
+    return parser.emitError(parser.getCurrentLocation())
+           << "expected lock action keyword";
+  }
+  auto actionSymbol = xilinx::AIE::symbolizeLockAction(actionKeyword);
+  if (!actionSymbol)
+    return parser.emitError(parser.getCurrentLocation())
+           << "invalid lock action: " << actionKeyword;
+  action = xilinx::AIE::LockActionAttr::get(parser.getContext(), *actionSymbol);
+
+  // The lock value SSA operand.
+  if (parser.parseComma() || parser.parseOperand(value))
+    return failure();
+
+  // Optionally parse a trailing blocking keyword.
+  if (!parser.parseOptionalComma())
+    return parseLockBlockingKeyword(parser, blocking);
+  return success();
+}
+
+// Custom printer for the value slot of aie.use_lock.
+static void printUseLockValue(OpAsmPrinter &printer, Operation *op,
+                              xilinx::AIE::LockActionAttr action, Value value,
+                              xilinx::AIE::LockBlockingAttr blocking) {
+  printer << xilinx::AIE::stringifyLockAction(action.getValue());
+  printer << ", ";
+  printer.printOperand(value);
+  if (blocking)
+    printer << ", "
+            << xilinx::AIE::stringifyLockBlocking(blocking.getValue());
+}
+
 #define GET_OP_CLASSES
 #include "aie/Dialect/AIE/IR/AIEOps.cpp.inc"
 
+FailureOr<int32_t> UseLockOp::getConstantValue() {
+  if (auto value = getConstantLockValue(*this))
+    return *value;
+  return emitOpError("expected the lock value to be a compile-time constant "
+                     "(defined by an arith.constant).");
+}
+
 TileOp SwitchboxOp::getTileOp() {
   return cast<TileElement>(this->getOperation()).getTileOp();
 }

lib/Dialect/AIE/IR/CMakeLists.txt

@@ -21,6 +21,7 @@ add_mlir_dialect_library(AIE
 
   LINK_LIBS PUBLIC
   MLIRAIEUtil
+  MLIRArithDialect
   MLIRIR
   MLIRSupport
 )

lib/Dialect/AIE/Transforms/AIECoreToStandard.cpp

@@ -443,19 +443,19 @@ struct AIEUseLockToStdLowering : OpConversionPattern<UseLockOp> {
         return useLock.emitOpError("Could not find the intrinsic function!");
 
       SmallVector<Value, 2> args;
-      auto lockValue = useLock.getLockValue();
+      auto i32Ty = IntegerType::get(rewriter.getContext(), 32);
+      args.push_back(arith::IndexCastOp::create(rewriter, useLock.getLoc(),
+                                                i32Ty, useLock.getLock()));
 
-      // AIE2 acquire greater equal is encoded as a negative value.
+      Value value = adaptor.getValue();
+      // AIE2 acquire-greater-equal is encoded as a negative value, so negate
+      // it at runtime.
       if (useLock.acquireGE()) {
-        lockValue = -lockValue;
+        Value zero = arith::ConstantOp::create(
+            rewriter, useLock.getLoc(), i32Ty, rewriter.getI32IntegerAttr(0));
+        value = arith::SubIOp::create(rewriter, useLock.getLoc(), zero, value);
       }
-      args.push_back(arith::IndexCastOp::create(
-          rewriter, useLock.getLoc(),
-          IntegerType::get(rewriter.getContext(), 32), useLock.getLock()));
-      args.push_back(
-          arith::ConstantOp::create(rewriter, useLock.getLoc(),
-                                    IntegerType::get(rewriter.getContext(), 32),
-                                    rewriter.getI32IntegerAttr(lockValue)));
+      args.push_back(value);
 
       func::CallOp::create(rewriter, useLock.getLoc(), useLockFunc, args);
     }