tensorflow-metal-plugin 0.1.2__tar.gz → 0.2.0__tar.gz

{tensorflow_metal_plugin-0.1.2 → tensorflow_metal_plugin-0.2.0}/Makefile

@@ -23,6 +23,19 @@ SDK := $(shell xcrun --sdk macosx --show-sdk-path 2>/dev/null)
 BUILD := build
 OUT := $(BUILD)/libmetal_plugin.dylib
 
+# Which TensorFlow the objects in $(BUILD) were compiled against.
+#
+# The header dependencies recorded by -MMD name paths inside that TensorFlow,
+# and those paths keep existing when PYTHON is pointed at a different one, so
+# make sees nothing to redo and links objects built against two different
+# TensorFlows. The result is a library that fails to load with an absl symbol
+# missing, which says nothing about the cause. This makes the include path
+# itself a dependency.
+STAMP := $(BUILD)/tf-include-path
+$(shell mkdir -p $(BUILD); \
+        [ "$$(cat $(STAMP) 2>/dev/null)" = "$(TF_INCLUDE)" ] || \
+        printf '%s' "$(TF_INCLUDE)" > $(STAMP))
+
 SOURCES := src/plugin_init.cc \
            $(wildcard src/tensorflow/core/common_runtime/metal/*.mm) \
            $(wildcard src/tensorflow/core/common_runtime/metal/kernels/*.mm)
@@ -69,7 +82,7 @@ LDFLAGS := -dynamiclib $(FRAMEWORKS) \
            -Wl,-undefined,dynamic_lookup \
            -Wl,-rpath,$(TF_LIB)
 
-.PHONY: all clean test sweep install
+.PHONY: all clean test test-stream sweep install
 
 all: $(OUT)
 
@@ -78,7 +91,7 @@ $(OUT): $(OBJECTS)
 	$(CXX) $(OBJECTS) $(LDFLAGS) -o $@
 	@echo "built $@"
 
-$(BUILD)/%.o: src/%
+$(BUILD)/%.o: src/% $(STAMP)
 	@mkdir -p $(dir $@)
 	$(CXX) $(CXXFLAGS) -c $< -o $@
 
@@ -86,9 +99,21 @@ $(BUILD)/%.o: src/%
 check-symbols: $(OUT)
 	@bash tools/check_symbols.sh $(OUT)
 
-test: $(OUT)
+test: $(OUT) test-stream
 	$(PYTHON) tests/run_tests.py
 
+# The StreamExecutor C API driven directly, without TensorFlow. Reaches what
+# no op can: memset32 with a pattern that is not four equal bytes has one
+# caller in the whole tree and it is a CUDA-only kernel.
+$(BUILD)/stream_executor_test: tests/stream_executor_test.mm $(OUT)
+	@mkdir -p $(dir $@)
+	$(CXX) $(CXXFLAGS) tests/stream_executor_test.mm -o $@ \
+	  -framework Foundation -L$(TF_LIB) -ltensorflow_framework.2 \
+	  -Wl,-rpath,$(TF_LIB)
+
+test-stream: $(BUILD)/stream_executor_test
+	$(BUILD)/stream_executor_test $(OUT)
+
 # Every registered op, through TensorFlow's own dispatch, against the CPU.
 sweep: $(OUT)
 	PYTHONPATH=tools $(PYTHON) tools/op_sweep.py

{tensorflow_metal_plugin-0.1.2/tensorflow_metal_plugin.egg-info → tensorflow_metal_plugin-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tensorflow-metal-plugin
-Version: 0.1.2
+Version: 0.2.0
 Summary: Metal GPU backend for TensorFlow on Apple silicon
 License: Apache-2.0
 Classifier: Development Status :: 3 - Alpha
@@ -28,10 +28,15 @@ PluggableDevice. It loads into a stock TensorFlow wheel and adds
 `/physical_device:GPU:0`.
 
 This is the out-of-tree form of the backend proposed in
-[tensorflow/tensorflow#126254](https://github.com/tensorflow/tensorflow/pull/126254).
+[tensorflow/tensorflow#126384](https://github.com/tensorflow/tensorflow/pull/126384).
 The sources are the same; the only difference is this repository exports
-`SE_InitPlugin` and `TF_InitKernel` from a shared object, where the in-tree
-form hands the same function pointers to `RegisterPluggableDevicePlugin`.
+`SE_InitPlugin`, `TF_InitKernel`, `TF_InitProfiler` and `TF_InitGraph` from a
+shared object, where the in-tree form hands the same function pointers to
+`RegisterPluggableDevicePlugin`.
+
+Four modules, then: the device and its memory, the kernels, a profiler that
+puts Metal work on the TensorFlow timeline, and a graph pass that fuses a bias
+and an activation into the convolution or matrix multiply in front of them.
 
 ## Status
 
@@ -126,8 +131,15 @@ error raised. `model.fit` reported `[nan, nan, nan]` and carried on.
 
 While those entry points are missing, every Metal kernel waits for the GPU
 before returning, which closes the window. It is announced in a warning at
-load, and `TF_METAL_SYNCHRONOUS` forces it either way. When the entry points
-come back the plugin returns to running asynchronously with no change here.
+load, and `TF_METAL_SYNCHRONOUS` forces it either way.
+
+Where the entry points do exist, the plugin implements `AssignVariableOp`,
+`AssignAddVariableOp` and `AssignSubVariableOp` itself, on the device, and
+nothing touches a variable from the host at all. That is what makes running
+asynchronously safe rather than merely faster: without those kernels, an
+asynchronous run reproduced the same `nan` on TensorFlow 2.19.1, where the C
+API is present. With them, a training step is 11.48 ms against 21.04 ms on the
+CPU, and correct.
 
 ## Is it faster than the CPU
 
@@ -136,21 +148,25 @@ on an M4 Max against TensorFlow 2.21.0, median of ten runs each, both devices
 in the same process on the same data, waiting for the device before stopping
 the clock:
 
-Two columns of speedup, because the wait above costs most of it. "Today" is
-what you get from a released TensorFlow; "async" is the same machine with
-`TF_METAL_SYNCHRONOUS=0`, which is what the plugin does once
+Two columns, because the wait above costs most of it. "Today" is a released
+TensorFlow, 2.21.0, where the kernel C API for resource variables is missing
+and every kernel therefore waits. "Async" is 2.19.1, which still exports that
+API, so the same code runs asynchronously; it is what a released TensorFlow
+does again once
 [#126374](https://github.com/tensorflow/tensorflow/issues/126374) is fixed.
+Both columns are measured, not projected.
 
 | | GPU today | CPU | today | async |
 | --- | ---: | ---: | ---: | ---: |
-| MatMul 2048x2048 | 3.80 ms | 14.42 ms | **3.8x** | 6.5x |
-| Conv2D, batch 64, 64x64x32 to 64 | 4.58 ms | 10.91 ms | **2.4x** | 3.0x |
-| MatMul 1024x1024 | 1.12 ms | 2.17 ms | 1.9x | 1.5x |
-| CNN training step, SGD, batch 128 | 19.46 ms | 20.61 ms | 1.1x | 1.7x |
-| CNN forward, batch 128 | 5.53 ms | 6.00 ms | 1.1x | 1.5x |
-| MatMul 512x512 | 0.44 ms | 0.36 ms | 0.8x | 0.9x |
-| ReduceSum 4096x4096 | 0.44 ms | 0.34 ms | 0.8x | 0.8x |
-| Elementwise 4096x4096 | 2.88 ms | 1.54 ms | 0.5x | 0.5x |
+| MatMul 2048x2048 | 1.78 ms | 11.97 ms | **6.7x** | 5.5x |
+| Conv2D, batch 64, 64x64x32 to 64 | 2.25 ms | 8.95 ms | **4.0x** | 4.2x |
+| CNN training step, SGD, batch 128 | 12.60 ms | 18.80 ms | 1.5x | 1.8x |
+| MatMul 1024x1024 | 1.19 ms | 1.89 ms | 1.6x | 2.0x |
+| Conv2D, batch 16 | 1.37 ms | 2.31 ms | 1.7x | 1.7x |
+| CNN forward, batch 128 | 4.55 ms | 5.17 ms | 1.1x | 1.5x |
+| MatMul 512x512 | 0.37 ms | 0.34 ms | 0.9x | 1.0x |
+| ReduceSum 4096x4096 | 0.46 ms | 0.28 ms | 0.6x | 0.7x |
+| Elementwise 4096x4096 | 3.05 ms | 1.36 ms | 0.5x | 0.5x |
 
 The pattern is the ordinary one and worth stating plainly: the GPU wins where
 there is arithmetic to do per byte moved, and loses where there is not. A
@@ -159,9 +175,12 @@ operations per element, so it is bound by memory on a machine whose CPU shares
 that same memory. Small matrices lose to the cost of getting work to the
 device at all.
 
-A training step is 1.7x the CPU when the plugin can run asynchronously, and
-barely ahead of it while it cannot. That gap is the cost of the missing entry
-points, not of the backend.
+The convolution numbers owe as much to the graph pass as to the kernels. It
+folds the bias and the activation into the convolution, and it turns
+TensorFlow's layout optimizer off, which was inserting an NHWC to NCHW
+transpose on either side of every convolution: on a 4x16x16x8 case those cost
+42.8 and 44.2 microseconds around a 43.0 microsecond convolution. MPSGraph
+takes either layout, so the rewrite was pure loss.
 
 `benchmarks/benchmark.py` reproduces the table.
 
@@ -264,7 +283,7 @@ Metal kernels with their dtypes is in
 ## Layout
 
 ```
-src/plugin_init.cc                          the two exported entry points
+src/plugin_init.cc                          the exported entry points
 src/tensorflow/core/common_runtime/metal/   the backend, verbatim from the
                                             TensorFlow tree
 tools/                                      build probes and the symbol check

{tensorflow_metal_plugin-0.1.2 → tensorflow_metal_plugin-0.2.0}/README.md

@@ -5,10 +5,15 @@ PluggableDevice. It loads into a stock TensorFlow wheel and adds
 `/physical_device:GPU:0`.
 
 This is the out-of-tree form of the backend proposed in
-[tensorflow/tensorflow#126254](https://github.com/tensorflow/tensorflow/pull/126254).
+[tensorflow/tensorflow#126384](https://github.com/tensorflow/tensorflow/pull/126384).
 The sources are the same; the only difference is this repository exports
-`SE_InitPlugin` and `TF_InitKernel` from a shared object, where the in-tree
-form hands the same function pointers to `RegisterPluggableDevicePlugin`.
+`SE_InitPlugin`, `TF_InitKernel`, `TF_InitProfiler` and `TF_InitGraph` from a
+shared object, where the in-tree form hands the same function pointers to
+`RegisterPluggableDevicePlugin`.
+
+Four modules, then: the device and its memory, the kernels, a profiler that
+puts Metal work on the TensorFlow timeline, and a graph pass that fuses a bias
+and an activation into the convolution or matrix multiply in front of them.
 
 ## Status
 
@@ -103,8 +108,15 @@ error raised. `model.fit` reported `[nan, nan, nan]` and carried on.
 
 While those entry points are missing, every Metal kernel waits for the GPU
 before returning, which closes the window. It is announced in a warning at
-load, and `TF_METAL_SYNCHRONOUS` forces it either way. When the entry points
-come back the plugin returns to running asynchronously with no change here.
+load, and `TF_METAL_SYNCHRONOUS` forces it either way.
+
+Where the entry points do exist, the plugin implements `AssignVariableOp`,
+`AssignAddVariableOp` and `AssignSubVariableOp` itself, on the device, and
+nothing touches a variable from the host at all. That is what makes running
+asynchronously safe rather than merely faster: without those kernels, an
+asynchronous run reproduced the same `nan` on TensorFlow 2.19.1, where the C
+API is present. With them, a training step is 11.48 ms against 21.04 ms on the
+CPU, and correct.
 
 ## Is it faster than the CPU
 
@@ -113,21 +125,25 @@ on an M4 Max against TensorFlow 2.21.0, median of ten runs each, both devices
 in the same process on the same data, waiting for the device before stopping
 the clock:
 
-Two columns of speedup, because the wait above costs most of it. "Today" is
-what you get from a released TensorFlow; "async" is the same machine with
-`TF_METAL_SYNCHRONOUS=0`, which is what the plugin does once
+Two columns, because the wait above costs most of it. "Today" is a released
+TensorFlow, 2.21.0, where the kernel C API for resource variables is missing
+and every kernel therefore waits. "Async" is 2.19.1, which still exports that
+API, so the same code runs asynchronously; it is what a released TensorFlow
+does again once
 [#126374](https://github.com/tensorflow/tensorflow/issues/126374) is fixed.
+Both columns are measured, not projected.
 
 | | GPU today | CPU | today | async |
 | --- | ---: | ---: | ---: | ---: |
-| MatMul 2048x2048 | 3.80 ms | 14.42 ms | **3.8x** | 6.5x |
-| Conv2D, batch 64, 64x64x32 to 64 | 4.58 ms | 10.91 ms | **2.4x** | 3.0x |
-| MatMul 1024x1024 | 1.12 ms | 2.17 ms | 1.9x | 1.5x |
-| CNN training step, SGD, batch 128 | 19.46 ms | 20.61 ms | 1.1x | 1.7x |
-| CNN forward, batch 128 | 5.53 ms | 6.00 ms | 1.1x | 1.5x |
-| MatMul 512x512 | 0.44 ms | 0.36 ms | 0.8x | 0.9x |
-| ReduceSum 4096x4096 | 0.44 ms | 0.34 ms | 0.8x | 0.8x |
-| Elementwise 4096x4096 | 2.88 ms | 1.54 ms | 0.5x | 0.5x |
+| MatMul 2048x2048 | 1.78 ms | 11.97 ms | **6.7x** | 5.5x |
+| Conv2D, batch 64, 64x64x32 to 64 | 2.25 ms | 8.95 ms | **4.0x** | 4.2x |
+| CNN training step, SGD, batch 128 | 12.60 ms | 18.80 ms | 1.5x | 1.8x |
+| MatMul 1024x1024 | 1.19 ms | 1.89 ms | 1.6x | 2.0x |
+| Conv2D, batch 16 | 1.37 ms | 2.31 ms | 1.7x | 1.7x |
+| CNN forward, batch 128 | 4.55 ms | 5.17 ms | 1.1x | 1.5x |
+| MatMul 512x512 | 0.37 ms | 0.34 ms | 0.9x | 1.0x |
+| ReduceSum 4096x4096 | 0.46 ms | 0.28 ms | 0.6x | 0.7x |
+| Elementwise 4096x4096 | 3.05 ms | 1.36 ms | 0.5x | 0.5x |
 
 The pattern is the ordinary one and worth stating plainly: the GPU wins where
 there is arithmetic to do per byte moved, and loses where there is not. A
@@ -136,9 +152,12 @@ operations per element, so it is bound by memory on a machine whose CPU shares
 that same memory. Small matrices lose to the cost of getting work to the
 device at all.
 
-A training step is 1.7x the CPU when the plugin can run asynchronously, and
-barely ahead of it while it cannot. That gap is the cost of the missing entry
-points, not of the backend.
+The convolution numbers owe as much to the graph pass as to the kernels. It
+folds the bias and the activation into the convolution, and it turns
+TensorFlow's layout optimizer off, which was inserting an NHWC to NCHW
+transpose on either side of every convolution: on a 4x16x16x8 case those cost
+42.8 and 44.2 microseconds around a 43.0 microsecond convolution. MPSGraph
+takes either layout, so the rewrite was pure loss.
 
 `benchmarks/benchmark.py` reproduces the table.
 
@@ -241,7 +260,7 @@ Metal kernels with their dtypes is in
 ## Layout
 
 ```
-src/plugin_init.cc                          the two exported entry points
+src/plugin_init.cc                          the exported entry points
 src/tensorflow/core/common_runtime/metal/   the backend, verbatim from the
                                             TensorFlow tree
 tools/                                      build probes and the symbol check

{tensorflow_metal_plugin-0.1.2 → tensorflow_metal_plugin-0.2.0}/setup.py

@@ -75,7 +75,7 @@ class BinaryDistribution(Distribution):
 
 setup(
     name="tensorflow-metal-plugin",
-    version="0.1.2",
+    version="0.2.0",
     description="Metal GPU backend for TensorFlow on Apple silicon",
     long_description=(HERE / "README.md").read_text(),
     long_description_content_type="text/markdown",

{tensorflow_metal_plugin-0.1.2 → tensorflow_metal_plugin-0.2.0}/src/plugin_init.cc

@@ -27,7 +27,9 @@ limitations under the License.
 #include "tensorflow/c/experimental/stream_executor/stream_executor.h"
 #include "tensorflow/c/tf_status.h"
 #include "tensorflow/core/common_runtime/metal/kernels/metal_kernels.h"
+#include "tensorflow/core/common_runtime/metal/metal_graph.h"
 #include "tensorflow/core/common_runtime/metal/metal_platform.h"
+#include "tensorflow/core/common_runtime/metal/metal_profiler.h"
 
 namespace {
 
@@ -55,4 +57,28 @@ void SE_InitPlugin(SE_PlatformRegistrationParams* params, TF_Status* status) {
 
 void TF_InitKernel() { tensorflow::metal::RegisterAllMetalKernels(); }
 
+// Optional: TensorFlow looks this up and carries on without it if absent. It
+// is what puts a Metal row in the trace viewer next to the host's.
+void TF_InitProfiler(TF_ProfilerRegistrationParams* params,
+                     TF_Status* status) {
+  if (MetalDisabledByEnvironment()) {
+    TF_SetStatus(status, TF_FAILED_PRECONDITION,
+                 "Metal: backend disabled by TF_DISABLE_METAL.");
+    return;
+  }
+  tensorflow::metal::MetalInitProfiler(params, status);
+}
+
+// Optional, like the profiler. Fuses a bias and an activation into the
+// convolution or matrix multiply in front of them, which this backend has
+// kernels for and which nothing else produces for a pluggable device.
+void TF_InitGraph(TP_OptimizerRegistrationParams* params, TF_Status* status) {
+  if (MetalDisabledByEnvironment()) {
+    TF_SetStatus(status, TF_FAILED_PRECONDITION,
+                 "Metal: backend disabled by TF_DISABLE_METAL.");
+    return;
+  }
+  tensorflow::metal::MetalInitGraph(params, status);
+}
+
 }  // extern "C"