Add macOS arm64 build compatibility fixes

build_macos.sh

@@ -0,0 +1,60 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "$0")" && pwd)"
+TMP_ROOT=/private/tmp
+
+export BAZELISK_HOME="$TMP_ROOT/tf-bazelisk"
+export CCACHE_DIR="$TMP_ROOT/tf-ccache"
+export CCACHE_BASEDIR="$ROOT_DIR"
+export CCACHE_NOHASHDIR=true
+export CCACHE_COMPILERCHECK=content
+export TF_PYTHON_VERSION=3.11
+export WHEEL_NAME=tensorflow_cpu
+export USE_PYWRAP_RULES=1
+export DEVELOPER_DIR=/Library/Developer/CommandLineTools
+
+# Pinned toolchain/build assumptions for this macOS build flow.
+BAZEL_BIN=/opt/homebrew/bin/bazel
+OUTPUT_USER_ROOT="$TMP_ROOT/tf-bazel-root"
+MACOS_MINIMUM_OS=12.0
+MACOS_SDK_VERSION=26.5
+JOBS=12
+TARGET=//tensorflow/tools/pip_package:wheel
+
+if command -v ccache >/dev/null 2>&1; then
+  CCACHE_PREFIX_DIR="$(dirname "$(command -v ccache)")/libexec"
+  if [ -d "$CCACHE_PREFIX_DIR" ]; then
+    export PATH="$CCACHE_PREFIX_DIR:$PATH"
+  fi
+fi
+
+if [ ! -x "$BAZEL_BIN" ]; then
+  echo "bazel not found at $BAZEL_BIN" >&2
+  exit 1
+fi
+
+"$BAZEL_BIN" --output_user_root="$OUTPUT_USER_ROOT" build \
+  --cpu=darwin_arm64 \
+  --macos_minimum_os="$MACOS_MINIMUM_OS" \
+  --macos_sdk_version="$MACOS_SDK_VERSION" \
+  --copt=-Wno-invalid-specialization \
+  --host_copt=-Wno-invalid-specialization \
+  --repo_env=TF_PYTHON_VERSION="$TF_PYTHON_VERSION" \
+  --repo_env=WHEEL_NAME="$WHEEL_NAME" \
+  --repo_env=USE_PYWRAP_RULES="$USE_PYWRAP_RULES" \
+  --action_env=CCACHE_DIR \
+  --host_action_env=CCACHE_DIR \
+  --action_env=CCACHE_BASEDIR \
+  --host_action_env=CCACHE_BASEDIR \
+  --action_env=CCACHE_NOHASHDIR \
+  --host_action_env=CCACHE_NOHASHDIR \
+  --action_env=CCACHE_COMPILERCHECK \
+  --host_action_env=CCACHE_COMPILERCHECK \
+  --action_env=DEVELOPER_DIR="$DEVELOPER_DIR" \
+  --host_action_env=DEVELOPER_DIR="$DEVELOPER_DIR" \
+  --strategy=CppCompile=local \
+  --jobs="$JOBS" \
+  --verbose_failures \
+  --show_progress_rate_limit=30 \
+  "$TARGET"

tensorflow/compiler/jit/BUILD

@@ -498,6 +498,14 @@ cc_library(
     ],
 )
 
+# Declares flags.h without linking the flags implementation or exposing its
+# transitive op registrations to wrapper generation.
+cc_library(
+    name = "flags_link_free_headers",
+    hdrs = ["flags.h"],
+    visibility = ["//tensorflow/compiler/tf2xla/ops:__pkg__"],
+)
+
 cc_header_only_library(
     name = "flags_headers_only",
     features = [

tensorflow/compiler/jit/encapsulate_subgraphs_pass.cc

@@ -139,6 +139,17 @@ std::string ExprProtoToString(const ExpressionProto& e) {
     case ExpressionProto::kDivNode:
       return absl::StrCat("(", ExprProtoToString(e.div_node().lhs()), " / ",
                          ExprProtoToString(e.div_node().rhs()), ")");
+    case ExpressionProto::kMaxNode:
+      return absl::StrCat("max(", ExprProtoToString(e.max_node().lhs()), ", ",
+                          ExprProtoToString(e.max_node().rhs()), ")");
+    case ExpressionProto::kGtNode:
+      return absl::StrCat("(", ExprProtoToString(e.gt_node().lhs()), " > ",
+                          ExprProtoToString(e.gt_node().rhs()), ")");
+    case ExpressionProto::kSelectNode:
+      return absl::StrCat("select(", ExprProtoToString(e.select_node().pred()),
+                          ", ", ExprProtoToString(e.select_node().on_true()),
+                          ", ", ExprProtoToString(e.select_node().on_false()),
+                          ")");
     default:
       return "<none>";
   }

tensorflow/compiler/jit/kernels/BUILD

@@ -1,3 +1,4 @@
+load("//tensorflow:tensorflow.bzl", "tf_cc_test")
 load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
 
 package(
@@ -48,7 +49,10 @@ XLA_OPS_DEPS = [
 cc_library(
     name = "xla_ops_no_jit_rewrite_registration",
     srcs = ["xla_ops.cc"],
-    hdrs = ["xla_ops.h"],
+    hdrs = [
+        "xla_ops.h",
+        "xla_ops_internal.h",
+    ],
     deps = XLA_OPS_DEPS + [
         "//tensorflow/compiler/jit:device_compilation_cache",
         "//tensorflow/compiler/jit:device_compilation_profiler",
@@ -75,6 +79,20 @@ cc_library(
     alwayslink = 1,
 )
 
+tf_cc_test(
+    name = "xla_ops_test",
+    srcs = ["xla_ops_test.cc"],
+    env = {"TF_XLA_FLAGS": "--tf_xla_enable_dynamic_sizes"},
+    deps = [
+        ":xla_ops_no_jit_rewrite_registration",
+        "//tensorflow/compiler/tf2xla:xla_argument",
+        "//tensorflow/core:framework",
+        "//tensorflow/core/platform:test",
+        "@com_google_googletest//:gtest_main",
+        "@local_xla//xla:shape_util",
+    ],
+)
+
 cc_library(
     name = "xla_ops",
     hdrs = ["xla_ops.h"],