mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 21:06:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Mosaic] Add a macro to convert abseil StatusOr to LLVM FailureOr

Browse Source 
PiperOrigin-RevId: 715943314
This commit is contained in:
Naums Mogers 2025-01-15 14:18:51 -08:00 committed by jax authors
parent 8a053af1ce
commit d3ba1eb339
1 changed files with 28 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
jaxlib/mosaic/dialect/tpu/util.h
View File

@ -17,8 +17,10 @@
#include "mlir/IR/Location.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Support/LogicalResult.h"
#include "absl/status/status.h"
#include "absl/types/span.h"
#include "mlir/include/mlir/IR/Attributes.h"
#include "mlir/include/mlir/IR/Diagnostics.h"
#include "mlir/include/mlir/IR/Value.h"
#include "jaxlib/mosaic/dialect/tpu/layout.h"
#include "jaxlib/mosaic/dialect/tpu/tpu_dialect.h"
@ -68,6 +70,32 @@
} \
} while (false)
template <typename Op>
class StatusToDiagnosticAdapter {
public:
// Returns an adapter that converts a non-OK absl::Status to an
// mlir::InFlightDiagnostic.
explicit StatusToDiagnosticAdapter(Op op) : op_(op) {}
// Converts a non-OK absl::Status to an mlir::InFlightDiagnostic.
mlir::InFlightDiagnostic operator()(const absl::Status &status) const {
return op_->emitOpError(status.ToString());
}
private:
Op op_;
};
// Returns a callable adapter that converts a non-OK absl::Status to an
// mlir::InFlightDiagnostic.
//
// Example usage:
// ASSIGN_OR_RETURN(T result, DoSomething(), _.With(StatusToDiagnostic(&op)));
template <typename Op>
inline StatusToDiagnosticAdapter<Op> StatusToDiagnostic(Op op) {
return StatusToDiagnosticAdapter<Op>(op);
}
namespace mlir::tpu {
// TPU_ASSERT_* macros should be understood as an assert, i.e. use it to check