[GLUON] Set proper location on restoring the insert point in gluon (#8531)

`warp_specialize` ops currently have unknown location set in the TTGIR
due to a quirk in the code emission in `_semantic.py`: for
`warp_specialize` we need save and then restore insert point. Location
is being inferred from the insert point, however if insert point happens
to be in a place that doesn't have location assigned (end of a block),
we set unknown loc. This change is a minimal fix that adds a helper that
gets the location from block's parent in such a case.
Alternatively we could also save location along with insert point, and
then restore it accordingly. This approach is simpler and should help
for most cases I could have think of however.
This change is important for consan changes I am working on, as it
breaks the LLVM backend if we create instrumentation function calls with
unknown location inferred from warp_specialize op.

Author: pawelszczerbuk

python/src/ir.h

@@ -35,15 +35,15 @@ class TritonOpBuilder {
     if (!block.empty())
       setLastLoc(block.begin()->getLoc());
     else
-      setLastLoc(builder->getUnknownLoc());
+      setLastLoc(getLocForBlock(&block));
     builder->setInsertionPointToStart(&block);
   }
 
   void setInsertionPointToEnd(mlir::Block &block) {
     if (!block.empty())
       setLastLoc(block.back().getLoc());
     else
-      setLastLoc(builder->getUnknownLoc());
+      setLastLoc(getLocForBlock(&block));
     builder->setInsertionPointToEnd(&block);
   }
 
@@ -53,10 +53,14 @@ class TritonOpBuilder {
   }
 
   void restoreInsertionPoint(mlir::OpBuilder::InsertPoint pt) {
-    if (pt.isSet() && pt.getPoint() != pt.getBlock()->end())
-      setLastLoc(pt.getPoint()->getLoc());
-    else
-      setLastLoc(builder->getUnknownLoc());
+    setLastLoc(builder->getUnknownLoc());
+    if (pt.isSet()) {
+      if (pt.getPoint() != pt.getBlock()->end())
+        setLastLoc(pt.getPoint()->getLoc());
+      else
+        setLastLoc(getLocForBlock(pt.getBlock()));
+    }
+
     builder->restoreInsertionPoint(pt);
   }
 
@@ -87,4 +91,10 @@ class TritonOpBuilder {
   std::unique_ptr<mlir::Location> lastLoc;
   bool lineInfoEnabled =
       !mlir::triton::tools::getBoolEnv("TRITON_DISABLE_LINE_INFO");
+
+  mlir::Location getLocForBlock(mlir::Block *block) {
+    if (auto parentOp = block->getParentOp())
+      return parentOp->getLoc();
+    return builder->getUnknownLoc();
+  }
 };