Skip to content

Commit

Permalink
[Mosaic] Extend macros to handle parentheses.
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 706937151
  • Loading branch information
WindQAQ authored and Google-ML-Automation committed Dec 19, 2024
1 parent 16712b5 commit 76b0711
Showing 1 changed file with 31 additions and 7 deletions.
38 changes: 31 additions & 7 deletions jaxlib/mosaic/dialect/tpu/util.h
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,11 @@

#include <array>
#include <cstdint>
#include <optional>
#include <sstream>
#include <string>
#include <type_traits>
#include <utility>

#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/IR/Builders.h"
Expand All @@ -15,8 +17,9 @@
#include "mlir/Support/LLVM.h"
#include "mlir/Support/LogicalResult.h"
#include "absl/types/span.h"
#include "jaxlib/mosaic/dialect/tpu/tpu_dialect.h"
#include "mlir/include/mlir/IR/Value.h"
#include "jaxlib/mosaic/dialect/tpu/tpu_dialect.h"
#include "tsl/platform/statusor.h"

// TODO: Instead of CHECK_EQs, can we do something like TF_RET_CHECK but with
// MLIR diagnostics?
Expand All @@ -31,16 +34,37 @@
// } \
// } while (false)

#define FAILUREOR_ASSIGN_OR_RETURN_IMPL(failureor, lhs, rhs) \
auto failureor = rhs; \
if (failed(failureor)) { \
return failure(); \
} \
lhs = std::move(failureor).value();
// All the macros below here are to handle the case in
// FAILUREOR_ASSIGN_OR_RETURN where the LHS is wrapped in parentheses. See a
// more detailed discussion at https://stackoverflow.com/a/62984543
#define FAILUREOR_ASSIGN_OR_RETURN_UNPARENTHESIZE_IF_PARENTHESIZED(X) \
FAILUREOR_ASSIGN_OR_RETURN_ESCAPE(FAILUREOR_ASSIGN_OR_RETURN_EMPTY X)
#define FAILUREOR_ASSIGN_OR_RETURN_EMPTY(...) \
FAILUREOR_ASSIGN_OR_RETURN_EMPTY __VA_ARGS__
#define FAILUREOR_ASSIGN_OR_RETURN_ESCAPE(...) \
FAILUREOR_ASSIGN_OR_RETURN_ESCAPE_(__VA_ARGS__)
#define FAILUREOR_ASSIGN_OR_RETURN_ESCAPE_(...) \
FAILUREOR_ASSIGN_OR_RETURN_##__VA_ARGS__
#define FAILUREOR_ASSIGN_OR_RETURN_FAILUREOR_ASSIGN_OR_RETURN_EMPTY

#define FAILUREOR_ASSIGN_OR_RETURN_IMPL(failureor, lhs, rhs) \
auto failureor = rhs; \
if (failed(failureor)) { \
return failure(); \
} \
FAILUREOR_ASSIGN_OR_RETURN_UNPARENTHESIZE_IF_PARENTHESIZED(lhs) = \
(std::move(failureor).value());
#define FAILUREOR_ASSIGN_OR_RETURN(lhs, rhs) \
FAILUREOR_ASSIGN_OR_RETURN_IMPL( \
TF_STATUS_MACROS_CONCAT_NAME(failureor, __COUNTER__), lhs, rhs)

#define RETURN_IF_FAILED(...) \
do { \
if (failed(__VA_ARGS__)) { \
return failure(); \
} \
} while (false)

namespace mlir::tpu {

template <bool adjust_bool = false>
Expand Down

0 comments on commit 76b0711

Please sign in to comment.