Get file, line and col from MLIR instead of inspecting python frames. Function name will be always set to "unknown" for now. (#4797)

Currently, `device_assert` takes in separate `filename`, `funcname`,
`lineno` and it gets them by inspecting python frames in core.py. It's
clunky and in practice it is often unknown. Since the loc information is
already in the builder, it could be used directly, which is what this PR
does.

Author: arakhmati-openai

include/triton/Dialect/Triton/IR/TritonOps.td

@@ -844,11 +844,11 @@ def TT_PrintOp : TT_Op<"print", [SameVariadicOperandSize, MemoryEffects<[MemWrit
 def TT_AssertOp : TT_Op<"assert", [MemoryEffects<[MemWrite<GlobalMemory>]>]> {
   let summary = "Device-side assert, as in CUDA for correctness checking";
   let description = [{
-    `tt.assert` takes a condition tensor, a message string, a file string, a function string, and a line number.
+    `tt.assert` takes a condition tensor and a message string.
     If the condition is false, the message is printed, and the program is aborted.
   }];
-  let arguments = (ins TT_Tensor:$condition, StrAttr:$message, StrAttr:$file, StrAttr:$func, I32Attr:$line);
-  let assemblyFormat = "$condition `,` $message `,` $file `,` $func `,` $line attr-dict `:` type($condition)";
+  let arguments = (ins TT_Tensor:$condition, StrAttr:$message);
+  let assemblyFormat = "$condition `,` $message attr-dict `:` type($condition)";
 }
 
 //

lib/Conversion/TritonGPUToLLVM/AssertOpToLLVM.cpp

@@ -34,19 +34,30 @@ struct AssertOpConversion : public ConvertOpToLLVMPattern<triton::AssertOp> {
         return failure();
       }
     }
-    llAssert(op, condition, adaptor.getMessage(), adaptor.getFile(),
-             adaptor.getFunc(), adaptor.getLine(), rewriter);
+    llAssert(op, condition, adaptor.getMessage(), rewriter);
     rewriter.eraseOp(op);
     return success();
   }
   // op: the op at which the assert is inserted. Unlike printf, we need to
   // know about the op to split the block.
   void llAssert(Operation *op, Value condition, StringRef message,
-                StringRef file, StringRef func, int line,
                 ConversionPatternRewriter &rewriter) const {
     ConversionPatternRewriter::InsertionGuard guard(rewriter);
+
     auto ctx = rewriter.getContext();
     auto loc = op->getLoc();
+
+    StringRef file = "unknown";
+    StringRef func = "unknown";
+    int line = 0;
+    int col = 0;
+
+    if (auto fileLineColLoc = dyn_cast<FileLineColLoc>(loc)) {
+      file = fileLineColLoc.getFilename();
+      line = fileLineColLoc.getLine();
+      col = fileLineColLoc.getColumn();
+    }
+
     // #block1
     // if (condition) {
     //   #block2

python/src/ir.cc

@@ -1546,17 +1546,10 @@ void init_triton_ir(py::module &&m) {
            })
       .def("create_assert",
            [](TritonOpBuilder &self, Value &condition,
-              const std::string &message, const std::string &fileName,
-              const std::string &funcName, unsigned lineNo) -> void {
+              const std::string &message) -> void {
              auto messageAttr = StringAttr::get(self.getBuilder().getContext(),
                                                 llvm::StringRef(message));
-             auto fileNameAttr = StringAttr::get(self.getBuilder().getContext(),
-                                                 llvm::StringRef(fileName));
-             auto funcNameAttr = StringAttr::get(self.getBuilder().getContext(),
-                                                 llvm::StringRef(funcName));
-             auto lineNoAttr = self.getBuilder().getI32IntegerAttr(lineNo);
-             self.create<AssertOp>(condition, messageAttr, fileNameAttr,
-                                   funcNameAttr, lineNoAttr);
+             self.create<AssertOp>(condition, messageAttr);
            })
       .def("create_assume",
            [](TritonOpBuilder &self, Value &condition) {

python/triton/language/core.py

@@ -2308,25 +2308,7 @@ def device_assert(cond, msg="", _builder=None):
     :param msg: the message to print if the assertion fails. This is required to be a string literal.
     '''
     msg = _constexpr_to_value(msg)
-    import inspect
-    frame = inspect.currentframe()
-    module = inspect.getmodule(frame)
-    # The triton function module doesn't have the name attribute.
-    # We use this trick to find the caller.
-    while hasattr(module, "__name__"):
-        frame = frame.f_back
-        module = inspect.getmodule(frame)
-    lineno = 0
-    func_name = 'unknown'
-    file_name = 'unknown'
-    if frame is not None and frame.f_back is not None:
-        func_name = frame.f_code.co_name
-        file_name = frame.f_back.f_code.co_filename
-        # TODO: The line number currently indicates the line
-        # where the triton function is called but not where the
-        # device_assert is called. Need to enhance this.
-        lineno = frame.f_back.f_lineno
-    return semantic.device_assert(semantic.to_tensor(cond, _builder), msg, file_name, func_name, lineno, _builder)
+    return semantic.device_assert(semantic.to_tensor(cond, _builder), msg, _builder)
 
 
 @builtin

python/triton/language/semantic.py

@@ -1631,12 +1631,12 @@ def device_print(prefix: str, args: List[tl.tensor], hex: bool, builder: ir.buil
     return tl.tensor(builder.create_print(prefix, hex, new_args, is_signed), tl.void)
 
 
-def device_assert(cond: tl.tensor, msg: str, file_name: str, func_name, lineno: int, builder: ir.builder) -> tl.tensor:
+def device_assert(cond: tl.tensor, msg: str, builder: ir.builder) -> tl.tensor:
     cond_ty = cond.type
     if not cond_ty.is_block():
         cond_ty = tl.block_type(cond_ty.scalar, (1, ))
         cond = tl.tensor(builder.create_splat(cond.handle, (1, )), cond_ty)
-    return tl.tensor(builder.create_assert(cond.handle, msg, file_name, func_name, lineno), tl.void)
+    return tl.tensor(builder.create_assert(cond.handle, msg), tl.void)
 
 
 def assume(cond, builder: ir.builder) -> tl.tensor:

python/triton/runtime/interpreter.py

@@ -643,9 +643,9 @@ def create_print(self, prefix, hex, values, isSigned):
         if hex:
             np.set_printoptions(formatter=None)
 
-    def create_assert(self, condition, message, fileName, funcName, lineNo):
+    def create_assert(self, condition, message):
         # Interpreter's device_assert function has a different format than Triton's device_assert
-        assert condition, f"{message} in {fileName}:{funcName}:{lineNo}"
+        assert condition, f"{message}"
 
     def create_assume(self, condition):
         assert condition, "Assume failed"

test/TritonGPU/combine.mlir

@@ -2307,7 +2307,7 @@ module attributes {"triton_gpu.num-warps" = 4 : i32, "triton_gpu.num-ctas" = 1 :
 tt.func @assertop(%ptr: tensor<1024x!tt.ptr<i1>, #blocked>) {
   %0 = tt.load %ptr : tensor<1024x!tt.ptr<i1>, #blocked>
   %1 = triton_gpu.convert_layout %0 : tensor<1024xi1, #blocked> -> tensor<1024xi1, #blocked1>
-  tt.assert %1, "cond must be true ", "unknown", "unknown", 0 : tensor<1024xi1, #blocked1>
+  tt.assert %1, "cond must be true " : tensor<1024xi1, #blocked1>
   tt.return
 }
 }