Enable dynamic lock acquisition for ObjectFIFO acquire/release in cores (#3308)

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 `,` $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,11 @@ 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());
+    //  - Make sure SSA values for aie.use_lock operands in
+    //    aie.mem/aie.memtile_dma/aie.shim_dma bodies do not get
+    //    hoisted.
+    return isa<CoreOp, RuntimeSequenceOp, MemOp, MemTileDMAOp, ShimDMAOp>(
+        region->getParentOp());
   }
 };
 
@@ -2679,6 +2685,15 @@ LogicalResult LockOp::verify() {
   return success();
 }
 
+// Look up for compile-time constant lock values, if any.
+// Returns std::nullopt if lock value does not reference an `arith.constant`.
+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 +2715,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 +2757,15 @@ LogicalResult UseLockOp::verify() {
     return (*this)->emitOpError(
         "AcquireGreaterEqual is not supported in AIE1.");
 
+  // Locks used inside a DMA/BD block are configured via static register writes
+  // and therefore require a compile-time constant value.
+  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 +2850,17 @@ static void printTraceRegValue(OpAsmPrinter &printer, Operation *op,
   xilinx::AIE::printTraceEventEnum(printer, value);
 }
 
+// Helper to parse a LockBlocking enum keyword.
 #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);
     }