Forward source locations across AIE/AIEX transform passes

lib/Conversion/AIEToConfiguration/AIEToConfiguration.cpp

@@ -450,12 +450,10 @@ static LogicalResult generateTransactions(AIERTControl &ctl,
 // Translate vector of TransactionBinaryOperation to a sequence of transaction
 // ops (npu.write32, npu.maskwrite32, npu.blockwrite).
 static LogicalResult
-emitTransactionOps(OpBuilder &builder,
+emitTransactionOps(OpBuilder &builder, Location loc,
                    std::vector<TransactionBinaryOperation> &operations,
                    std::vector<memref::GlobalOp> &global_data) {
 
-  auto loc = builder.getUnknownLoc();
-
   // create the txn ops
   for (auto [op, payload] : llvm::zip(operations, global_data)) {
 
@@ -531,11 +529,10 @@ emitTransactionOps(OpBuilder &builder,
 // Translate vector of TransactionBinaryOperation to a sequence of control
 // packet ops.
 static LogicalResult
-emitControlPacketOps(OpBuilder &builder,
+emitControlPacketOps(OpBuilder &builder, Location loc,
                      std::vector<TransactionBinaryOperation> &operations,
                      std::vector<memref::GlobalOp> &global_data) {
 
-  auto loc = builder.getUnknownLoc();
   auto ctx = builder.getContext();
 
   // create the control packet ops
@@ -636,17 +633,21 @@ static LogicalResult convertTransactionOpsToMLIR(
     std::vector<TransactionBinaryOperation> &operations,
     std::string blockwrite_prefix = "config_blockwrite_data_") {
 
-  auto loc = builder.getUnknownLoc();
-
   // for each blockwrite in the binary, create a GlobalOp with the data at the
   // device level
   std::vector<memref::GlobalOp> global_data;
+  Operation *parentOp = builder.getBlock()->getParentOp();
+  DeviceOp device = llvm::dyn_cast<DeviceOp>(parentOp);
+  if (!device) {
+    device = parentOp->getParentOfType<DeviceOp>();
+  }
+  if (!device) {
+    parentOp->emitError(
+        "expected insertion point to be nested under an aie.device op");
+    return failure();
+  }
+  Location loc = device.getLoc();
   {
-    DeviceOp device =
-        llvm::dyn_cast<DeviceOp>(builder.getBlock()->getParentOp());
-    if (!device) {
-      device = builder.getBlock()->getParentOp()->getParentOfType<DeviceOp>();
-    }
     OpBuilder::InsertionGuard guard(builder);
     builder.setInsertionPointToStart(device.getBody());
     int id = 0;
@@ -676,10 +677,10 @@ static LogicalResult convertTransactionOpsToMLIR(
 
   // create the txn ops
   if (outputType == AIE::AIEToConfigurationOutputType::Transaction) {
-    if (failed(emitTransactionOps(builder, operations, global_data)))
+    if (failed(emitTransactionOps(builder, loc, operations, global_data)))
       return failure();
   } else if (outputType == AIE::AIEToConfigurationOutputType::ControlPacket) {
-    if (failed(emitControlPacketOps(builder, operations, global_data)))
+    if (failed(emitControlPacketOps(builder, loc, operations, global_data)))
       return failure();
     // resolve mask writes; control packet doesn't natively support mask write.
     if (failed(orConsecutiveWritesOnSameAddr(builder.getBlock())))
@@ -784,7 +785,7 @@ convertAIEToConfiguration(AIE::DeviceOp device, StringRef clElfDir,
   // and collect them in a vector. If there are none, create a new runtime
   // sequence. Otherwise assume the insertion point is the first
   // aiex.configure op.
-  auto loc = builder.getUnknownLoc();
+  auto loc = device.getLoc();
   SmallVector<AIEX::ConfigureOp> configureOps;
   device.walk([&](AIEX::ConfigureOp op) { configureOps.push_back(op); });
 

lib/Dialect/AIE/IR/AIEDialect.cpp

@@ -1610,8 +1610,8 @@ TileOp TileOp::getOrCreate(mlir::OpBuilder builder, DeviceOp device, int col,
     OpBuilder::InsertionGuard guard(builder);
     mlir::Block &device_start_block = *device.getBodyRegion().begin();
     builder.setInsertionPointToStart(&device_start_block);
-    tile = TileOp::create(builder, builder.getUnknownLoc(),
-                          builder.getIndexType(), col, row);
+    tile = TileOp::create(builder, device.getLoc(), builder.getIndexType(), col,
+                          row);
   }
   return tile;
 }

lib/Dialect/AIE/Transforms/AIECanonicalizeDevice.cpp

@@ -42,7 +42,7 @@ struct AIECanonicalizeDevicePass
     // the new op quite yet.
     OpBuilder builder(moduleOp->getContext());
 
-    Location location = builder.getUnknownLoc();
+    Location location = moduleOp.getLoc();
     auto deviceOp = DeviceOp::create(
         builder, location,
         AIEDeviceAttr::get(builder.getContext(), AIEDevice::xcvc1902),

lib/Dialect/AIE/Transforms/AIECoreToStandard.cpp

@@ -211,7 +211,7 @@ struct AIEDebugOpToStdLowering : OpConversionPattern<DebugOp> {
              << funcName;
     SmallVector<Value, 1> args;
     args.push_back(op.getArg());
-    func::CallOp::create(rewriter, rewriter.getUnknownLoc(), func, args);
+    func::CallOp::create(rewriter, op.getLoc(), func, args);
     rewriter.eraseOp(op);
     return success();
   }
@@ -259,7 +259,7 @@ struct AIEPutStreamToStdLowering : OpConversionPattern<PutStreamOp> {
           rewriter, op.getLoc(), IntegerType::get(rewriter.getContext(), 32),
           rewriter.getI32IntegerAttr(0))); // tlast
     }
-    func::CallOp::create(rewriter, rewriter.getUnknownLoc(), putMSFunc, args);
+    func::CallOp::create(rewriter, op.getLoc(), putMSFunc, args);
     rewriter.eraseOp(op);
     return success();
   }
@@ -300,8 +300,8 @@ struct AIEGetStreamToStdLowering : OpConversionPattern<GetStreamOp> {
     SmallVector<Value, 2> args;
     if (targetModel.getTargetArch() == AIEArch::AIE1)
       args.push_back(op.getChannel());
-    auto getSSCall = func::CallOp::create(rewriter, rewriter.getUnknownLoc(),
-                                          getSSFunc, args);
+    auto getSSCall =
+        func::CallOp::create(rewriter, op.getLoc(), getSSFunc, args);
     rewriter.replaceOp(op, getSSCall.getResult(0));
     // Capture TLAST in AIEv2?
     return success();
@@ -352,7 +352,7 @@ struct AIEPutCascadeToStdLowering : OpConversionPattern<PutCascadeOp> {
           rewriter, op.getLoc(), IntegerType::get(rewriter.getContext(), 32),
           rewriter.getI32IntegerAttr(1))); // enable
 
-    func::CallOp::create(rewriter, rewriter.getUnknownLoc(), putMCDFunc, args);
+    func::CallOp::create(rewriter, op.getLoc(), putMCDFunc, args);
     rewriter.eraseOp(op);
     return success();
   }
@@ -388,8 +388,8 @@ struct AIEGetCascadeToStdLowering : OpConversionPattern<GetCascadeOp> {
           rewriter, op.getLoc(), IntegerType::get(rewriter.getContext(), 32),
           rewriter.getI32IntegerAttr(1))); // enable
 
-    auto getSCDCall = func::CallOp::create(rewriter, rewriter.getUnknownLoc(),
-                                           getSCDFunc, args);
+    auto getSCDCall =
+        func::CallOp::create(rewriter, op.getLoc(), getSCDFunc, args);
     Value result = getSCDCall.getResult(0);
 
     // Check if we need a bitcast
@@ -459,8 +459,7 @@ struct AIEUseLockToStdLowering : OpConversionPattern<UseLockOp> {
                                     IntegerType::get(rewriter.getContext(), 32),
                                     rewriter.getI32IntegerAttr(lockValue)));
 
-      func::CallOp::create(rewriter, rewriter.getUnknownLoc(), useLockFunc,
-                           args);
+      func::CallOp::create(rewriter, useLock.getLoc(), useLockFunc, args);
     }
     rewriter.eraseOp(useLock);
     return success();
@@ -490,19 +489,19 @@ struct AIEBufferToStandard : OpConversionPattern<BufferOp> {
     // prevent duplication in the data section of the elf/object file)
     if ((tileRow != row && tileRow != -1) || (tileCol != col && tileCol != -1))
       initValue = nullptr;
-    memref::GlobalOp::create(rewriter, rewriter.getUnknownLoc(), symName,
+    memref::GlobalOp::create(rewriter, buffer.getLoc(), symName,
                              rewriter.getStringAttr("public"), buffer.getType(),
                              initValue, /*constant*/ false,
                              /*alignment*/ nullptr);
 
     for (auto &use : make_early_inc_range(buffer.getResult().getUses())) {
       Operation *user = use.getOwner();
       rewriter.setInsertionPoint(user);
-      auto allocated = memref::GetGlobalOp::create(
-          rewriter, rewriter.getUnknownLoc(), t, symName);
+      auto allocated =
+          memref::GetGlobalOp::create(rewriter, buffer.getLoc(), t, symName);
       // Assume that buffers are aligned so they can be vectorized.
-      memref::AssumeAlignmentOp::create(rewriter, rewriter.getUnknownLoc(),
-                                        allocated, 32);
+      memref::AssumeAlignmentOp::create(rewriter, buffer.getLoc(), allocated,
+                                        32);
 
       use.set(allocated.getResult());
     }
@@ -547,7 +546,7 @@ struct AIECoreToStandardFunc : OpConversionPattern<CoreOp> {
     std::string coreName("core_" + std::to_string(col) + "_" +
                          std::to_string(row));
     auto coreFunc =
-        func::FuncOp::create(rewriter, rewriter.getUnknownLoc(), coreName,
+        func::FuncOp::create(rewriter, op.getLoc(), coreName,
                              FunctionType::get(rewriter.getContext(), {}, {}));
 
     rewriter.cloneRegionBefore(op.getBody(), coreFunc.getBody(),
@@ -641,8 +640,7 @@ struct AIECoreToStandardFunc : OpConversionPattern<CoreOp> {
       rewriter.setInsertionPointAfter(childOp);
 
       if (isa<EndOp>(childOp)) {
-        func::ReturnOp::create(rewriter, rewriter.getUnknownLoc(),
-                               ValueRange({}));
+        func::ReturnOp::create(rewriter, childOp->getLoc(), ValueRange({}));
         rewriter.eraseOp(childOp);
       }
     });
@@ -701,7 +699,7 @@ struct AIEEventOpToStdLowering : OpConversionPattern<EventOp> {
     if (!eventFunc)
       return op.emitOpError("Could not find the intrinsic function ")
              << funcName;
-    func::CallOp::create(rewriter, rewriter.getUnknownLoc(), eventFunc, args);
+    func::CallOp::create(rewriter, op.getLoc(), eventFunc, args);
     rewriter.eraseOp(op);
     return success();
   }

lib/Dialect/AIE/Transforms/AIECreatePathFindFlows.cpp

@@ -46,8 +46,8 @@ struct ConvertFlowsToInterconnect : OpConversionPattern<FlowOp> {
     auto point = rewriter.saveInsertionPoint();
     rewriter.setInsertionPoint(b.getTerminator());
 
-    ConnectOp::create(rewriter, rewriter.getUnknownLoc(), inBundle, inIndex,
-                      outBundle, outIndex);
+    ConnectOp::create(rewriter, flowOp.getLoc(), inBundle, inIndex, outBundle,
+                      outIndex);
 
     rewriter.restoreInsertionPoint(point);
 
@@ -832,13 +832,13 @@ AIEPathfinderPass::runOnPacketFlow(DeviceOp device, OpBuilder &builder,
     Operation *tileOp = map.second;
     TileOp tile = cast<TileOp>(map.second);
     TileID tileId = tile.getTileID();
+    Location tileLoc = tile.getLoc();
 
     // Create a switchbox for the routes and insert inside it.
     builder.setInsertionPointAfter(tileOp);
     SwitchboxOp swbox =
         analyzer.getSwitchbox(builder, tile.colIndex(), tile.rowIndex());
-    SwitchboxOp::ensureTerminator(swbox.getConnections(), builder,
-                                  builder.getUnknownLoc());
+    SwitchboxOp::ensureTerminator(swbox.getConnections(), builder, tileLoc);
     Block &b = swbox.getConnections().front();
     builder.setInsertionPoint(b.getTerminator());
 
@@ -859,8 +859,7 @@ AIEPathfinderPass::runOnPacketFlow(DeviceOp device, OpBuilder &builder,
         if (amselOpNeededVector[amselValue]) {
           int arbiterID = a;
           int msel = i;
-          auto amsel = AMSelOp::create(builder, builder.getUnknownLoc(),
-                                       arbiterID, msel);
+          auto amsel = AMSelOp::create(builder, tileLoc, arbiterID, msel);
           amselOps[amselValue] = amsel;
         }
       }
@@ -885,8 +884,8 @@ AIEPathfinderPass::runOnPacketFlow(DeviceOp device, OpBuilder &builder,
         amsels.push_back(amselOps[msel]);
       }
 
-      MasterSetOp::create(builder, builder.getUnknownLoc(),
-                          builder.getIndexType(), bundle, channel, amsels,
+      MasterSetOp::create(builder, tileLoc, builder.getIndexType(), bundle,
+                          channel, amsels,
                           keepPktHeaderAttr[{tileId, tileMaster}]);
     }
 
@@ -915,10 +914,9 @@ AIEPathfinderPass::runOnPacketFlow(DeviceOp device, OpBuilder &builder,
 
       PacketRulesOp packetrules;
       if (slaveRules.count(slave) == 0) {
-        packetrules = PacketRulesOp::create(builder, builder.getUnknownLoc(),
-                                            bundle, channel);
+        packetrules = PacketRulesOp::create(builder, tileLoc, bundle, channel);
         PacketRulesOp::ensureTerminator(packetrules.getRules(), builder,
-                                        builder.getUnknownLoc());
+                                        tileLoc);
         slaveRules[slave] = packetrules;
       } else
         packetrules = slaveRules[slave];
@@ -939,7 +937,7 @@ AIEPathfinderPass::runOnPacketFlow(DeviceOp device, OpBuilder &builder,
       }
 
       builder.setInsertionPoint(rules.getTerminator());
-      PacketRuleOp::create(builder, builder.getUnknownLoc(), mask, ID, amsel);
+      PacketRuleOp::create(builder, tileLoc, mask, ID, amsel);
     }
   }
 
@@ -996,13 +994,13 @@ AIEPathfinderPass::runOnPacketFlow(DeviceOp device, OpBuilder &builder,
           pktrules.setSourceBundle(WireBundle::South);
           if (pktrules.getSourceChannel() == 0) {
             pktrules.setSourceChannel(3);
-            ConnectOp::create(builder, builder.getUnknownLoc(), WireBundle::DMA,
-                              0, WireBundle::North, 3);
+            ConnectOp::create(builder, tileOp.getLoc(), WireBundle::DMA, 0,
+                              WireBundle::North, 3);
           }
           if (pktrules.getSourceChannel() == 1) {
             pktrules.setSourceChannel(7);
-            ConnectOp::create(builder, builder.getUnknownLoc(), WireBundle::DMA,
-                              1, WireBundle::North, 7);
+            ConnectOp::create(builder, tileOp.getLoc(), WireBundle::DMA, 1,
+                              WireBundle::North, 7);
           }
         }
       }
@@ -1027,13 +1025,13 @@ AIEPathfinderPass::runOnPacketFlow(DeviceOp device, OpBuilder &builder,
           mtset.setDestBundle(WireBundle::South);
           if (mtset.getDestChannel() == 0) {
             mtset.setDestChannel(2);
-            ConnectOp::create(builder, builder.getUnknownLoc(),
-                              WireBundle::North, 2, WireBundle::DMA, 0);
+            ConnectOp::create(builder, tileOp.getLoc(), WireBundle::North, 2,
+                              WireBundle::DMA, 0);
           }
           if (mtset.getDestChannel() == 1) {
             mtset.setDestChannel(3);
-            ConnectOp::create(builder, builder.getUnknownLoc(),
-                              WireBundle::North, 3, WireBundle::DMA, 1);
+            ConnectOp::create(builder, tileOp.getLoc(), WireBundle::North, 3,
+                              WireBundle::DMA, 1);
           }
         }
       }
@@ -1084,54 +1082,55 @@ void AIEPathfinderPass::runOnOperation() {
         sw = analyzer.coordToSwitchbox[{col, row}];
       else
         continue;
+      Location loc = tile.getLoc();
       if (col > 0) {
         // connections east-west between stream switches
         if (analyzer.coordToSwitchbox.count({col - 1, row})) {
           auto westsw = analyzer.coordToSwitchbox[{col - 1, row}];
-          WireOp::create(builder, builder.getUnknownLoc(), westsw,
-                         WireBundle::East, sw, WireBundle::West);
+          WireOp::create(builder, loc, westsw, WireBundle::East, sw,
+                         WireBundle::West);
         }
       }
       if (row > 0) {
         // connections between abstract 'core' of tile
-        WireOp::create(builder, builder.getUnknownLoc(), tile, WireBundle::Core,
-                       sw, WireBundle::Core);
+        WireOp::create(builder, loc, tile, WireBundle::Core, sw,
+                       WireBundle::Core);
         // connections between abstract 'dma' of tile
-        WireOp::create(builder, builder.getUnknownLoc(), tile, WireBundle::DMA,
-                       sw, WireBundle::DMA);
+        WireOp::create(builder, loc, tile, WireBundle::DMA, sw,
+                       WireBundle::DMA);
         // connections north-south inside array ( including connection to shim
         // row)
         if (analyzer.coordToSwitchbox.count({col, row - 1})) {
           auto southsw = analyzer.coordToSwitchbox[{col, row - 1}];
-          WireOp::create(builder, builder.getUnknownLoc(), southsw,
-                         WireBundle::North, sw, WireBundle::South);
+          WireOp::create(builder, loc, southsw, WireBundle::North, sw,
+                         WireBundle::South);
         }
       } else if (row == 0) {
         if (tile.isShimNOCTile()) {
           if (analyzer.coordToShimMux.count({col, 0})) {
             auto shimsw = analyzer.coordToShimMux[{col, 0}];
             WireOp::create(
-                builder, builder.getUnknownLoc(), shimsw,
+                builder, loc, shimsw,
                 WireBundle::North, // Changed to connect into the north
                 sw, WireBundle::South);
             // PLIO is attached to shim mux
             if (analyzer.coordToPLIO.count(col)) {
               auto plio = analyzer.coordToPLIO[col];
-              WireOp::create(builder, builder.getUnknownLoc(), plio,
-                             WireBundle::North, shimsw, WireBundle::South);
+              WireOp::create(builder, loc, plio, WireBundle::North, shimsw,
+                             WireBundle::South);
             }
 
             // abstract 'DMA' connection on tile is attached to shim mux ( in
             // row 0 )
-            WireOp::create(builder, builder.getUnknownLoc(), tile,
-                           WireBundle::DMA, shimsw, WireBundle::DMA);
+            WireOp::create(builder, loc, tile, WireBundle::DMA, shimsw,
+                           WireBundle::DMA);
           }
         } else if (tile.isShimPLTile()) {
           // PLIO is attached directly to switch
           if (analyzer.coordToPLIO.count(col)) {
             auto plio = analyzer.coordToPLIO[col];
-            WireOp::create(builder, builder.getUnknownLoc(), plio,
-                           WireBundle::North, sw, WireBundle::South);
+            WireOp::create(builder, loc, plio, WireBundle::North, sw,
+                           WireBundle::South);
           }
         }
       }