Merge pull request #3142 from ScottTodd:main-to-google PiperOrigin-RevId: 331661857

diff --git a/build_tools/docker/build_and_update_gcr.py b/build_tools/docker/build_and_update_gcr.py
deleted file mode 100755
index 6de59a4..0000000
--- a/build_tools/docker/build_and_update_gcr.py
+++ /dev/null

@@ -1,158 +0,0 @@
-#!/usr/bin/env python3
-
-# Copyright 2020 Google LLC
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#      https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Builds the specified Docker images and optionally pushes them to GCR.
-
-Example usage:
-  python3 build_tools/docker/build_and_update_gcr.py --image cmake
-"""
-
-import argparse
-import functools
-import os
-import subprocess
-import sys
-
-IREE_GCR_URL = 'gcr.io/iree-oss/'
-DOCKER_DIR = 'build_tools/docker/'
-
-# Map from image names to images that they depend on.
-IMAGES_TO_DEPENDENCIES = {
-    'bazel': [],
-    'bazel-bindings': ['bazel'],
-    'bazel-tensorflow': ['bazel-bindings'],
-    'bazel-nvidia': ['bazel-tensorflow'],
-    'bazel-swiftshader': ['bazel-tensorflow'],
-    'cmake': [],
-    'cmake-android': ['cmake'],
-    'cmake-nvidia': ['cmake'],
-    'cmake-vulkan': ['cmake'],
-    'cmake-swiftshader': ['cmake-vulkan'],
-    'rbe-toolchain': [],
-}
-
-IMAGES_TO_DEPENDENT_IMAGES = {k: [] for k in IMAGES_TO_DEPENDENCIES.keys()}
-for image, dependencies in IMAGES_TO_DEPENDENCIES.items():
-  for dependency in dependencies:
-    IMAGES_TO_DEPENDENT_IMAGES[dependency].append(image)
-
-IMAGES_HELP = [f'`{name}`' for name in IMAGES_TO_DEPENDENCIES.keys()]
-IMAGES_HELP = f'{", ".join(IMAGES_HELP)} or `all`'
-
-RBE_MESSAGE = """
-Remember to update the `rbe_default` digest in the `WORKSPACE` file to reflect
-the new digest for the container.
-
-Use `docker images --digests` to view the digest."""
-
-
-def parse_arguments():
-  """Parses command-line options."""
-  parser = argparse.ArgumentParser(
-      description="Build IREE's Docker images and optionally push them to GCR.")
-  parser.add_argument(
-      '--image',
-      dest='images',
-      type=str,
-      required=True,
-      action='append',
-      help=f'Name of the image to build: {IMAGES_HELP}.')
-  parser.add_argument(
-      '--tag',
-      type=str,
-      default='latest',
-      help='Tags for the images to build. Defaults to `latest` (which is good '
-      'for testing changes in a PR). Use `prod` to update the images that the '
-      'OSS CI uses.')
-  parser.add_argument(
-      '--push',
-      action='store_true',
-      help='Push the built images to GCR. Requires gcloud authorization.')
-
-  args = parser.parse_args()
-  for image in args.images:
-    if image == 'all':
-      args.images = IMAGES_TO_DEPENDENCIES.keys()
-    elif image not in IMAGES_TO_DEPENDENCIES.keys():
-      raise parser.error('Expected --image to be one of:\n'
-                         f'  {IMAGES_HELP}\n'
-                         f'but got `{image}`.')
-  return args
-
-
-def cmp_images_by_dependency(image1, image2):
-  if image2 in IMAGES_TO_DEPENDENT_IMAGES[image1]:
-    return -1
-  if image1 in IMAGES_TO_DEPENDENT_IMAGES[image2]:
-    return 1
-  return (image1 > image2) - (image1 < image2)
-
-
-def run_command(command):
-  print(f'Running: {" ".join(command)}')
-  process = subprocess.Popen(
-      command,
-      bufsize=1,
-      stderr=subprocess.STDOUT,
-      stdout=subprocess.PIPE,
-      text=True)
-  for line in process.stdout:
-    print(line, end='')
-
-  return process.poll()
-
-
-def check_command(command):
-  exit_code = run_command(command)
-  if exit_code != 0:
-    print(f'Command failed: {" ".join(command)}')
-    sys.exit(exit_code)
-
-
-if __name__ == '__main__':
-  args = parse_arguments()
-
-  # Ensure the user has the correct authorization if they try to push to GCR.
-  if args.push:
-    if run_command(['which', 'gcloud']) != 0:
-      print('gcloud not found.'
-            ' See https://cloud.google.com/sdk/install for installation.')
-      sys.exit(1)
-    check_command(['gcloud', 'auth', 'configure-docker'])
-
-  # Check if any images depend on `args.images` and update them if they do.
-  images_to_update_set = set()
-  to_check = list(args.images)
-  while to_check:
-    image = to_check.pop()
-    if image not in images_to_update_set:
-      images_to_update_set.add(image)
-      to_check.extend(IMAGES_TO_DEPENDENT_IMAGES[image])
-
-  # Topo sort by image dependency
-  images_to_update = sorted(
-      images_to_update_set, key=functools.cmp_to_key(cmp_images_by_dependency))
-
-  print(f'Also updating dependent images. Will update: {images_to_update}')
-  for image in images_to_update:
-    print(f'Updating image {image}')
-    image_url = os.path.join(IREE_GCR_URL, f'{image}:{args.tag}')
-    image_path = os.path.join(DOCKER_DIR, image.replace('-', '_'))
-    check_command(['docker', 'build', '--tag', image_url, image_path])
-    if args.push:
-      check_command(['docker', 'push', image_url])
-
-  if 'rbe-toolchain' in images_to_update:
-    print(RBE_MESSAGE)

diff --git a/build_tools/docker/manage_images.py b/build_tools/docker/manage_images.py
new file mode 100755
index 0000000..9e7b7e3
--- /dev/null
+++ b/build_tools/docker/manage_images.py

@@ -0,0 +1,272 @@
+#!/usr/bin/env python3
+
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Manages IREE Docker image definitions.
+
+Includes information on their dependency graph and GCR URL.
+
+Example usage:
+
+Rebuild the cmake image and all images that transitiviely on depend on it,
+tagging them with `latest`:
+  python3 build_tools/docker/manage_images.py --build --image cmake
+
+Print out output for rebuilding the cmake image and all images that
+transitiviely on depend on it, but don't take side-effecting actions:
+  python3 build_tools/docker/manage_images.py --build --image cmake --dry-run
+
+Push all `prod` images to GCR:
+  python3 build_tools/docker/manage_images.py --push --tag prod --images all
+
+Rebuild and push all images and update references to them in the repository:
+  python3 build_tools/docker/manage_images.py --push --images all
+  --update-references
+"""
+
+import argparse
+import fileinput
+import os
+import posixpath
+import re
+import subprocess
+import sys
+
+IREE_GCR_URL = 'gcr.io/iree-oss/'
+DOCKER_DIR = 'build_tools/docker/'
+
+# Map from image names to images that they depend on.
+IMAGES_TO_DEPENDENCIES = {
+    'bazel': [],
+    'bazel-bindings': ['bazel'],
+    'bazel-tensorflow': ['bazel-bindings'],
+    'bazel-nvidia': ['bazel-tensorflow'],
+    'bazel-swiftshader': ['bazel-tensorflow'],
+    'cmake': [],
+    'cmake-android': ['cmake'],
+    'cmake-nvidia': ['cmake'],
+    'cmake-vulkan': ['cmake'],
+    'cmake-swiftshader': ['cmake-vulkan'],
+    'rbe-toolchain': [],
+}
+
+IMAGES_TO_DEPENDENT_IMAGES = {k: [] for k in IMAGES_TO_DEPENDENCIES}
+for image, dependencies in IMAGES_TO_DEPENDENCIES.items():
+  for dependency in dependencies:
+    IMAGES_TO_DEPENDENT_IMAGES[dependency].append(image)
+
+IMAGES_HELP = [f'`{name}`' for name in IMAGES_TO_DEPENDENCIES]
+IMAGES_HELP = f'{", ".join(IMAGES_HELP)} or `all`'
+
+
+def parse_arguments():
+  """Parses command-line options."""
+  parser = argparse.ArgumentParser(
+      description="Build IREE's Docker images and optionally push them to GCR.")
+  parser.add_argument(
+      '--images',
+      '--image',
+      type=str,
+      required=True,
+      action='append',
+      help=f'Name of the image to build: {IMAGES_HELP}.')
+  parser.add_argument(
+      '--tag',
+      type=str,
+      default='latest',
+      help='Tag for the images to build. Defaults to `latest` (which is good '
+      'for testing changes in a PR). Use `prod` to update the images that the '
+      'CI caches.')
+  parser.add_argument(
+      '--pull',
+      action='store_true',
+      help='Pull the specified image before building.')
+  parser.add_argument(
+      '--build',
+      action='store_true',
+      help='Build new images from the current Dockerfiles.')
+  parser.add_argument(
+      '--push',
+      action='store_true',
+      help='Push the built images to GCR. Requires gcloud authorization.')
+  parser.add_argument(
+      '--update_references',
+      '--update-references',
+      action='store_true',
+      help='Update all references to the specified images to point at the new'
+      ' digest.')
+  parser.add_argument(
+      '--dry_run',
+      '--dry-run',
+      '-n',
+      action='store_true',
+      help='Print output without building or pushing any images.')
+
+  args = parser.parse_args()
+  for image in args.images:
+    if image == 'all':
+      # Sort for a determinstic order
+      args.images = sorted(IMAGES_TO_DEPENDENCIES.keys())
+    elif image not in IMAGES_TO_DEPENDENCIES:
+      raise parser.error('Expected --image to be one of:\n'
+                         f'  {IMAGES_HELP}\n'
+                         f'but got `{image}`.')
+  return args
+
+
+def get_ordered_images_to_process(images):
+  unmarked_images = list(images)
+  # Python doesn't have a builtin OrderedSet
+  marked_images = set()
+  order = []
+
+  def visit(image):
+    if image in marked_images:
+      return
+    for dependent_images in IMAGES_TO_DEPENDENT_IMAGES[image]:
+      visit(dependent_images)
+    marked_images.add(image)
+    order.append(image)
+
+  while unmarked_images:
+    visit(unmarked_images.pop())
+
+  order.reverse()
+  return order
+
+
+def stream_command(command, dry_run=False):
+  print(f'Running: `{" ".join(command)}`')
+  if dry_run:
+    return 0
+  process = subprocess.Popen(
+      command,
+      bufsize=1,
+      stderr=subprocess.STDOUT,
+      stdout=subprocess.PIPE,
+      universal_newlines=True)
+  for line in process.stdout:
+    print(line, end='')
+
+  if process.poll() is None:
+    raise RuntimeError('Unexpected end of output while process is not finished')
+  return process.poll()
+
+
+def check_stream_command(command, dry_run=False):
+  exit_code = stream_command(command, dry_run=dry_run)
+  if exit_code != 0:
+    print(f'Command failed with exit code {exit_code}: `{" ".join(command)}`')
+    sys.exit(exit_code)
+
+
+def get_repo_digest(image):
+  inspect_command = [
+      'docker',
+      'image',
+      'inspect',
+      f'{image}',
+      '-f',
+      '{{index .RepoDigests 0}}',
+  ]
+  inspect_process = subprocess.run(
+      inspect_command,
+      universal_newlines=True,
+      stdout=subprocess.PIPE,
+      stderr=subprocess.PIPE,
+      timeout=10)
+  if inspect_process.returncode != 0:
+    print(f'Computing the repository digest for {image} failed.'
+          ' Has it been pushed to GCR?')
+    print(f'Output from `{" ".join(inspect_command)}`:')
+    print(inspect_process.stdout, end='')
+    print(inspect_process.stderr, end='')
+    sys.exit(inspect_process.returncode)
+  _, repo_digest = inspect_process.stdout.strip().split('@')
+  return repo_digest
+
+
+def update_rbe_reference(digest, dry_run=False):
+  print('Updating WORKSPACE file for rbe-toolchain')
+  for line in fileinput.input(files=['WORKSPACE'], inplace=(not dry_run)):
+    if line.strip().startswith('digest ='):
+      print(re.sub('sha256:[a-zA-Z0-9]+', digest, line), end='')
+    else:
+      print(line, end='')
+
+
+def update_references(image_name, digest, dry_run=False):
+  print(f'Updating references to {image_name}')
+
+  grep_command = ['git', 'grep', '-l', f'{image_name}@sha256']
+  grep_process = subprocess.run(
+      grep_command,
+      stdout=subprocess.PIPE,
+      stderr=subprocess.PIPE,
+      timeout=5,
+      universal_newlines=True)
+  if grep_process.returncode > 1:
+    print(f'{" ".join(grep_command)} '
+          f'failed with exit code {grep_process.returncode}')
+    sys.exit(grep_process.returncode)
+  if grep_process.returncode == 1:
+    print(f'Found no references to {image_name}')
+    return
+
+  files = grep_process.stdout.split()
+  print(f'Updating references in {len(files)} files: {files}')
+  for line in fileinput.input(files=files, inplace=(not dry_run)):
+    print(
+        re.sub(f'{image_name}@sha256:[a-zA-Z0-9]+', f'{image_name}@{digest}',
+               line),
+        end='')
+
+
+if __name__ == '__main__':
+  args = parse_arguments()
+
+  # Ensure the user has the correct authorization if they try to push to GCR.
+  if args.push:
+    if stream_command(['which', 'gcloud']) != 0:
+      print('gcloud not found.'
+            ' See https://cloud.google.com/sdk/install for installation.')
+      sys.exit(1)
+    check_stream_command(['gcloud', 'auth', 'configure-docker'],
+                         dry_run=args.dry_run)
+
+  images_to_process = get_ordered_images_to_process(args.images)
+  print(f'Also processing dependent images. Will process: {images_to_process}')
+
+  for image in images_to_process:
+    print(f'Processing image {image}')
+    image_name = posixpath.join(IREE_GCR_URL, image)
+    image_tag = f'{image_name}:{args.tag}'
+    image_path = os.path.join(DOCKER_DIR, image.replace('-', '_'))
+
+    if args.pull:
+      check_stream_command(['docker', 'pull', image_tag], dry_run=args.dry_run)
+
+    if args.build:
+      check_stream_command(['docker', 'build', '--tag', image_tag, image_path],
+                           dry_run=args.dry_run)
+
+    if args.push:
+      check_stream_command(['docker', 'push', image_tag], dry_run=args.dry_run)
+
+    if args.update_references:
+      digest = get_repo_digest(image_tag)
+      # Just hardcode this oddity
+      if image == 'rbe-toolchain':
+        update_rbe_reference(digest, dry_run=args.dry_run)
+      update_references(image_name, digest, dry_run=args.dry_run)

diff --git a/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/integrations/build_kokoro.sh b/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/integrations/build_kokoro.sh
index 63c5893..b8372a2 100755
--- a/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/integrations/build_kokoro.sh
+++ b/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/integrations/build_kokoro.sh

@@ -33,7 +33,7 @@
 
 docker run "${DOCKER_RUN_ARGS[@]?}" \
   --env IREE_VULKAN_DISABLE=0 \
-  gcr.io/iree-oss/bazel-swiftshader:prod \
+  gcr.io/iree-oss/bazel-swiftshader@sha256:59ca639199c1548d3fd2c9f6bcc04ccd91c6a32fd514a4ea4e2b8542e7b0eed2 \
   build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-swiftshader/integrations/build.sh
 
 # Kokoro will rsync this entire directory back to the executor orchestrating the

diff --git a/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-turing/integrations/build_kokoro.sh b/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-turing/integrations/build_kokoro.sh
index 0712a66..4f6f453 100755
--- a/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-turing/integrations/build_kokoro.sh
+++ b/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-turing/integrations/build_kokoro.sh

@@ -34,7 +34,7 @@
   --rm \
   --env IREE_VULKAN_DISABLE=0 \
   --gpus all \
-  gcr.io/iree-oss/bazel-nvidia:prod \
+  gcr.io/iree-oss/bazel-nvidia@sha256:77866668ac679de65f5008c3c3df0e5dbf2944a431c88a8c1b6b2e8ab7f8c65d \
   build_tools/kokoro/gcp_ubuntu/bazel/linux/x86-turing/integrations/build.sh
 
 # Kokoro will rsync this entire directory back to the executor orchestrating the

diff --git a/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86/bindings/build_kokoro.sh b/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86/bindings/build_kokoro.sh
index 24cbabf..ce85c9c 100755
--- a/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86/bindings/build_kokoro.sh
+++ b/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86/bindings/build_kokoro.sh

@@ -32,7 +32,7 @@
 docker_setup
 
 docker run "${DOCKER_RUN_ARGS[@]?}" \
-  gcr.io/iree-oss/bazel-bindings:prod \
+  gcr.io/iree-oss/bazel-bindings@sha256:1f5e59f10c35d0f9211c1a8821a931aca746f47a66fe1bc31b8b3bad4f38a0a7 \
   build_tools/kokoro/gcp_ubuntu/bazel/linux/x86/bindings/build.sh
 
 # Kokoro will rsync this entire directory back to the executor orchestrating the

diff --git a/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86/core/build_kokoro.sh b/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86/core/build_kokoro.sh
index da07cdd..da01f66 100755
--- a/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86/core/build_kokoro.sh
+++ b/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86/core/build_kokoro.sh

@@ -32,7 +32,7 @@
 docker_setup
 
 docker run "${DOCKER_RUN_ARGS[@]?}" \
-  gcr.io/iree-oss/bazel:prod \
+  gcr.io/iree-oss/bazel@sha256:b3ce4db78ccfc175cad0b071075cb4b26845980dacd7afa51f39625876062476 \
   build_tools/kokoro/gcp_ubuntu/bazel/linux/x86/core/build.sh
 
 # Kokoro will rsync this entire directory back to the executor orchestrating the

diff --git a/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86/integrations/build_kokoro.sh b/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86/integrations/build_kokoro.sh
index c35d897..ca204ca 100755
--- a/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86/integrations/build_kokoro.sh
+++ b/build_tools/kokoro/gcp_ubuntu/bazel/linux/x86/integrations/build_kokoro.sh

@@ -32,7 +32,7 @@
 docker_setup
 
 docker run "${DOCKER_RUN_ARGS[@]?}" \
-  gcr.io/iree-oss/bazel-tensorflow:prod \
+  gcr.io/iree-oss/bazel-tensorflow@sha256:97045a41e101c7870112e59e39a94e0e3a6dbe376a5bb189cc85b9474bff75a0 \
   build_tools/kokoro/gcp_ubuntu/bazel/linux/x86/integrations/build.sh
 
 # Kokoro will rsync this entire directory back to the executor orchestrating the

diff --git a/build_tools/kokoro/gcp_ubuntu/cmake/android/arm64-v8a/build_kokoro.sh b/build_tools/kokoro/gcp_ubuntu/cmake/android/arm64-v8a/build_kokoro.sh
index 329e226..3f22a1c 100755
--- a/build_tools/kokoro/gcp_ubuntu/cmake/android/arm64-v8a/build_kokoro.sh
+++ b/build_tools/kokoro/gcp_ubuntu/cmake/android/arm64-v8a/build_kokoro.sh

@@ -32,7 +32,7 @@
 docker_setup
 
 docker run "${DOCKER_RUN_ARGS[@]?}" \
-  gcr.io/iree-oss/cmake-android:prod \
+  gcr.io/iree-oss/cmake-android@sha256:5efb3d61e26be5ea8f26313d858c11faff5c1ed1fd83b7051d6e1c3309265e01 \
   build_tools/kokoro/gcp_ubuntu/cmake/android/build.sh arm64-v8a
 
 # Kokoro will rsync this entire directory back to the executor orchestrating the

diff --git a/build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-swiftshader/build_kokoro.sh b/build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-swiftshader/build_kokoro.sh
index 82426ec..362eaf9 100755
--- a/build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-swiftshader/build_kokoro.sh
+++ b/build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-swiftshader/build_kokoro.sh

@@ -33,7 +33,7 @@
 
 docker run "${DOCKER_RUN_ARGS[@]?}" \
   --env IREE_VULKAN_DISABLE=0 \
-  gcr.io/iree-oss/cmake-swiftshader:prod \
+  gcr.io/iree-oss/cmake-swiftshader@sha256:1912ed3a5f85c8d9abd0729834711905ef4ef03b381eb7f99d9fdb7867932d30 \
   build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-swiftshader/build.sh
 
 # Kokoro will rsync this entire directory back to the executor orchestrating the

diff --git a/build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-turing/build_kokoro.sh b/build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-turing/build_kokoro.sh
index 0b2364a..518dcd8 100755
--- a/build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-turing/build_kokoro.sh
+++ b/build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-turing/build_kokoro.sh

@@ -34,7 +34,7 @@
 docker run "${DOCKER_RUN_ARGS[@]?}" \
   --env IREE_VULKAN_DISABLE=0 \
   --gpus all \
-  gcr.io/iree-oss/cmake-nvidia:prod \
+  gcr.io/iree-oss/cmake-nvidia@sha256:26b6fcc4005b4cbb23988b9a8d5484355dd58f5a4a8df10537df8fceb1d7e26a \
   build_tools/kokoro/gcp_ubuntu/cmake/linux/x86-turing/build.sh
 
 # Kokoro will rsync this entire directory back to the executor orchestrating the

diff --git a/build_tools/kokoro/gcp_ubuntu/cmake/linux/x86/build_kokoro.sh b/build_tools/kokoro/gcp_ubuntu/cmake/linux/x86/build_kokoro.sh
index ee459ff..b3d684e 100755
--- a/build_tools/kokoro/gcp_ubuntu/cmake/linux/x86/build_kokoro.sh
+++ b/build_tools/kokoro/gcp_ubuntu/cmake/linux/x86/build_kokoro.sh

@@ -31,7 +31,7 @@
 docker_setup
 
 docker run "${DOCKER_RUN_ARGS[@]?}" \
-  gcr.io/iree-oss/cmake:prod \
+  gcr.io/iree-oss/cmake@sha256:da2de0066bf5e9607fe35c4e7816b19e779d47cd29a0a5f889a66c9c64aded57 \
   build_tools/kokoro/gcp_ubuntu/cmake/linux/x86/build.sh
 
 # Kokoro will rsync this entire directory back to the executor orchestrating the

diff --git a/docs/design_docs/metal_hal_driver.md b/docs/design_docs/metal_hal_driver.md
index 89ade9a..0729d9a 100644
--- a/docs/design_docs/metal_hal_driver.md
+++ b/docs/design_docs/metal_hal_driver.md

@@ -58,6 +58,8 @@
 [`hal::CommandQueue`][hal-command-queue]   | [`MTLCommandQueue`][mtl-command-queue]
 [`hal::CommandBuffer`][hal-command-buffer] | [`MTLCommandBuffer`][mtl-command-buffer]
 [`hal::Semaphore`][hal-semaphore]          | [`MTLSharedEvent`][mtl-shared-event]
+[`hal::Allocator`][hal-allocator]          | N/A
+[`hal::Buffer`][hal-buffer]                | [`MTLBuffer`][mtl-buffer]
 
 In the following subsections, we go over each pair to provide more details.
 
@@ -124,25 +126,77 @@
 `notifyListener:atValue:block:` to singal a semaphore to wake the current
 thread, which is put into sleep by waiting on the semaphore.
 
+### Allocator
+
+At the moment the Metal HAL driver just has a very simple
+[`hal::Allocator`][hal-allocator] implementation. It just wraps a `MTLDevice`
+and redirects all allocation requests to the `MTLDevice`. No page/pool/slab or
+whatever. This is only meant to get started. In the future we should have a
+better memory allocation library, probably by layering the
+[Vulkan Memory Allocator][vma] on top of [`MTLHeap`][mtl-heap].
+
+### Buffer
+
+IREE [`hal::Buffer`][hal-buffer] maps Metal `MTLBuffer`. See
+[Memory Management](#memory-management) for more details.
+
+## Memory Management
+
+### Storage type
+
+Metal provides four [`MTLStorageMode`][mtl-storage-mode] options:
+
+*   `MTLStorageModeShared`: The resource is stored in system memory and is
+    accessible to both the CPU and the GPU.
+*   `MTLStorageModeManaged`: The CPU and GPU may maintain separate copies of the
+    resource, and any changes must be explicitly synchronized.
+*   `MTLStorageModePrivate`: The resource can be accessed only by the GPU.
+*   `MTLStorageMemoryless`: The resource’s contents can be accessed only by the
+    GPU and only exist temporarily during a render pass.
+
+Among them, `MTLStorageModeManaged` is only available on macOS.
+
+IREE HAL defines serveral [`MemoryType`][hal-buffer]. They need to map to the
+above storage modes:
+
+*   If `kDeviceLocal` but not `kHostVisible`, `MTLStorageModePrivate` is chosen.
+*   If `kDeviceLocal` and `kHostVisible`:
+    *   If macOS, `MTLStorageModeManaged` can be chosen.
+    *   Otherwise, `MTLStorageModeShared` is chosen.
+*   If not `DeviceLocal` but `kDeviceVisible`, `MTLStorageModeShared` is chosen.
+*   If not `kDeviceLocal` and not `kDeviceVisible`, `MTLStorageModeShared` is
+    chosen. (TODO: We should probably use host buffer here.)
+
+IREE HAL also allows to create buffers with `kHostCoherent` bit. This may still
+be backed by `MTLStorageModeManaged` `MTLBuffer`s in macOS. To respect the
+`kHostCoherent` protocol, the Metal HAL driver will perform necessary
+`InValidate`/`Flush` operations automatically under the hood.
+
 [macos-version-share]: https://gs.statcounter.com/macos-version-market-share/desktop/worldwide
 [ios-version-share]: https://developer.apple.com/support/app-store/
 [iree-hal]: https://github.com/google/iree/tree/main/iree/hal
 [iree-metal]: https://github.com/google/iree/tree/main/iree/hal/metal
 [iree-refptr]: https://github.com/google/iree/blob/main/iree/base/ref_ptr.h
-[hal-driver]: https://github.com/google/iree/blob/main/iree/hal/driver.h
-[hal-device]: https://github.com/google/iree/blob/main/iree/hal/device.h
+[hal-allocator]: https://github.com/google/iree/blob/main/iree/hal/allocator.h
+[hal-buffer]: https://github.com/google/iree/blob/main/iree/hal/buffer.h
 [hal-command-queue]: https://github.com/google/iree/blob/main/iree/hal/command_queue.h
 [hal-command-buffer]: https://github.com/google/iree/blob/main/iree/hal/command_buffer.h
+[hal-device]: https://github.com/google/iree/blob/main/iree/hal/device.h
+[hal-driver]: https://github.com/google/iree/blob/main/iree/hal/driver.h
 [hal-semaphore]: https://github.com/google/iree/blob/main/iree/hal/semaphore.h
-[metal-driver]: https://github.com/google/iree/blob/main/iree/hal/metal/metal_driver.h
-[metal-device]: https://github.com/google/iree/blob/main/iree/hal/metal/metal_device.h
 [metal-command-queue]: https://github.com/google/iree/blob/main/iree/hal/metal/metal_command_queue.h
 [metal-command-buffer]: https://github.com/google/iree/blob/main/iree/hal/metal/metal_command_buffer.h
+[metal-device]: https://github.com/google/iree/blob/main/iree/hal/metal/metal_device.h
+[metal-driver]: https://github.com/google/iree/blob/main/iree/hal/metal/metal_driver.h
 [metal-shared-event]: https://github.com/google/iree/blob/main/iree/hal/metal/metal_shared_event.h
-[mtl-device]: https://developer.apple.com/documentation/metal/mtldevice?language=objc
-[mtl-command-queue]: https://developer.apple.com/documentation/metal/mtlcommandqueue?language=objc
+[mtl-buffer]: https://developer.apple.com/documentation/metal/mtlbuffer?language=objc
 [mtl-command-buffer]: https://developer.apple.com/documentation/metal/mtlcommandbuffer?language=objc
 [mtl-command-encoder]: https://developer.apple.com/documentation/metal/mtlcommandencoder?language=objc
+[mtl-command-queue]: https://developer.apple.com/documentation/metal/mtlcommandqueue?language=objc
+[mtl-device]: https://developer.apple.com/documentation/metal/mtldevice?language=objc
+[mtl-heap]: https://developer.apple.com/documentation/metal/mtlheap?language=objc
 [mtl-shared-event]: https://developer.apple.com/documentation/metal/mtlsharedevent?language=objc
+[mtl-storage-mode]: https://developer.apple.com/documentation/metal/mtlstoragemode?language=objc
 [objc-arc]: https://en.wikipedia.org/wiki/Automatic_Reference_Counting
 [objcxx]: https://en.wikipedia.org/wiki/Objective-C#Objective-C++
+[vma]: https://github.com/GPUOpen-LibrariesAndSDKs/VulkanMemoryAllocator

diff --git a/integrations/tensorflow/e2e/bool_test.py b/integrations/tensorflow/e2e/bool_test.py
index 2b29f2e..161f6ad 100644
--- a/integrations/tensorflow/e2e/bool_test.py
+++ b/integrations/tensorflow/e2e/bool_test.py

@@ -21,6 +21,10 @@
 
 class MathModule(tf.Module):
 
+  @tf.function(input_signature=[])
+  def constant(self):
+    return np.array([True, False, True], dtype=np.bool)
+
   @tf.function(input_signature=[tf.TensorSpec([4], tf.float32)])
   def greater_than(self, x):
     return x > 1.0
@@ -36,6 +40,13 @@
 @tf_test_utils.compile_module(MathModule)
 class BooleanTest(tf_test_utils.TracedModuleTestCase):
 
+  def test_constant(self):
+
+    def constant(module):
+      module.constant()
+
+    self.compare_backends(constant)
+
   def test_greater_than(self):
 
     def greater_than(module):

diff --git a/iree/compiler/Dialect/HAL/Conversion/FlowToHAL/ConvertTensorOps.cpp b/iree/compiler/Dialect/HAL/Conversion/FlowToHAL/ConvertTensorOps.cpp
index 6fcadbd..6efc149 100644
--- a/iree/compiler/Dialect/HAL/Conversion/FlowToHAL/ConvertTensorOps.cpp
+++ b/iree/compiler/Dialect/HAL/Conversion/FlowToHAL/ConvertTensorOps.cpp

@@ -57,9 +57,25 @@
         IREE::HAL::BufferUsageBitfield::All |
         IREE::HAL::BufferUsageBitfield::Constant;
 
+    auto elementsAttr = constantOp.getValue().cast<ElementsAttr>();
+    auto elementsTy = elementsAttr.getType().cast<ShapedType>();
+
+    // Expand boolean elements to the minimum bit widht supported by the HAL
+    // (8-bits).
+    // To improve memory bandwidth and increase computae we should prefer to
+    // pack 1-bit tensors into wider storage before this lossy conversion. For
+    // example bitwise ops on 8x32xi1 can be converted to ops on tensor<8xi32>.
+    if (elementsTy.getElementType().isInteger(1)) {
+      elementsAttr =
+          elementsAttr.mapValues(rewriter.getIntegerType(8),
+                                 llvm::function_ref<APInt(const APInt &val)>(
+                                     [](const APInt &val) -> APInt {
+                                       return APInt(8, val.getBoolValue());
+                                     }));
+    }
+
     auto buffer = rewriter.createOrFold<IREE::HAL::AllocatorAllocateConstOp>(
-        constantOp.getLoc(), allocator, memoryTypes, bufferUsage,
-        constantOp.getValue().cast<ElementsAttr>());
+        constantOp.getLoc(), allocator, memoryTypes, bufferUsage, elementsAttr);
 
     // TODO(benvanik): implement resource sets.
     rewriter.create<IREE::HAL::ExDeferReleaseOp>(constantOp.getLoc(), buffer);

diff --git a/iree/compiler/Dialect/HAL/Conversion/FlowToHAL/test/stream_ops.mlir b/iree/compiler/Dialect/HAL/Conversion/FlowToHAL/test/stream_ops.mlir
index aad814a..084af52 100644
--- a/iree/compiler/Dialect/HAL/Conversion/FlowToHAL/test/stream_ops.mlir
+++ b/iree/compiler/Dialect/HAL/Conversion/FlowToHAL/test/stream_ops.mlir

@@ -27,14 +27,13 @@
   // CHECK: %[[CMD:.+]] = hal.command_buffer.create {{.+}}, "OneShot", "Transfer|Dispatch"
   // CHECK-NEXT: hal.command_buffer.begin %[[CMD]]
   %0 = flow.ex.stream.fragment(%arg1 = %cst : index, %arg2 = %arg0 : tensor<128xf32>) -> tensor<128xf32> {
-    //  CHECK-DAG: %[[EXE:.+]] = hal.executable.lookup {{.+}}, @ex0 : !hal.executable
     //  CHECK-DAG: %[[EXE_LAYOUT:.+]] = hal.executable_layout.lookup
     //      CHECK: hal.command_buffer.push_descriptor_set %[[CMD]], %[[EXE_LAYOUT]], set=0, bindings=[0 = (%arg0, %c0, %sz_3), 1 = (%buffer_1, %c0, %sz_4)]
-    //      CHECK: hal.command_buffer.dispatch.symbol {{.+}}, entry_point = @ex0::@vmla::@entry0, workgroup_xyz
+    //      CHECK: hal.command_buffer.dispatch.symbol {{.+}}, @ex0::@vmla::@entry0, workgroup_xyz
     //      CHECK: hal.command_buffer.execution_barrier
     %1 = flow.dispatch @ex0::@entry0[%arg1 : index](%arg2) : (tensor<128xf32>) -> tensor<128xf32>
     //      CHECK: hal.command_buffer.push_descriptor_set
-    //      CHECK: hal.command_buffer.dispatch.symbol {{.+}}, entry_point = @ex0::@vmla::@entry0, workgroup_xyz
+    //      CHECK: hal.command_buffer.dispatch.symbol {{.+}}, @ex0::@vmla::@entry0, workgroup_xyz
     //      CHECK: hal.command_buffer.execution_barrier
     %2 = flow.dispatch @ex0::@entry0[%arg1 : index](%1) : (tensor<128xf32>) -> tensor<128xf32>
     flow.return %2 : tensor<128xf32>

diff --git a/iree/compiler/Dialect/HAL/Conversion/FlowToHAL/test/tensor_ops.mlir b/iree/compiler/Dialect/HAL/Conversion/FlowToHAL/test/tensor_ops.mlir
index 52d67f8..89bd590 100644
--- a/iree/compiler/Dialect/HAL/Conversion/FlowToHAL/test/tensor_ops.mlir
+++ b/iree/compiler/Dialect/HAL/Conversion/FlowToHAL/test/tensor_ops.mlir

@@ -15,7 +15,7 @@
 func @constantTensor1() {
   // CHECK-NEXT: %dev = hal.ex.shared_device
   // CHECK-NEXT: %allocator = hal.device.allocator %dev
-  // CHECK-NEXT: %cbuffer = hal.allocator.allocate.const %allocator, {{.+}} = dense<[true, false]> : tensor<2xi1>
+  // CHECK-NEXT: %cbuffer = hal.allocator.allocate.const %allocator, {{.+}} = dense<[1, 0]> : tensor<2xi8>
   %0 = constant dense<[1, 0]> : tensor<2xi1>
   return
 }

diff --git a/iree/compiler/Dialect/HAL/IR/HALOpFolders.cpp b/iree/compiler/Dialect/HAL/IR/HALOpFolders.cpp
index efdd27a..626fc01 100644
--- a/iree/compiler/Dialect/HAL/IR/HALOpFolders.cpp
+++ b/iree/compiler/Dialect/HAL/IR/HALOpFolders.cpp

@@ -256,6 +256,36 @@
 }
 
 //===----------------------------------------------------------------------===//
+// iree::hal::CommandBuffer
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+/// Skips a hal.command_buffer.device accessor when the device was created in
+/// the same scope.
+struct SkipCommandBufferDeviceOp
+    : public OpRewritePattern<CommandBufferDeviceOp> {
+  using OpRewritePattern<CommandBufferDeviceOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(CommandBufferDeviceOp op,
+                                PatternRewriter &rewriter) const override {
+    if (auto createOp = dyn_cast_or_null<CommandBufferCreateOp>(
+            op.command_buffer().getDefiningOp())) {
+      rewriter.replaceOp(op, createOp.device());
+      return success();
+    }
+    return failure();
+  }
+};
+
+}  // namespace
+
+void CommandBufferDeviceOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &results, MLIRContext *context) {
+  results.insert<SkipCommandBufferDeviceOp>(context);
+}
+
+//===----------------------------------------------------------------------===//
 // hal.device.switch
 //===----------------------------------------------------------------------===//
 

diff --git a/iree/compiler/Dialect/HAL/IR/HALOps.cpp b/iree/compiler/Dialect/HAL/IR/HALOps.cpp
index 1f9e448..540234b 100644
--- a/iree/compiler/Dialect/HAL/IR/HALOps.cpp
+++ b/iree/compiler/Dialect/HAL/IR/HALOps.cpp

@@ -853,10 +853,9 @@
 
 void CommandBufferDispatchSymbolOp::build(
     OpBuilder &builder, OperationState &state, Value commandBuffer,
-    Value executable, IREE::HAL::ExecutableEntryPointOp entryPoint,
-    Value workgroupX, Value workgroupY, Value workgroupZ) {
-  state.addOperands(
-      {commandBuffer, executable, workgroupX, workgroupY, workgroupZ});
+    IREE::HAL::ExecutableEntryPointOp entryPoint, Value workgroupX,
+    Value workgroupY, Value workgroupZ) {
+  state.addOperands({commandBuffer, workgroupX, workgroupY, workgroupZ});
   // Construct Executable::Target::EntryPoint nested reference.
   StringRef executableOpSymName =
       entryPoint.getParentOp()

diff --git a/iree/compiler/Dialect/HAL/IR/HALOps.td b/iree/compiler/Dialect/HAL/IR/HALOps.td
index 3f3f70e..42856a8 100644
--- a/iree/compiler/Dialect/HAL/IR/HALOps.td
+++ b/iree/compiler/Dialect/HAL/IR/HALOps.td

@@ -1001,6 +1001,24 @@
   let assemblyFormat = "$command_buffer attr-dict";
 }
 
+def HAL_CommandBufferDeviceOp : HAL_PureOp<"command_buffer.device"> {
+  let summary = [{command buffer device query operation}];
+  let description = [{
+    Used during conversion to access the device used to create a command buffer.
+  }];
+
+  let arguments = (ins
+    HAL_CommandBuffer:$command_buffer
+  );
+  let results = (outs
+    HAL_Device:$device
+  );
+
+  let assemblyFormat = "$command_buffer attr-dict `:` type($device)";
+
+  let hasCanonicalizer = 1;
+}
+
 def HAL_CommandBufferExecutionBarrierOp : HAL_Op<"command_buffer.execution_barrier", [
     AttrSizedOperandSegments,
   ]> {
@@ -1179,15 +1197,13 @@
     %x = constant 128 : index
     %y = constant 32 : index
     %z = constant 1 : index
-    hal.command_buffer.dispatch.symbol %cmd, %executable,
-                                       entry_point = @executable::@target::@entry0,
+    hal.command_buffer.dispatch.symbol %cmd, @executable::@target::@entry,
                                        workgroup_xyz = [%x, %y, %z]
     ```
   }];
 
   let arguments = (ins
     HAL_CommandBuffer:$command_buffer,
-    HAL_Executable:$executable,  // TODO(scotttodd): remove this and extract from nested ref?
     SymbolRefAttr:$entry_point,
     HAL_Dim:$workgroup_x,
     HAL_Dim:$workgroup_y,
@@ -1195,7 +1211,7 @@
   );
 
   let assemblyFormat = [{
-    $command_buffer `,` $executable `,` `entry_point` `=` $entry_point `,`
+    $command_buffer `,` $entry_point `,`
     `workgroup_xyz` `=` `[` $workgroup_x `,` $workgroup_y `,` $workgroup_z `]`
     attr-dict
   }];
@@ -1204,7 +1220,7 @@
   let builders = [
     OpBuilder<[{
       OpBuilder &builder, OperationState &state, Value commandBuffer,
-      Value executable, IREE::HAL::ExecutableEntryPointOp entryPoint,
+      IREE::HAL::ExecutableEntryPointOp entryPoint,
       Value workgroupX, Value workgroupY, Value workgroupZ
     }]>,
   ];
@@ -1248,22 +1264,20 @@
     given buffer, using using a nested symbol reference to the entry point.
 
     ```mlir
-    hal.command_buffer.dispatch.indirect.symbol %cmd, %executable,
-                                                entry_point = @executable::@target::@entry0,
+    hal.command_buffer.dispatch.indirect.symbol %cmd, @executable::@target::@entry,
                                                 workgroups = %buffer[%offset]
     ```
   }];
 
   let arguments = (ins
     HAL_CommandBuffer:$command_buffer,
-    HAL_Executable:$executable,  // TODO(scotttodd): remove this and extract from nested ref?
     SymbolRefAttr:$entry_point,
     HAL_Buffer:$workgroups_buffer,
     HAL_DeviceSize:$workgroups_offset
   );
 
   let assemblyFormat = [{
-    $command_buffer `,` $executable `,` `entry_point` `=` $entry_point `,`
+    $command_buffer `,` $entry_point `,`
     `workgroups` `=` $workgroups_buffer `[` $workgroups_offset `]` attr-dict
   }];
 }

diff --git a/iree/compiler/Dialect/HAL/IR/test/command_buffer_folding.mlir b/iree/compiler/Dialect/HAL/IR/test/command_buffer_folding.mlir
new file mode 100644
index 0000000..a004780
--- /dev/null
+++ b/iree/compiler/Dialect/HAL/IR/test/command_buffer_folding.mlir

@@ -0,0 +1,16 @@
+// Tests folding and canonicalization of HAL command buffer ops.
+
+// RUN: iree-opt -split-input-file -canonicalize %s | iree-opt -split-input-file | IreeFileCheck %s
+
+// CHECK-LABEL: @skip_command_buffer_device
+func @skip_command_buffer_device() -> !hal.executable {
+  %dev = hal.ex.shared_device : !hal.device
+  %cmd = hal.command_buffer.create %dev, "OneShot", "Transfer|Dispatch" : !hal.command_buffer
+
+  // CHECK-NOT: hal.command_buffer.device
+  // CHECK: %[[EXECUTABLE:.+]] = hal.executable.lookup %dev, @executable_name : !hal.executable
+  %0 = hal.command_buffer.device %cmd : !hal.device
+  %exe = hal.executable.lookup %0, @executable_name : !hal.executable
+
+  return %exe : !hal.executable
+}

diff --git a/iree/compiler/Dialect/HAL/IR/test/command_buffer_ops.mlir b/iree/compiler/Dialect/HAL/IR/test/command_buffer_ops.mlir
index 168725d..c2e77bc 100644
--- a/iree/compiler/Dialect/HAL/IR/test/command_buffer_ops.mlir
+++ b/iree/compiler/Dialect/HAL/IR/test/command_buffer_ops.mlir

@@ -42,6 +42,15 @@
 
 // -----
 
+// CHECK-LABEL: @command_buffer_device
+func @command_buffer_device(%arg0 : !hal.command_buffer) {
+  // CHECK: %0 = hal.command_buffer.device %arg0 : !hal.device
+  %0 = hal.command_buffer.device %arg0 : !hal.device
+  return
+}
+
+// -----
+
 // CHECK-LABEL: @command_buffer_execution_barrier
 func @command_buffer_execution_barrier(%arg0 : !hal.command_buffer) {
   %0 = "test_hal.buffer"() : () -> !hal.buffer
@@ -115,12 +124,11 @@
       }
     }
   }
-  %0 = "test_hal.executable"() : () -> !hal.executable
-  %1 = "test_hal.workgroup_x"() : () -> index
-  %2 = "test_hal.workgroup_y"() : () -> index
-  %3 = "test_hal.workgroup_z"() : () -> index
-  // CHECK: hal.command_buffer.dispatch.symbol %arg0, %0, entry_point = @ex::@backend::@entry0, workgroup_xyz = [%1, %2, %3]
-  hal.command_buffer.dispatch.symbol %arg0, %0, entry_point = @ex::@backend::@entry0, workgroup_xyz = [%1, %2, %3]
+  %0 = "test_hal.workgroup_x"() : () -> index
+  %1 = "test_hal.workgroup_y"() : () -> index
+  %2 = "test_hal.workgroup_z"() : () -> index
+  // CHECK: hal.command_buffer.dispatch.symbol %arg0, @ex::@backend::@entry0, workgroup_xyz = [%0, %1, %2]
+  hal.command_buffer.dispatch.symbol %arg0, @ex::@backend::@entry0, workgroup_xyz = [%0, %1, %2]
   return
 }
 
@@ -135,10 +143,9 @@
       }
     }
   }
-  %0 = "test_hal.executable"() : () -> !hal.executable
-  %1 = "test_hal.buffer"() : () -> !hal.buffer
-  %2 = "test_hal.offset"() : () -> index
-  // CHECK: hal.command_buffer.dispatch.indirect.symbol %arg0, %0, entry_point = @ex::@backend::@entry0, workgroups = %1[%2]
-  hal.command_buffer.dispatch.indirect.symbol %arg0, %0, entry_point = @ex::@backend::@entry0, workgroups = %1[%2]
+  %0 = "test_hal.buffer"() : () -> !hal.buffer
+  %1 = "test_hal.offset"() : () -> index
+  // CHECK: hal.command_buffer.dispatch.indirect.symbol %arg0, @ex::@backend::@entry0, workgroups = %0[%1]
+  hal.command_buffer.dispatch.indirect.symbol %arg0, @ex::@backend::@entry0, workgroups = %0[%1]
   return
 }

diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/AOT/LLVMAOTTarget.cpp b/iree/compiler/Dialect/HAL/Target/LLVM/AOT/LLVMAOTTarget.cpp
index 0107c79..6fa091b 100644
--- a/iree/compiler/Dialect/HAL/Target/LLVM/AOT/LLVMAOTTarget.cpp
+++ b/iree/compiler/Dialect/HAL/Target/LLVM/AOT/LLVMAOTTarget.cpp

@@ -42,8 +42,9 @@
   LLVMAOTTargetBackend(LLVMTargetOptions options)
       : options_(std::move(options)) {}
 
-  // NOTE: we could vary this based on the options, such as by arch/etc.
-  std::string name() const override { return "dylib*"; }
+  // NOTE: we could vary these based on the options, such as by arch/etc.
+  std::string name() const override { return "llvm_aot"; }
+  std::string filter_pattern() const override { return "dylib*"; }
 
   void getDependentDialects(DialectRegistry& registry) const override {
     // clang-format off
@@ -55,17 +56,6 @@
     // clang-format on
   }
 
-  void declareTargetOps(IREE::Flow::ExecutableOp sourceOp,
-                        IREE::HAL::ExecutableOp executableOp) override {
-    OpBuilder targetBuilder(&executableOp.getBlock().back());
-    auto targetContainerOp =
-        targetBuilder.create<IREE::HAL::ExecutableTargetOp>(
-            sourceOp.getLoc(), /*name=*/"llvm_aot",
-            /*targetBackendFilter=*/name());
-    OpBuilder containerBuilder(&targetContainerOp.getBlock().back());
-    containerBuilder.create<ModuleOp>(sourceOp.getLoc());
-  }
-
   void buildTranslationPassPipeline(ExecutableTargetOp targetOp,
                                     OpPassManager& passManager) override {
     buildLLVMTransformPassPipeline(passManager);

diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/IR/LLVMIRTarget.cpp b/iree/compiler/Dialect/HAL/Target/LLVM/IR/LLVMIRTarget.cpp
index 9d32820..fac9612 100644
--- a/iree/compiler/Dialect/HAL/Target/LLVM/IR/LLVMIRTarget.cpp
+++ b/iree/compiler/Dialect/HAL/Target/LLVM/IR/LLVMIRTarget.cpp

@@ -39,8 +39,9 @@
   LLVMIRTargetBackend(LLVMTargetOptions options)
       : options_(std::move(options)) {}
 
-  // NOTE: we could vary this based on the options, such as by arch/etc.
-  std::string name() const override { return "llvm-ir*"; }
+  // NOTE: we could vary these based on the options, such as by arch/etc.
+  std::string name() const override { return "llvm_ir"; }
+  std::string filter_pattern() const override { return "llvm-ir*"; }
 
   void getDependentDialects(DialectRegistry& registry) const override {
     // clang-format off
@@ -52,17 +53,6 @@
     // clang-format on
   }
 
-  void declareTargetOps(IREE::Flow::ExecutableOp sourceOp,
-                        IREE::HAL::ExecutableOp executableOp) override {
-    OpBuilder targetBuilder(&executableOp.getBlock().back());
-    auto targetContainerOp =
-        targetBuilder.create<IREE::HAL::ExecutableTargetOp>(
-            sourceOp.getLoc(), /*name=*/"llvm_ir",
-            /*targetBackendFilter=*/name());
-    OpBuilder containerBuilder(&targetContainerOp.getBlock().back());
-    containerBuilder.create<ModuleOp>(sourceOp.getLoc());
-  }
-
   void buildTranslationPassPipeline(ExecutableTargetOp targetOp,
                                     OpPassManager& passManager) override {
     buildLLVMTransformPassPipeline(passManager);

diff --git a/iree/compiler/Dialect/HAL/Target/TargetBackend.cpp b/iree/compiler/Dialect/HAL/Target/TargetBackend.cpp
index ac21f6b..b126991 100644
--- a/iree/compiler/Dialect/HAL/Target/TargetBackend.cpp
+++ b/iree/compiler/Dialect/HAL/Target/TargetBackend.cpp

@@ -75,7 +75,7 @@
                                      IREE::HAL::ExecutableOp executableOp) {
   OpBuilder targetBuilder(&executableOp.getBlock().back());
   auto targetContainerOp = targetBuilder.create<IREE::HAL::ExecutableTargetOp>(
-      sourceOp.getLoc(), /*name=*/name(), /*targetBackendFilter=*/name());
+      sourceOp.getLoc(), name(), filter_pattern());
   OpBuilder containerBuilder(&targetContainerOp.getBlock().back());
   containerBuilder.create<ModuleOp>(sourceOp.getLoc());
 }
@@ -135,24 +135,22 @@
     Location loc, DispatchState dispatchState,
     DeviceSwitchBuilder &switchBuilder) {
   auto *region = switchBuilder.addConditionRegion(
-      IREE::HAL::DeviceMatchIDAttr::get(name(), loc.getContext()),
+      IREE::HAL::DeviceMatchIDAttr::get(filter_pattern(), loc.getContext()),
       {
           dispatchState.workload,
           dispatchState.commandBuffer,
-          dispatchState.executable,
       });
   auto &entryBlock = region->front();
   auto workload = entryBlock.getArgument(0);
   auto commandBuffer = entryBlock.getArgument(1);
-  auto executable = entryBlock.getArgument(2);
 
   auto builder = OpBuilder::atBlockBegin(&entryBlock);
   auto workgroupCount = calculateDispatchWorkgroupCount(
       loc, dispatchState.executableOp, dispatchState.entryPointOp, workload,
       builder);
   builder.create<IREE::HAL::CommandBufferDispatchSymbolOp>(
-      loc, commandBuffer, executable, dispatchState.entryPointOp,
-      workgroupCount[0], workgroupCount[1], workgroupCount[2]);
+      loc, commandBuffer, dispatchState.entryPointOp, workgroupCount[0],
+      workgroupCount[1], workgroupCount[2]);
 
   builder.create<IREE::HAL::ReturnOp>(loc);
   return success();

diff --git a/iree/compiler/Dialect/HAL/Target/TargetBackend.h b/iree/compiler/Dialect/HAL/Target/TargetBackend.h
index f087a48..7182a4d 100644
--- a/iree/compiler/Dialect/HAL/Target/TargetBackend.h
+++ b/iree/compiler/Dialect/HAL/Target/TargetBackend.h

@@ -118,9 +118,11 @@
 
   virtual ~TargetBackend() = default;
 
-  // Returns the name of the backend as expected to be matched with a call to
-  // matchPattern. For example, 'vulkan-v1.1' or 'vmla*'.
+  // Returns a name for the backend used to differentiate between other targets.
   virtual std::string name() const = 0;
+  // Returns a filter pattern for the backend as expected to be matched with a
+  // call to matchPattern. For example, 'vulkan-v1.1' or 'vmla*'.
+  virtual std::string filter_pattern() const = 0;
 
   // Creates an interface representing the bindings and push constants required
   // to dispatch the executable. Interfaces used across backends and executables

diff --git a/iree/compiler/Dialect/HAL/Target/VMLA/VMLATarget.cpp b/iree/compiler/Dialect/HAL/Target/VMLA/VMLATarget.cpp
index 3db3986..ae429f5 100644
--- a/iree/compiler/Dialect/HAL/Target/VMLA/VMLATarget.cpp
+++ b/iree/compiler/Dialect/HAL/Target/VMLA/VMLATarget.cpp

@@ -48,6 +48,7 @@
   VMLATargetBackend(VMLATargetOptions options) : options_(std::move(options)) {}
 
   std::string name() const override { return "vmla"; }
+  std::string filter_pattern() const override { return "vmla"; }
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<VM::VMDialect, VMLA::VMLADialect>();

diff --git a/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/VulkanSPIRVTarget.cpp b/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/VulkanSPIRVTarget.cpp
index 2161067..6851ec2 100644
--- a/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/VulkanSPIRVTarget.cpp
+++ b/iree/compiler/Dialect/HAL/Target/VulkanSPIRV/VulkanSPIRVTarget.cpp

@@ -233,8 +233,9 @@
   VulkanSPIRVTargetBackend(VulkanSPIRVTargetOptions options)
       : options_(std::move(options)) {}
 
-  // NOTE: we could vary this based on the options such as 'vulkan-v1.1'.
-  std::string name() const override { return "vulkan*"; }
+  // NOTE: we could vary these based on the options such as 'vulkan-v1.1'.
+  std::string name() const override { return "vulkan_spirv"; }
+  std::string filter_pattern() const override { return "vulkan*"; }
 
   void getDependentDialects(DialectRegistry &registry) const override {
     // clang-format off
@@ -252,8 +253,7 @@
                         IREE::HAL::ExecutableOp executableOp) override {
     OpBuilder targetBuilder(&executableOp.getBlock().back());
     auto targetOp = targetBuilder.create<IREE::HAL::ExecutableTargetOp>(
-        sourceOp.getLoc(), /*name=*/"vulkan_any",
-        /*targetBackendFilter=*/name());
+        sourceOp.getLoc(), name(), filter_pattern());
     OpBuilder containerBuilder(&targetOp.getBlock().back());
 
     auto innerModuleOp = containerBuilder.create<ModuleOp>(sourceOp.getLoc());
@@ -281,7 +281,8 @@
     IREE::HAL::ExecutableOp executableOp = dispatchState.executableOp;
     for (auto executableTargetOp :
          executableOp.getBlock().getOps<IREE::HAL::ExecutableTargetOp>()) {
-      if (matchPattern(executableTargetOp.target_backend_filter(), name())) {
+      if (matchPattern(executableTargetOp.target_backend_filter(),
+                       filter_pattern())) {
         ModuleOp innerModuleOp = executableTargetOp.getInnerModule();
         auto spvModuleOps = innerModuleOp.getOps<spirv::ModuleOp>();
         assert(llvm::hasSingleElement(spvModuleOps));
@@ -324,7 +325,7 @@
     }
 
     auto *region = switchBuilder.addConditionRegion(
-        IREE::HAL::DeviceMatchIDAttr::get(name(), loc.getContext()),
+        IREE::HAL::DeviceMatchIDAttr::get(filter_pattern(), loc.getContext()),
         {
             dispatchState.workload,
             dispatchState.commandBuffer,
@@ -410,7 +411,8 @@
     spirv::ModuleOp spvModuleOp;
     for (auto executableTargetOp :
          executableOp.getBlock().getOps<IREE::HAL::ExecutableTargetOp>()) {
-      if (matchPattern(executableTargetOp.target_backend_filter(), name())) {
+      if (matchPattern(executableTargetOp.target_backend_filter(),
+                       filter_pattern())) {
         ModuleOp innerModuleOp = executableTargetOp.getInnerModule();
         assert(!innerModuleOp.getAttr(
             iree_compiler::getEntryPointScheduleAttrName()));

diff --git a/iree/compiler/Dialect/HAL/Transforms/MaterializeResourceCaches.cpp b/iree/compiler/Dialect/HAL/Transforms/MaterializeResourceCaches.cpp
index 1c1c9e4..94b567d 100644
--- a/iree/compiler/Dialect/HAL/Transforms/MaterializeResourceCaches.cpp
+++ b/iree/compiler/Dialect/HAL/Transforms/MaterializeResourceCaches.cpp

@@ -65,7 +65,7 @@
     // loads from variables.
     for (auto funcOp : moduleOp.getOps<FuncOp>()) {
       for (auto &block : funcOp) {
-        for (auto &op : llvm::make_early_inc_range(block)) {
+        block.walk([&](Operation *op) {
           if (auto lookupOp = dyn_cast<DescriptorSetLayoutLookupOp>(op)) {
             replaceDescriptorSetLayoutLookupOp(lookupOp);
           } else if (auto lookupOp = dyn_cast<ExecutableLayoutLookupOp>(op)) {
@@ -73,7 +73,7 @@
           } else if (auto lookupOp = dyn_cast<ExecutableLookupOp>(op)) {
             replaceExecutableLookupOp(lookupOp);
           }
-        }
+        });
       }
     }
 

diff --git a/iree/compiler/Dialect/HAL/Transforms/Passes.cpp b/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
index d17f726..1a9440d 100644
--- a/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
+++ b/iree/compiler/Dialect/HAL/Transforms/Passes.cpp

@@ -75,6 +75,13 @@
   // been expanded to primitives.
   passManager.addPass(createPublicABIGenerationPass());
 
+  // Resolve entry point ordinals from nested symbol references prior to
+  // serialization. As this pass creates lookup ops it should run before
+  // MaterializeResourceCachesPass.
+  passManager.addPass(createResolveEntryPointOrdinalsPass());
+  passManager.addNestedPass<FuncOp>(createCanonicalizerPass());
+  passManager.addNestedPass<FuncOp>(createCSEPass());
+
   // Gather cachable resources such as executables and descriptor sets and
   // cache them at initialization-time.
   passManager.addPass(createMaterializeResourceCachesPass(targetOptions));
@@ -86,10 +93,6 @@
   passManager.addNestedPass<FuncOp>(createCanonicalizerPass());
   passManager.addNestedPass<FuncOp>(createCSEPass());
 
-  // Resolve entry point ordinals from nested symbol references prior to
-  // serialization.
-  passManager.addPass(createResolveEntryPointOrdinalsPass());
-
   // TODO(#1036): run this once per hal.executable.target in a nested pass
   // manager so that we have as many passes as hal.executable.target ops.
   if (transformOptions.serializeExecutables) {

diff --git a/iree/compiler/Dialect/HAL/Transforms/ResolveEntryPointOrdinals.cpp b/iree/compiler/Dialect/HAL/Transforms/ResolveEntryPointOrdinals.cpp
index 51f07da..03030d2 100644
--- a/iree/compiler/Dialect/HAL/Transforms/ResolveEntryPointOrdinals.cpp
+++ b/iree/compiler/Dialect/HAL/Transforms/ResolveEntryPointOrdinals.cpp

@@ -31,8 +31,19 @@
                                 PatternRewriter &rewriter) const override {
     auto entryPointOp = dyn_cast<IREE::HAL::ExecutableEntryPointOp>(
         SymbolTable::lookupNearestSymbolFrom(op, op.entry_point()));
+
+    // Lookup the device for our command buffer, then the executable from the
+    // entry point's nested reference.
+    auto device = rewriter.createOrFold<IREE::HAL::CommandBufferDeviceOp>(
+        op.getLoc(), IREE::HAL::DeviceType::get(rewriter.getContext()),
+        op.command_buffer());
+    auto executableOp = dyn_cast<IREE::HAL::ExecutableOp>(
+        entryPointOp.getParentOp()->getParentOp());
+    auto executable = rewriter.createOrFold<IREE::HAL::ExecutableLookupOp>(
+        op.getLoc(), device, executableOp.sym_name());
+
     rewriter.replaceOpWithNewOp<IREE::HAL::CommandBufferDispatchOp>(
-        op, op.command_buffer(), op.executable(), entryPointOp.ordinalAttr(),
+        op, op.command_buffer(), executable, entryPointOp.ordinalAttr(),
         op.workgroup_x(), op.workgroup_y(), op.workgroup_z());
     return success();
   }
@@ -49,8 +60,19 @@
       PatternRewriter &rewriter) const override {
     auto entryPointOp = dyn_cast<IREE::HAL::ExecutableEntryPointOp>(
         SymbolTable::lookupNearestSymbolFrom(op, op.entry_point()));
+
+    // Lookup the device for our command buffer, then the executable from the
+    // entry point's nested reference.
+    auto device = rewriter.createOrFold<IREE::HAL::CommandBufferDeviceOp>(
+        op.getLoc(), IREE::HAL::DeviceType::get(rewriter.getContext()),
+        op.command_buffer());
+    auto executableOp = dyn_cast<IREE::HAL::ExecutableOp>(
+        entryPointOp.getParentOp()->getParentOp());
+    auto executable = rewriter.createOrFold<IREE::HAL::ExecutableLookupOp>(
+        op.getLoc(), device, executableOp.sym_name());
+
     rewriter.replaceOpWithNewOp<IREE::HAL::CommandBufferDispatchIndirectOp>(
-        op, op.command_buffer(), op.executable(), entryPointOp.ordinalAttr(),
+        op, op.command_buffer(), executable, entryPointOp.ordinalAttr(),
         op.workgroups_buffer(), op.workgroups_offset());
     return success();
   }

diff --git a/iree/compiler/Dialect/HAL/Transforms/test/resolve_entry_point_ordinals.mlir b/iree/compiler/Dialect/HAL/Transforms/test/resolve_entry_point_ordinals.mlir
index f254b6e..c5855ca 100644
--- a/iree/compiler/Dialect/HAL/Transforms/test/resolve_entry_point_ordinals.mlir
+++ b/iree/compiler/Dialect/HAL/Transforms/test/resolve_entry_point_ordinals.mlir

@@ -19,12 +19,13 @@
 
   func @dispatch_with_nested_references() {
     %cmd = "test_hal.command_buffer"() : () -> !hal.command_buffer
-    %exe = "test_hal.executable"() : () -> !hal.executable
     %x = "test_hal.workgroup_x"() : () -> index
     %y = "test_hal.workgroup_y"() : () -> index
     %z = "test_hal.workgroup_z"() : () -> index
-    // CHECK: hal.command_buffer.dispatch %0, %1, entry_point = 0, workgroup_xyz = [%2, %3, %4]
-    hal.command_buffer.dispatch.symbol %cmd, %exe, entry_point = @exe::@target::@entry, workgroup_xyz = [%x, %y, %z]
+    // CHECK: %[[DEVICE:.+]] = hal.command_buffer.device %0
+    // CHECK: %[[EXE:.+]] = hal.executable.lookup %[[DEVICE]], @exe
+    // CHECK: hal.command_buffer.dispatch %0, %[[EXE]], entry_point = 0, workgroup_xyz = [%1, %2, %3]
+    hal.command_buffer.dispatch.symbol %cmd, @exe::@target::@entry, workgroup_xyz = [%x, %y, %z]
     return
   }
 }
@@ -66,11 +67,12 @@
 
   func @dispatch_indirect_with_nested_references() {
     %cmd = "test_hal.command_buffer"() : () -> !hal.command_buffer
-    %exe = "test_hal.executable"() : () -> !hal.executable
     %buffer = "test_hal.buffer"() : () -> !hal.buffer
     %offset = "test_hal.offset"() : () -> index
-    // CHECK: hal.command_buffer.dispatch.indirect %0, %1, entry_point = 0, workgroups = %2[%3]
-    hal.command_buffer.dispatch.indirect.symbol %cmd, %exe, entry_point = @exe::@target::@entry, workgroups = %buffer[%offset]
+    // CHECK: %[[DEVICE:.+]] = hal.command_buffer.device %0
+    // CHECK: %[[EXE:.+]] = hal.executable.lookup %[[DEVICE]], @exe
+    // CHECK: hal.command_buffer.dispatch.indirect %0, %[[EXE]], entry_point = 0, workgroups = %1[%2]
+    hal.command_buffer.dispatch.indirect.symbol %cmd, @exe::@target::@entry, workgroups = %buffer[%offset]
     return
   }
 }

diff --git a/iree/compiler/Dialect/VMLA/Conversion/StandardToVMLA/ConvertStandardToVMLA.cpp b/iree/compiler/Dialect/VMLA/Conversion/StandardToVMLA/ConvertStandardToVMLA.cpp
index 1c790d6..1d4b7e0 100644
--- a/iree/compiler/Dialect/VMLA/Conversion/StandardToVMLA/ConvertStandardToVMLA.cpp
+++ b/iree/compiler/Dialect/VMLA/Conversion/StandardToVMLA/ConvertStandardToVMLA.cpp

@@ -42,6 +42,15 @@
       ConversionPatternRewriter &rewriter) const override {
     auto value = srcOp.value().dyn_cast<ElementsAttr>();
     if (!value) return failure();
+
+    if (value.getType().getElementType().isInteger(1)) {
+      value = value.mapValues(rewriter.getIntegerType(8),
+                              llvm::function_ref<APInt(const APInt &val)>(
+                                  [](const APInt &val) -> APInt {
+                                    return APInt(8, val.getBoolValue());
+                                  }));
+    }
+
     rewriter.replaceOpWithNewOp<IREE::VMLA::ConstantOp>(srcOp, value);
     return success();
   }

diff --git a/iree/compiler/Dialect/VMLA/Conversion/StandardToVMLA/test/constant_ops.mlir b/iree/compiler/Dialect/VMLA/Conversion/StandardToVMLA/test/constant_ops.mlir
index 4971289..87307e9 100644
--- a/iree/compiler/Dialect/VMLA/Conversion/StandardToVMLA/test/constant_ops.mlir
+++ b/iree/compiler/Dialect/VMLA/Conversion/StandardToVMLA/test/constant_ops.mlir

@@ -15,3 +15,12 @@
   %0 = constant dense<[-1.0, -2.0, 3.0, 4.0]> : tensor<4xf32>
   return %0 : tensor<4xf32>
 }
+
+// -----
+
+// CHECK-LABEL: @constant_tensor_bool
+func @constant_tensor_bool() -> tensor<4xi1> attributes { sym_visibility = "private" } {
+  // CHECK: = vmla.constant dense<[0, 1, 1, 0]> : tensor<4xi8> -> !vmla.buffer
+  %0 = constant dense<[false, true, true, false]> : tensor<4xi1>
+  return %0 : tensor<4xi1>
+}

diff --git a/iree/hal/buffer_test.cc b/iree/hal/buffer_test.cc
index 4e627da..bf58539 100644
--- a/iree/hal/buffer_test.cc
+++ b/iree/hal/buffer_test.cc

@@ -43,7 +43,7 @@
 
   // We don't currently do any padding on the host.
   // Other implementations may differ.
-  EXPECT_GE(14, buffer->allocation_size());
+  EXPECT_LE(14, buffer->allocation_size());
   EXPECT_EQ(0, buffer->byte_offset());
   EXPECT_EQ(14, buffer->byte_length());
 
@@ -69,7 +69,7 @@
       HeapBuffer::AllocateCopy(BufferUsage::kTransfer | BufferUsage::kMapping,
                                src_data.data(), src_data.size());
   EXPECT_NE(nullptr, buffer->allocator());
-  EXPECT_GE(src_data.size(), buffer->allocation_size());
+  EXPECT_LE(src_data.size(), buffer->allocation_size());
 
   // Data should have been copied.
   std::vector<uint8_t> actual_data(src_data.size());
@@ -99,7 +99,7 @@
   EXPECT_NE(nullptr, buffer->allocator());
   EXPECT_EQ(MemoryType::kHostLocal, buffer->memory_type());
   EXPECT_EQ(BufferUsage::kTransfer | BufferUsage::kMapping, buffer->usage());
-  EXPECT_GE(src_data.size() * sizeof(int32_t), buffer->allocation_size());
+  EXPECT_LE(src_data.size() * sizeof(int32_t), buffer->allocation_size());
 
   // Data should have been copied.
   std::vector<int32_t> actual_data(src_data.size());

diff --git a/iree/hal/cts/BUILD b/iree/hal/cts/BUILD
index 445e5af..cbee0cf 100644
--- a/iree/hal/cts/BUILD
+++ b/iree/hal/cts/BUILD

@@ -56,6 +56,18 @@
 )
 
 cc_test(
+    name = "buffer_test",
+    srcs = ["buffer_test.cc"],
+    deps = [
+        ":cts_test_base",
+        "//iree/base:status",
+        "//iree/hal:driver_registry",
+        "//iree/testing:gtest",
+        "//iree/testing:gtest_main",
+    ],
+)
+
+cc_test(
     name = "command_buffer_test",
     srcs = ["command_buffer_test.cc"],
     deps = [

diff --git a/iree/hal/cts/CMakeLists.txt b/iree/hal/cts/CMakeLists.txt
index 3de079d..1271a0c 100644
--- a/iree/hal/cts/CMakeLists.txt
+++ b/iree/hal/cts/CMakeLists.txt

@@ -47,6 +47,18 @@
 
 iree_cc_test(
   NAME
+    buffer_test
+  SRCS
+    "buffer_test.cc"
+  DEPS
+    ::cts_test_base
+    iree::base::status
+    iree::hal::driver_registry
+    iree::testing::gtest
+)
+
+iree_cc_test(
+  NAME
     command_buffer_test
   SRCS
     "command_buffer_test.cc"

diff --git a/iree/hal/cts/allocator_test.cc b/iree/hal/cts/allocator_test.cc
index 00d72b4..1608c62 100644
--- a/iree/hal/cts/allocator_test.cc
+++ b/iree/hal/cts/allocator_test.cc

@@ -65,6 +65,18 @@
   EXPECT_GE(buffer->allocation_size(), allocation_size);  // Larger is okay.
 }
 
+TEST_P(AllocatorTest, CanUseBufferLike) {
+  MemoryType memory_type = MemoryType::kHostLocal | MemoryType::kDeviceVisible;
+  BufferUsage usage = BufferUsage::kMapping;
+  size_t allocation_size = 1024;
+
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto buffer, allocator_->Allocate(memory_type, usage, allocation_size));
+  // Using the buffer for its original requested purpose should be fine.
+  EXPECT_TRUE(
+      allocator_->CanUseBufferLike(allocator_, memory_type, usage, usage));
+}
+
 INSTANTIATE_TEST_SUITE_P(AllDrivers, AllocatorTest,
                          ::testing::ValuesIn(DriverRegistry::shared_registry()
                                                  ->EnumerateAvailableDrivers()),

diff --git a/iree/hal/cts/buffer_test.cc b/iree/hal/cts/buffer_test.cc
new file mode 100644
index 0000000..26eecc7
--- /dev/null
+++ b/iree/hal/cts/buffer_test.cc

@@ -0,0 +1,385 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "iree/base/status.h"
+#include "iree/hal/cts/cts_test_base.h"
+#include "iree/hal/driver_registry.h"
+#include "iree/testing/gtest.h"
+#include "iree/testing/status_matchers.h"
+
+namespace iree {
+namespace hal {
+namespace cts {
+
+using ::testing::_;
+using ::testing::ElementsAre;
+using ::testing::Eq;
+
+// Note: this file only covers hal::Buffer APIs that can be overridden by
+// subclasses. Errors caught by hal::Buffer's common validations are not
+// covered as they are already tested in iree/hal/buffer_test.cc.
+
+class BufferTest : public CtsTestBase {};
+
+TEST_P(BufferTest, Allocate) {
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto buffer, device_->allocator()->Allocate(
+                       MemoryType::kHostLocal | MemoryType::kDeviceVisible,
+                       BufferUsage::kTransfer | BufferUsage::kMapping, 14));
+
+  EXPECT_NE(nullptr, buffer->allocator());
+  EXPECT_EQ(MemoryAccess::kAll, buffer->allowed_access());
+  EXPECT_EQ(MemoryType::kHostLocal | MemoryType::kDeviceVisible,
+            buffer->memory_type());
+  EXPECT_EQ(BufferUsage::kTransfer | BufferUsage::kMapping, buffer->usage());
+
+  EXPECT_LE(14, buffer->allocation_size());
+  EXPECT_EQ(0, buffer->byte_offset());
+  EXPECT_EQ(14, buffer->byte_length());
+}
+
+TEST_P(BufferTest, AllocateZeroLength) {
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto buffer, device_->allocator()->Allocate(
+                       MemoryType::kHostLocal | MemoryType::kDeviceVisible,
+                       BufferUsage::kTransfer | BufferUsage::kMapping, 0));
+  EXPECT_LE(0, buffer->allocation_size());
+}
+
+TEST_P(BufferTest, Fill8) {
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto buffer, device_->allocator()->Allocate(
+                       MemoryType::kHostLocal | MemoryType::kDeviceVisible,
+                       BufferUsage::kTransfer | BufferUsage::kMapping, 5));
+
+  std::vector<uint8_t> actual_data(buffer->allocation_size());
+
+  // Fill with a sentinel.
+  IREE_EXPECT_OK(buffer->Fill8(0, buffer->allocation_size(), 0x33u));
+  IREE_EXPECT_OK(buffer->ReadData(0, actual_data.data(), actual_data.size()));
+  EXPECT_THAT(actual_data, ElementsAre(0x33, 0x33, 0x33, 0x33, 0x33));
+
+  // Zero fills are fine.
+  IREE_EXPECT_OK(buffer->Fill8(0, 0, 0x44u));
+  IREE_EXPECT_OK(buffer->ReadData(0, actual_data.data(), actual_data.size()));
+  EXPECT_THAT(actual_data, ElementsAre(0x33, 0x33, 0x33, 0x33, 0x33));
+
+  // Fill the remaining parts of the buffer by using kWholeBuffer.
+  IREE_EXPECT_OK(buffer->Fill8(2, kWholeBuffer, 0x55u));
+  IREE_EXPECT_OK(buffer->ReadData(0, actual_data.data(), actual_data.size()));
+  EXPECT_THAT(actual_data, ElementsAre(0x33, 0x33, 0x55, 0x55, 0x55));
+
+  // Fill a small region of the buffer.
+  IREE_EXPECT_OK(buffer->Fill8(1, 1, 0x66u));
+  IREE_EXPECT_OK(buffer->ReadData(0, actual_data.data(), actual_data.size()));
+  EXPECT_THAT(actual_data, ElementsAre(0x33, 0x66, 0x55, 0x55, 0x55));
+
+  // Whole buffer helper.
+  IREE_EXPECT_OK(buffer->Fill8(0x99u));
+  IREE_EXPECT_OK(buffer->ReadData(0, actual_data.data(), actual_data.size()));
+  EXPECT_THAT(actual_data, ElementsAre(0x99, 0x99, 0x99, 0x99, 0x99));
+}
+
+TEST_P(BufferTest, Fill16) {
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto buffer, device_->allocator()->Allocate(
+                       MemoryType::kHostLocal | MemoryType::kDeviceVisible,
+                       BufferUsage::kTransfer | BufferUsage::kMapping, 9));
+
+  std::vector<uint8_t> actual_data(buffer->allocation_size());
+
+  // Fill with a sentinel.
+  IREE_EXPECT_OK(buffer->Fill16(0, 4, 0x1122u));
+  IREE_EXPECT_OK(buffer->ReadData(0, actual_data.data(), actual_data.size()));
+  EXPECT_THAT(actual_data, ElementsAre(0x22, 0x11, 0x22, 0x11, 0, 0, 0, 0, 0));
+
+  // Zero fills are fine.
+  IREE_EXPECT_OK(buffer->Fill16(0, 0, 0x5566u));
+  IREE_EXPECT_OK(buffer->ReadData(0, actual_data.data(), actual_data.size()));
+  EXPECT_THAT(actual_data, ElementsAre(0x22, 0x11, 0x22, 0x11, 0, 0, 0, 0, 0));
+
+  // Fill the remaining parts of the buffer by using kWholeBuffer.
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto aligned_buffer,
+      device_->allocator()->Allocate(
+          MemoryType::kHostLocal | MemoryType::kDeviceVisible,
+          BufferUsage::kTransfer | BufferUsage::kMapping, 8));
+  IREE_EXPECT_OK(aligned_buffer->Fill16(4, kWholeBuffer, 0x5566u));
+  std::vector<uint8_t> aligned_actual_data(aligned_buffer->allocation_size());
+  IREE_EXPECT_OK(aligned_buffer->ReadData(0, aligned_actual_data.data(),
+                                          aligned_actual_data.size()));
+  EXPECT_THAT(aligned_actual_data,
+              ElementsAre(0, 0, 0, 0, 0x66, 0x55, 0x66, 0x55));
+
+  // Whole buffer helper.
+  IREE_EXPECT_OK(aligned_buffer->Fill16(0x5566u));
+  IREE_EXPECT_OK(aligned_buffer->ReadData(0, aligned_actual_data.data(),
+                                          aligned_actual_data.size()));
+  EXPECT_THAT(aligned_actual_data,
+              ElementsAre(0x66, 0x55, 0x66, 0x55, 0x66, 0x55, 0x66, 0x55));
+}
+
+TEST_P(BufferTest, Fill32) {
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto buffer, device_->allocator()->Allocate(
+                       MemoryType::kHostLocal | MemoryType::kDeviceVisible,
+                       BufferUsage::kTransfer | BufferUsage::kMapping, 9));
+
+  std::vector<uint8_t> actual_data(buffer->allocation_size());
+
+  // Fill with a sentinel.
+  IREE_EXPECT_OK(buffer->Fill32(0, 8, 0x11223344u));
+  IREE_EXPECT_OK(buffer->ReadData(0, actual_data.data(), actual_data.size()));
+  EXPECT_THAT(actual_data,
+              ElementsAre(0x44, 0x33, 0x22, 0x11, 0x44, 0x33, 0x22, 0x11, 0));
+
+  // Zero fills are fine.
+  IREE_EXPECT_OK(buffer->Fill32(0, 0, 0x55667788u));
+  IREE_EXPECT_OK(buffer->ReadData(0, actual_data.data(), actual_data.size()));
+  EXPECT_THAT(actual_data,
+              ElementsAre(0x44, 0x33, 0x22, 0x11, 0x44, 0x33, 0x22, 0x11, 0));
+
+  // Fill the remaining parts of the buffer by using kWholeBuffer.
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto aligned_buffer,
+      device_->allocator()->Allocate(
+          MemoryType::kHostLocal | MemoryType::kDeviceVisible,
+          BufferUsage::kTransfer | BufferUsage::kMapping, 8));
+  IREE_EXPECT_OK(aligned_buffer->Fill32(4, kWholeBuffer, 0x55667788u));
+  std::vector<uint8_t> aligned_actual_data(aligned_buffer->allocation_size());
+  IREE_EXPECT_OK(aligned_buffer->ReadData(0, aligned_actual_data.data(),
+                                          aligned_actual_data.size()));
+  EXPECT_THAT(aligned_actual_data,
+              ElementsAre(0, 0, 0, 0, 0x88, 0x77, 0x66, 0x55));
+
+  // Whole buffer helper.
+  IREE_EXPECT_OK(aligned_buffer->Fill32(0x55667788u));
+  IREE_EXPECT_OK(aligned_buffer->ReadData(0, aligned_actual_data.data(),
+                                          aligned_actual_data.size()));
+  EXPECT_THAT(aligned_actual_data,
+              ElementsAre(0x88, 0x77, 0x66, 0x55, 0x88, 0x77, 0x66, 0x55));
+}
+
+TEST_P(BufferTest, ReadWriteData) {
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto buffer, device_->allocator()->Allocate(
+                       MemoryType::kHostLocal | MemoryType::kDeviceVisible,
+                       BufferUsage::kTransfer | BufferUsage::kMapping, 4));
+
+  std::vector<uint8_t> actual_data(4);
+
+  // Write over the entire buffer.
+  std::vector<uint8_t> new_data = {10, 20, 30, 40};
+  IREE_EXPECT_OK(buffer->WriteData(0, new_data.data(), new_data.size()));
+  IREE_EXPECT_OK(buffer->ReadData(0, actual_data.data(), actual_data.size()));
+  EXPECT_THAT(actual_data, Eq(new_data));
+
+  // Writing zero bytes is valid.
+  std::vector<uint8_t> zero_data;
+  IREE_EXPECT_OK(buffer->WriteData(0, zero_data.data(), 0));
+  IREE_EXPECT_OK(buffer->ReadData(0, actual_data.data(), actual_data.size()));
+  EXPECT_THAT(actual_data, Eq(new_data));
+
+  // Write over a portion of the buffer.
+  std::vector<uint8_t> partial_data = {99};
+  IREE_EXPECT_OK(
+      buffer->WriteData(1, partial_data.data(), partial_data.size()));
+  IREE_EXPECT_OK(buffer->ReadData(0, actual_data.data(), actual_data.size()));
+  EXPECT_THAT(actual_data, ElementsAre(10, 99, 30, 40));
+}
+
+TEST_P(BufferTest, CopyData) {
+  std::vector<uint8_t> src_data = {0, 1, 2, 3};
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto src_buffer,
+      device_->allocator()->Allocate(
+          MemoryType::kHostLocal | MemoryType::kDeviceVisible,
+          BufferUsage::kTransfer | BufferUsage::kMapping, src_data.size()));
+  IREE_EXPECT_OK(src_buffer->WriteData(0, src_data.data(), src_data.size()));
+
+  std::vector<uint8_t> dst_data = {0, 1, 2, 3, 4};
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto dst_buffer,
+      device_->allocator()->Allocate(
+          MemoryType::kHostLocal | MemoryType::kDeviceVisible,
+          BufferUsage::kTransfer | BufferUsage::kMapping, dst_data.size()));
+  IREE_EXPECT_OK(dst_buffer->WriteData(0, dst_data.data(), dst_data.size()));
+
+  // Copy of length 0 should not change the dest buffer.
+  IREE_EXPECT_OK(dst_buffer->CopyData(0, src_buffer.get(), 0, 0));
+  std::vector<uint8_t> actual_data(dst_data.size());
+  IREE_EXPECT_OK(
+      dst_buffer->ReadData(0, actual_data.data(), actual_data.size()));
+  EXPECT_THAT(actual_data, Eq(dst_data));
+
+  // Copy a subrange of the buffer.
+  IREE_EXPECT_OK(dst_buffer->CopyData(1, src_buffer.get(), 2, 2));
+  IREE_EXPECT_OK(
+      dst_buffer->ReadData(0, actual_data.data(), actual_data.size()));
+  EXPECT_THAT(actual_data, ElementsAre(0, 2, 3, 3, 4));
+
+  // Copy the entire buffer using kWholeBuffer. This will adjust sizes
+  // to ensure that the min buffer is taken. We test both src and dst buffer
+  // offset/length calculations (note that some may end up as 0 copies).
+  IREE_EXPECT_OK(dst_buffer->CopyData(3, src_buffer.get(), 0, kWholeBuffer));
+  IREE_EXPECT_OK(
+      dst_buffer->ReadData(0, actual_data.data(), actual_data.size()));
+  EXPECT_THAT(actual_data, ElementsAre(0, 2, 3, 0, 1));
+  IREE_EXPECT_OK(dst_buffer->CopyData(0, src_buffer.get(), 2, kWholeBuffer));
+  IREE_EXPECT_OK(
+      dst_buffer->ReadData(0, actual_data.data(), actual_data.size()));
+  EXPECT_THAT(actual_data, ElementsAre(2, 3, 3, 0, 1));
+  IREE_EXPECT_OK(dst_buffer->CopyData(0, src_buffer.get(), 3, kWholeBuffer));
+  IREE_EXPECT_OK(
+      dst_buffer->ReadData(0, actual_data.data(), actual_data.size()));
+  EXPECT_THAT(actual_data, ElementsAre(3, 3, 3, 0, 1));
+  IREE_EXPECT_OK(dst_buffer->CopyData(4, src_buffer.get(), 0, kWholeBuffer));
+  IREE_EXPECT_OK(
+      dst_buffer->ReadData(0, actual_data.data(), actual_data.size()));
+  EXPECT_THAT(actual_data, ElementsAre(3, 3, 3, 0, 0));
+}
+
+TEST_P(BufferTest, MapMemory) {
+  std::vector<uint8_t> src_data = {0, 1, 2, 3, 4, 5, 6};
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto buffer,
+      device_->allocator()->Allocate(
+          MemoryType::kHostLocal | MemoryType::kDeviceVisible,
+          BufferUsage::kTransfer | BufferUsage::kMapping, src_data.size()));
+  IREE_EXPECT_OK(buffer->WriteData(0, src_data.data(), src_data.size()));
+
+  // 0-length mappings are valid.
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto mapping, buffer->MapMemory<uint8_t>(MemoryAccess::kRead, 0, 0));
+  EXPECT_TRUE(mapping.empty());
+  EXPECT_EQ(0, mapping.size());
+  EXPECT_EQ(0, mapping.byte_length());
+  EXPECT_NE(nullptr, mapping.data());
+  IREE_ASSERT_OK_AND_ASSIGN(auto span, mapping.Subspan());
+  EXPECT_TRUE(span.empty());
+  mapping.reset();
+
+  // Map the whole buffer for reading.
+  IREE_ASSERT_OK_AND_ASSIGN(mapping, buffer->MapMemory<uint8_t>(
+                                         MemoryAccess::kRead, 0, kWholeBuffer));
+  EXPECT_EQ(src_data.size(), mapping.size());
+  IREE_ASSERT_OK_AND_ASSIGN(span, mapping.Subspan());
+  EXPECT_THAT(span, ElementsAre(0, 1, 2, 3, 4, 5, 6));
+  mapping.reset();
+
+  // Map a portion of the buffer for reading.
+  IREE_ASSERT_OK_AND_ASSIGN(
+      mapping, buffer->MapMemory<uint8_t>(MemoryAccess::kRead, 1, 2));
+  EXPECT_EQ(2, mapping.size());
+  IREE_ASSERT_OK_AND_ASSIGN(span, mapping.Subspan());
+  EXPECT_THAT(span, ElementsAre(1, 2));
+  mapping.reset();
+}
+
+TEST_P(BufferTest, MapMemoryNonByte) {
+  std::vector<uint8_t> src_data = {0, 1, 2, 3, 4, 5, 6};
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto buffer,
+      device_->allocator()->Allocate(
+          MemoryType::kHostLocal | MemoryType::kDeviceVisible,
+          BufferUsage::kTransfer | BufferUsage::kMapping, src_data.size()));
+  IREE_EXPECT_OK(buffer->WriteData(0, src_data.data(), src_data.size()));
+
+  // Map the buffer as non-byte values.
+  // Note that we'll round down to the number of valid elements at the
+  // alignment.
+  IREE_ASSERT_OK_AND_ASSIGN(auto mapping16,
+                            buffer->MapMemory<uint16_t>(MemoryAccess::kRead));
+  EXPECT_EQ(3, mapping16.size());
+  EXPECT_LE(6, mapping16.byte_length());
+  IREE_ASSERT_OK_AND_ASSIGN(auto span16, mapping16.Subspan());
+  EXPECT_THAT(span16, ElementsAre(0x0100, 0x0302, 0x0504));
+  mapping16.reset();
+}
+
+TEST_P(BufferTest, MapMemoryWrite) {
+  std::vector<uint8_t> src_data = {0, 1, 2, 3, 4, 5, 6};
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto buffer,
+      device_->allocator()->Allocate(
+          MemoryType::kHostLocal | MemoryType::kDeviceVisible,
+          BufferUsage::kTransfer | BufferUsage::kMapping, src_data.size()));
+  IREE_EXPECT_OK(buffer->WriteData(0, src_data.data(), src_data.size()));
+
+  // Map and modify the data. We should see it when we read back.
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto mapping, buffer->MapMemory<uint8_t>(MemoryAccess::kWrite, 1, 2));
+  auto mutable_data = mapping.mutable_data();
+  mutable_data[0] = 0xAA;
+  mutable_data[1] = 0xBB;
+  mapping.reset();
+  std::vector<uint8_t> actual_data(src_data.size());
+  IREE_EXPECT_OK(buffer->ReadData(0, actual_data.data(), actual_data.size()));
+  EXPECT_THAT(actual_data, ElementsAre(0, 0xAA, 0xBB, 3, 4, 5, 6));
+}
+
+TEST_P(BufferTest, MapMemoryDiscard) {
+  std::vector<uint8_t> src_data = {0, 1, 2, 3, 4, 5, 6};
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto buffer,
+      device_->allocator()->Allocate(
+          MemoryType::kHostLocal | MemoryType::kDeviceVisible,
+          BufferUsage::kTransfer | BufferUsage::kMapping, src_data.size()));
+  IREE_EXPECT_OK(buffer->WriteData(0, src_data.data(), src_data.size()));
+
+  // Map for discard. Note that we can't really rely on the value of the data
+  // so we just trust that it's been discarded. It's a hint, anyway. We can be
+  // sure that the data we didn't want to discard is the same though.
+  std::vector<uint8_t> actual_data(src_data.size());
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto mapping,
+      buffer->MapMemory<uint8_t>(MemoryAccess::kDiscardWrite, 1, 2));
+  IREE_EXPECT_OK(buffer->ReadData(0, actual_data.data(), actual_data.size()));
+  EXPECT_THAT(actual_data, ElementsAre(0, _, _, 3, 4, 5, 6));
+  mapping.reset();
+}
+
+TEST_P(BufferTest, MapMemorySubspan) {
+  std::vector<uint8_t> src_data = {0, 1, 2, 3, 4, 5, 6};
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto parent_buffer,
+      device_->allocator()->Allocate(
+          MemoryType::kHostLocal | MemoryType::kDeviceVisible,
+          BufferUsage::kTransfer | BufferUsage::kMapping, src_data.size()));
+  IREE_EXPECT_OK(parent_buffer->WriteData(0, src_data.data(), src_data.size()));
+
+  IREE_ASSERT_OK_AND_ASSIGN(auto subspan_buffer,
+                            Buffer::Subspan(parent_buffer, 1, 3));
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto mapping,
+      subspan_buffer->MapMemory<uint8_t>(MemoryAccess::kDiscardWrite, 1, 2));
+  auto* mutable_data = mapping.mutable_data();
+  mutable_data[0] = 0xCC;
+  mutable_data[1] = 0xDD;
+  mapping.reset();
+
+  std::vector<uint8_t> actual_data(src_data.size());
+  IREE_EXPECT_OK(
+      parent_buffer->ReadData(0, actual_data.data(), actual_data.size()));
+  EXPECT_THAT(actual_data, ElementsAre(0, 1, 0xCC, 0xDD, 4, 5, 6));
+}
+
+INSTANTIATE_TEST_SUITE_P(AllDrivers, BufferTest,
+                         ::testing::ValuesIn(DriverRegistry::shared_registry()
+                                                 ->EnumerateAvailableDrivers()),
+                         GenerateTestName());
+
+}  // namespace cts
+}  // namespace hal
+}  // namespace iree

diff --git a/iree/hal/cts/command_buffer_test.cc b/iree/hal/cts/command_buffer_test.cc
index c146d89..5b6e4d6 100644
--- a/iree/hal/cts/command_buffer_test.cc
+++ b/iree/hal/cts/command_buffer_test.cc

@@ -12,6 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <cstring>
+#include <vector>
+
 #include "iree/base/status.h"
 #include "iree/hal/cts/cts_test_base.h"
 #include "iree/hal/driver_registry.h"
@@ -22,9 +25,24 @@
 namespace hal {
 namespace cts {
 
-class CommandBufferTest : public CtsTestBase {};
+using ::testing::ContainerEq;
 
-TEST_P(CommandBufferTest, CreateCommandBuffer) {
+class CommandBufferTest : public CtsTestBase {
+ protected:
+  static constexpr device_size_t kBufferNumBytes = 16;
+
+  void SubmitAndWait(CommandQueue* command_queue,
+                     CommandBuffer* command_buffer) {
+    IREE_ASSERT_OK_AND_ASSIGN(auto signal_semaphore,
+                              device_->CreateSemaphore(0ull));
+
+    IREE_ASSERT_OK(command_queue->Submit(
+        {{}, {command_buffer}, {{signal_semaphore.get(), 1ull}}}));
+    IREE_ASSERT_OK(signal_semaphore->Wait(1ull, InfiniteFuture()));
+  }
+};
+
+TEST_P(CommandBufferTest, Create) {
   IREE_ASSERT_OK_AND_ASSIGN(
       auto command_buffer,
       device_->CreateCommandBuffer(CommandBufferMode::kOneShot,
@@ -37,7 +55,187 @@
   EXPECT_FALSE(command_buffer->is_recording());
 }
 
-// TODO(scotttodd): Begin, End, UpdateBuffer, CopyBuffer, Dispatch, Sync, etc.
+TEST_P(CommandBufferTest, BeginEnd) {
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto command_buffer,
+      device_->CreateCommandBuffer(CommandBufferMode::kOneShot,
+                                   CommandCategory::kDispatch));
+
+  EXPECT_FALSE(command_buffer->is_recording());
+  IREE_EXPECT_OK(command_buffer->Begin());
+  EXPECT_TRUE(command_buffer->is_recording());
+  IREE_EXPECT_OK(command_buffer->End());
+  EXPECT_FALSE(command_buffer->is_recording());
+}
+
+TEST_P(CommandBufferTest, FillBufferWithRepeatedBytes) {
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto command_buffer,
+      device_->CreateCommandBuffer(CommandBufferMode::kOneShot,
+                                   CommandCategory::kTransfer));
+
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto device_buffer,
+      device_->allocator()->Allocate(
+          MemoryType::kDeviceLocal | MemoryType::kHostVisible,
+          BufferUsage::kAll, kBufferNumBytes));
+
+  std::vector<uint8_t> reference_buffer(kBufferNumBytes);
+
+  IREE_EXPECT_OK(command_buffer->Begin());
+
+  // Fill the device buffer with segments of different values so that we can
+  // test both fill and offset/size.
+
+  uint8_t val1 = 0x07;
+  IREE_EXPECT_OK(command_buffer->FillBuffer(device_buffer.get(),
+                                            /*target_offset=*/0,
+                                            /*length=*/kBufferNumBytes / 4,
+                                            &val1,
+                                            /*pattern_length=*/1));
+  std::memset(reference_buffer.data(), val1, kBufferNumBytes / 4);
+
+  uint8_t val2 = 0xbe;
+  IREE_EXPECT_OK(
+      command_buffer->FillBuffer(device_buffer.get(),
+                                 /*target_offset=*/kBufferNumBytes / 4,
+                                 /*length=*/kBufferNumBytes / 4, &val2,
+                                 /*pattern_length=*/1));
+  std::memset(reference_buffer.data() + kBufferNumBytes / 4, val2,
+              kBufferNumBytes / 4);
+
+  uint8_t val3 = 0x54;
+  IREE_EXPECT_OK(
+      command_buffer->FillBuffer(device_buffer.get(),
+                                 /*target_offset=*/kBufferNumBytes / 2,
+                                 /*length=*/kBufferNumBytes / 2, &val3,
+                                 /*pattern_length=*/1));
+  std::memset(reference_buffer.data() + kBufferNumBytes / 2, val3,
+              kBufferNumBytes / 2);
+
+  IREE_EXPECT_OK(command_buffer->End());
+
+  SubmitAndWait(device_->transfer_queues()[0], command_buffer.get());
+
+  // Read back the device buffer.
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto mapped_memory,
+      device_buffer->MapMemory<uint8_t>(MemoryAccess::kRead));
+  IREE_EXPECT_OK(mapped_memory.Invalidate());
+
+  std::vector<uint8_t> actual_data(mapped_memory.data(),
+                                   mapped_memory.data() + kBufferNumBytes);
+  EXPECT_THAT(actual_data, ContainerEq(reference_buffer));
+}
+
+TEST_P(CommandBufferTest, CopyWholeBuffer) {
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto command_buffer,
+      device_->CreateCommandBuffer(CommandBufferMode::kOneShot,
+                                   CommandCategory::kTransfer));
+
+  // Create a host buffer.
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto host_buffer, device_->allocator()->Allocate(
+                            MemoryType::kHostVisible | MemoryType::kHostCached |
+                                MemoryType::kDeviceVisible,
+                            BufferUsage::kAll, kBufferNumBytes));
+
+  // Fill the host buffer.
+  uint8_t i8_val = 0x55;
+  IREE_EXPECT_OK(host_buffer->Fill8(0, kWholeBuffer, i8_val));
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto host_mapped_memory,
+      // Cannot use kDiscard here given we filled in the above.
+      host_buffer->MapMemory<uint8_t>(MemoryAccess::kWrite));
+  IREE_EXPECT_OK(host_mapped_memory.Flush());
+
+  std::vector<uint8_t> reference_buffer(kBufferNumBytes);
+  std::memset(reference_buffer.data(), i8_val, kBufferNumBytes);
+
+  // Create a device buffer.
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto device_buffer,
+      device_->allocator()->Allocate(
+          MemoryType::kDeviceLocal | MemoryType::kHostVisible,
+          BufferUsage::kAll, kBufferNumBytes));
+
+  // Copy the host buffer to the device buffer.
+  IREE_EXPECT_OK(command_buffer->Begin());
+  IREE_EXPECT_OK(
+      command_buffer->CopyBuffer(host_buffer.get(), /*source_offset=*/0,
+                                 device_buffer.get(), /*target_offset=*/0,
+                                 /*length=*/kBufferNumBytes));
+  IREE_EXPECT_OK(command_buffer->End());
+
+  SubmitAndWait(device_->transfer_queues()[0], command_buffer.get());
+
+  // Read back the device buffer.
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto device_mapped_memory,
+      device_buffer->MapMemory<uint8_t>(MemoryAccess::kRead));
+  IREE_EXPECT_OK(device_mapped_memory.Invalidate());
+
+  std::vector<uint8_t> actual_data(
+      device_mapped_memory.data(),
+      device_mapped_memory.data() + kBufferNumBytes);
+  EXPECT_THAT(actual_data, ContainerEq(reference_buffer));
+}
+
+TEST_P(CommandBufferTest, CopySubBuffer) {
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto command_buffer,
+      device_->CreateCommandBuffer(CommandBufferMode::kOneShot,
+                                   CommandCategory::kTransfer));
+  // Create a device buffer.
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto device_buffer,
+      device_->allocator()->Allocate(
+          MemoryType::kDeviceLocal | MemoryType::kHostVisible,
+          BufferUsage::kAll, kBufferNumBytes));
+
+  // Create another host buffer with a smaller size.
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto host_buffer, device_->allocator()->Allocate(
+                            MemoryType::kHostVisible | MemoryType::kHostCached |
+                                MemoryType::kDeviceVisible,
+                            BufferUsage::kAll, kBufferNumBytes / 2));
+
+  // Fill the host buffer.
+  uint8_t i8_val = 0x88;
+  IREE_EXPECT_OK(host_buffer->Fill8(0, kWholeBuffer, i8_val));
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto host_mapped_memory,
+      // Cannot use kDiscard here given we filled in the above.
+      host_buffer->MapMemory<uint8_t>(MemoryAccess::kWrite));
+  IREE_EXPECT_OK(host_mapped_memory.Flush());
+
+  std::vector<uint8_t> reference_buffer(kBufferNumBytes);
+  std::memset(reference_buffer.data() + 8, i8_val, kBufferNumBytes / 2 - 4);
+
+  // Copy the host buffer to the device buffer.
+  IREE_EXPECT_OK(command_buffer->Begin());
+  IREE_EXPECT_OK(
+      command_buffer->CopyBuffer(host_buffer.get(), /*source_offset=*/4,
+                                 device_buffer.get(), /*target_offset=*/8,
+                                 /*length=*/kBufferNumBytes / 2 - 4));
+  IREE_EXPECT_OK(command_buffer->End());
+
+  SubmitAndWait(device_->transfer_queues()[0], command_buffer.get());
+
+  // Read back the device buffer.
+  IREE_ASSERT_OK_AND_ASSIGN(
+      auto device_mapped_memory,
+      device_buffer->MapMemory<uint8_t>(MemoryAccess::kRead));
+  IREE_EXPECT_OK(device_mapped_memory.Invalidate());
+
+  std::vector<uint8_t> actual_data(
+      device_mapped_memory.data(),
+      device_mapped_memory.data() + kBufferNumBytes);
+  EXPECT_THAT(actual_data, ContainerEq(reference_buffer));
+}
+
+// TODO(scotttodd): UpdateBuffer, Dispatch, Sync, etc.
 
 INSTANTIATE_TEST_SUITE_P(AllDrivers, CommandBufferTest,
                          ::testing::ValuesIn(DriverRegistry::shared_registry()

diff --git a/iree/hal/metal/CMakeLists.txt b/iree/hal/metal/CMakeLists.txt
index a7e643d..0f49362 100644
--- a/iree/hal/metal/CMakeLists.txt
+++ b/iree/hal/metal/CMakeLists.txt

@@ -61,6 +61,7 @@
   DEPS
     ::metal_command_buffer
     ::metal_command_queue
+    ::metal_direct_allocator
     ::metal_shared_event
     absl::strings
     absl::span
@@ -79,6 +80,26 @@
 
 iree_cc_library(
   NAME
+    metal_direct_allocator
+  HDRS
+    "metal_buffer.h"
+    "metal_direct_allocator.h"
+  SRCS
+    "metal_buffer.mm"
+    "metal_direct_allocator.mm"
+  DEPS
+    absl::memory
+    iree::base::logging
+    iree::base::status
+    iree::base::tracing
+    iree::hal::allocator
+  LINKOPTS
+    "-framework Metal"
+  PUBLIC
+)
+
+iree_cc_library(
+  NAME
     metal_driver
   HDRS
     "metal_driver.h"

diff --git a/iree/hal/metal/metal_buffer.h b/iree/hal/metal/metal_buffer.h
new file mode 100644
index 0000000..dd62e8e
--- /dev/null
+++ b/iree/hal/metal/metal_buffer.h

@@ -0,0 +1,103 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef IREE_HAL_METAL_METAL_BUFFER_H_
+#define IREE_HAL_METAL_METAL_BUFFER_H_
+
+#import <Metal/Metal.h>
+
+#include "iree/hal/buffer.h"
+
+namespace iree {
+namespace hal {
+namespace metal {
+
+class MetalDirectAllocator;
+
+// A buffer implementation for Metal that directly wraps a MTLBuffer.
+class MetalBuffer final : public Buffer {
+ public:
+  // Creates a MetalBuffer instance with retaining the given id<MTLBuffer>.
+  static StatusOr<ref_ptr<MetalBuffer>> Create(
+      MetalDirectAllocator* allocator, MemoryTypeBitfield memory_type,
+      MemoryAccessBitfield allowed_access, BufferUsageBitfield usage,
+      device_size_t allocation_size, device_size_t byte_offset,
+      device_size_t byte_length, id<MTLBuffer> buffer,
+      id<MTLCommandQueue> transfer_queue);
+
+  // Creates a MetalBuffer instance without retaining the given id<MTLBuffer>.
+  static StatusOr<ref_ptr<MetalBuffer>> CreateUnretained(
+      MetalDirectAllocator* allocator, MemoryTypeBitfield memory_type,
+      MemoryAccessBitfield allowed_access, BufferUsageBitfield usage,
+      device_size_t allocation_size, device_size_t byte_offset,
+      device_size_t byte_length, id<MTLBuffer> buffer,
+      id<MTLCommandQueue> transfer_queue);
+
+  ~MetalBuffer() override;
+
+  id<MTLBuffer> handle() const { return metal_handle_; }
+
+ private:
+  // Creates a MetalBuffer instance without retaining the given id<MTLBuffer>.
+  MetalBuffer(MetalDirectAllocator* allocator, MemoryTypeBitfield memory_type,
+              MemoryAccessBitfield allowed_access, BufferUsageBitfield usage,
+              device_size_t allocation_size, device_size_t byte_offset,
+              device_size_t byte_length, id<MTLBuffer> buffer,
+              id<MTLCommandQueue> transfer_queue);
+
+  Status FillImpl(device_size_t byte_offset, device_size_t byte_length,
+                  const void* pattern, device_size_t pattern_length) override;
+  Status ReadDataImpl(device_size_t source_offset, void* data,
+                      device_size_t data_length) override;
+  Status WriteDataImpl(device_size_t target_offset, const void* data,
+                       device_size_t data_length) override;
+  Status CopyDataImpl(device_size_t target_offset, Buffer* source_buffer,
+                      device_size_t source_offset,
+                      device_size_t data_length) override;
+
+  Status MapMemoryImpl(MappingMode mapping_mode,
+                       MemoryAccessBitfield memory_access,
+                       device_size_t local_byte_offset,
+                       device_size_t local_byte_length,
+                       void** out_data) override;
+  Status UnmapMemoryImpl(device_size_t local_byte_offset,
+                         device_size_t local_byte_length, void* data) override;
+  Status InvalidateMappedMemoryImpl(device_size_t local_byte_offset,
+                                    device_size_t local_byte_length) override;
+  Status FlushMappedMemoryImpl(device_size_t local_byte_offset,
+                               device_size_t local_byte_length) override;
+
+  // Returns true if we need to automatically invaliate/flush CPU caches to keep
+  // memory hierarchy consistent.
+  //
+  // Note: this is needed when the buffer is requested with
+  // MemoryType::kHostCoherent bit but under the hood we are using memory types
+  // that does not have that property natively, e.g., MTLStorageModeManaged.
+  // Under such circumstances, we need to perform the invalidate/flush operation
+  // "automatically" for users.
+  bool requires_autosync() const;
+
+  // We need to hold an reference to the queue so that we can encode
+  // synchronizeResource commands for synchronizing the buffer with
+  // MTLResourceStorageModeManaged.
+  id<MTLCommandQueue> metal_transfer_queue_;
+
+  id<MTLBuffer> metal_handle_;
+};
+
+}  // namespace metal
+}  // namespace hal
+}  // namespace iree
+
+#endif  // IREE_HAL_METAL_METAL_BUFFER_H_

diff --git a/iree/hal/metal/metal_buffer.mm b/iree/hal/metal/metal_buffer.mm
new file mode 100644
index 0000000..78e7e98
--- /dev/null
+++ b/iree/hal/metal/metal_buffer.mm

@@ -0,0 +1,203 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "iree/hal/metal/metal_buffer.h"
+
+#include "iree/base/status.h"
+#include "iree/base/tracing.h"
+#include "iree/hal/metal/metal_direct_allocator.h"
+
+namespace iree {
+namespace hal {
+namespace metal {
+
+// static
+StatusOr<ref_ptr<MetalBuffer>> MetalBuffer::Create(
+    MetalDirectAllocator* allocator, MemoryTypeBitfield memory_type,
+    MemoryAccessBitfield allowed_access, BufferUsageBitfield usage, device_size_t allocation_size,
+    device_size_t byte_offset, device_size_t byte_length, id<MTLBuffer> buffer,
+    id<MTLCommandQueue> transfer_queue) {
+  IREE_TRACE_SCOPE0("MetalBuffer::Create");
+  return assign_ref(new MetalBuffer(allocator, memory_type, allowed_access, usage, allocation_size,
+                                    byte_offset, byte_length, [buffer retain], transfer_queue));
+}
+
+// static
+StatusOr<ref_ptr<MetalBuffer>> MetalBuffer::CreateUnretained(
+    MetalDirectAllocator* allocator, MemoryTypeBitfield memory_type,
+    MemoryAccessBitfield allowed_access, BufferUsageBitfield usage, device_size_t allocation_size,
+    device_size_t byte_offset, device_size_t byte_length, id<MTLBuffer> buffer,
+    id<MTLCommandQueue> transfer_queue) {
+  IREE_TRACE_SCOPE0("MetalBuffer::Create");
+  return assign_ref(new MetalBuffer(allocator, memory_type, allowed_access, usage, allocation_size,
+                                    byte_offset, byte_length, buffer, transfer_queue));
+}
+
+MetalBuffer::MetalBuffer(MetalDirectAllocator* allocator, MemoryTypeBitfield memory_type,
+                         MemoryAccessBitfield allowed_access, BufferUsageBitfield usage,
+                         device_size_t allocation_size, device_size_t byte_offset,
+                         device_size_t byte_length, id<MTLBuffer> buffer,
+                         id<MTLCommandQueue> transfer_queue)
+    : Buffer(allocator, memory_type, allowed_access, usage, allocation_size, byte_offset,
+             byte_length),
+      metal_transfer_queue_([transfer_queue retain]),
+      metal_handle_(buffer) {}
+
+MetalBuffer::~MetalBuffer() {
+  IREE_TRACE_SCOPE0("MetalBuffer::dtor");
+  [metal_handle_ release];
+  [metal_transfer_queue_ release];
+}
+
+Status MetalBuffer::FillImpl(device_size_t byte_offset, device_size_t byte_length,
+                             const void* pattern, device_size_t pattern_length) {
+  IREE_ASSIGN_OR_RETURN(auto mapping,
+                        MapMemory<uint8_t>(MemoryAccess::kDiscardWrite, byte_offset, byte_length));
+  void* data_ptr = static_cast<void*>(mapping.mutable_data());
+  switch (pattern_length) {
+    case 1: {
+      uint8_t* data = static_cast<uint8_t*>(data_ptr);
+      uint8_t value_bits = *static_cast<const uint8_t*>(pattern);
+      std::fill_n(data, byte_length, value_bits);
+      break;
+    }
+    case 2: {
+      uint16_t* data = static_cast<uint16_t*>(data_ptr);
+      uint16_t value_bits = *static_cast<const uint16_t*>(pattern);
+      std::fill_n(data, byte_length / sizeof(uint16_t), value_bits);
+      break;
+    }
+    case 4: {
+      uint32_t* data = static_cast<uint32_t*>(data_ptr);
+      uint32_t value_bits = *static_cast<const uint32_t*>(pattern);
+      std::fill_n(data, byte_length / sizeof(uint32_t), value_bits);
+      break;
+    }
+    default:
+      return InvalidArgumentErrorBuilder(IREE_LOC)
+             << "Unsupported scalar data size: " << pattern_length;
+  }
+  return OkStatus();
+}
+
+Status MetalBuffer::ReadDataImpl(device_size_t source_offset, void* data,
+                                 device_size_t data_length) {
+  IREE_ASSIGN_OR_RETURN(auto mapping,
+                        MapMemory<uint8_t>(MemoryAccess::kRead, source_offset, data_length));
+  std::memcpy(data, mapping.data(), mapping.byte_length());
+  return OkStatus();
+}
+
+Status MetalBuffer::WriteDataImpl(device_size_t target_offset, const void* data,
+                                  device_size_t data_length) {
+  IREE_ASSIGN_OR_RETURN(
+      auto mapping, MapMemory<uint8_t>(MemoryAccess::kDiscardWrite, target_offset, data_length));
+  std::memcpy(mapping.mutable_data(), data, mapping.byte_length());
+  return OkStatus();
+}
+
+Status MetalBuffer::CopyDataImpl(device_size_t target_offset, Buffer* source_buffer,
+                                 device_size_t source_offset, device_size_t data_length) {
+  // This is pretty terrible. Let's not do this.
+  // TODO(benvanik): a way for allocators to indicate transfer compat.
+  IREE_ASSIGN_OR_RETURN(auto source_mapping, source_buffer->MapMemory<uint8_t>(
+                                                 MemoryAccess::kRead, source_offset, data_length));
+  CHECK_EQ(data_length, source_mapping.size());
+  IREE_ASSIGN_OR_RETURN(auto target_mapping, MapMemory<uint8_t>(MemoryAccess::kDiscardWrite,
+                                                                target_offset, data_length));
+  CHECK_EQ(data_length, target_mapping.size());
+  std::memcpy(target_mapping.mutable_data(), source_mapping.data(), data_length);
+  return OkStatus();
+}
+
+Status MetalBuffer::MapMemoryImpl(MappingMode mapping_mode, MemoryAccessBitfield memory_access,
+                                  device_size_t local_byte_offset, device_size_t local_byte_length,
+                                  void** out_data) {
+  uint8_t* data_ptr = reinterpret_cast<uint8_t*>([metal_handle_ contents]);
+  *out_data = data_ptr + local_byte_offset;
+
+  // If we mapped for discard scribble over the bytes. This is not a mandated
+  // behavior but it will make debugging issues easier. Alternatively for
+  // heap buffers we could reallocate them such that ASAN yells, but that
+  // would only work if the entire buffer was discarded.
+#ifndef NDEBUG
+  if (AnyBitSet(memory_access & MemoryAccess::kDiscard)) {
+    std::memset(data_ptr + local_byte_offset, 0xCD, local_byte_length);
+  }
+#endif  // !NDEBUG
+
+  if (requires_autosync()) {
+    IREE_RETURN_IF_ERROR(InvalidateMappedMemoryImpl(local_byte_offset, local_byte_length));
+  }
+
+  return OkStatus();
+}
+
+Status MetalBuffer::UnmapMemoryImpl(device_size_t local_byte_offset,
+                                    device_size_t local_byte_length, void* data) {
+  if (requires_autosync()) {
+    IREE_RETURN_IF_ERROR(FlushMappedMemoryImpl(local_byte_offset, local_byte_length));
+  }
+
+  return OkStatus();
+}
+
+Status MetalBuffer::InvalidateMappedMemoryImpl(device_size_t local_byte_offset,
+                                               device_size_t local_byte_length) {
+#ifdef IREE_PLATFORM_MACOS
+  // The following is only necessary for MTLStorageManaged.
+  if (metal_handle_.storageMode == MTLStorageModeManaged) {
+    @autoreleasepool {
+      id<MTLCommandBuffer> command_buffer =
+          [metal_transfer_queue_ commandBufferWithUnretainedReferences];
+
+      id<MTLBlitCommandEncoder> blit_encoder = [command_buffer blitCommandEncoder];
+      [blit_encoder synchronizeResource:metal_handle_];
+      [blit_encoder endEncoding];
+
+      [command_buffer commit];
+      [command_buffer waitUntilCompleted];
+    }
+  }
+#endif
+
+  return OkStatus();
+}
+
+Status MetalBuffer::FlushMappedMemoryImpl(device_size_t local_byte_offset,
+                                          device_size_t local_byte_length) {
+#ifdef IREE_PLATFORM_MACOS
+  // The following is only necessary for MTLStorageManaged.
+  if (metal_handle_.storageMode == MTLStorageModeManaged) {
+    [metal_handle_ didModifyRange:NSMakeRange(local_byte_offset, local_byte_length)];
+  }
+#endif
+
+  return OkStatus();
+}
+
+bool MetalBuffer::requires_autosync() const {
+  // We only need to perform "automatic" resource synchronization if it's MTLStorageModeManaged,
+  // which is only available on macOS.
+#ifdef IREE_PLATFORM_MACOS
+  return AllBitsSet(memory_type(), MemoryType::kHostCoherent) &&
+         metal_handle_.storageMode == MTLStorageModeManaged;
+#else
+  return false;
+#endif
+}
+
+}  // namespace metal
+}  // namespace hal
+}  // namespace iree

diff --git a/iree/hal/metal/metal_command_buffer.h b/iree/hal/metal/metal_command_buffer.h
index c3d532e..c58e5be 100644
--- a/iree/hal/metal/metal_command_buffer.h
+++ b/iree/hal/metal/metal_command_buffer.h

@@ -18,6 +18,7 @@
 #import <Metal/Metal.h>
 
 #include "iree/hal/command_buffer.h"
+#include "iree/hal/metal/metal_buffer.h"
 
 namespace iree {
 namespace hal {
@@ -91,8 +92,23 @@
                      CommandCategoryBitfield command_categories,
                      id<MTLCommandBuffer> command_buffer);
 
+  StatusOr<MetalBuffer*> CastBuffer(Buffer* buffer) const;
+
+  // Gets or begins an active MTLBlitCommandEncoder. This also ends all previous
+  // encoded compute commands if any.
+  id<MTLBlitCommandEncoder> GetOrBeginBlitEncoder();
+  void EndBlitEncoder();
+
+  // Gets or begins a new MTLComputeCommandEncoder. This also ends all previous
+  // encoded blit commands if any.
+  id<MTLComputeCommandEncoder> GetOrBeginComputeEncoder();
+  void EndComputeEncoder();
+
   bool is_recording_ = false;
   id<MTLCommandBuffer> metal_handle_;
+
+  id<MTLComputeCommandEncoder> current_compute_encoder_ = nil;
+  id<MTLBlitCommandEncoder> current_blit_encoder_ = nil;
 };
 
 }  // namespace metal

diff --git a/iree/hal/metal/metal_command_buffer.mm b/iree/hal/metal/metal_command_buffer.mm
index 17dec6c..8555793 100644
--- a/iree/hal/metal/metal_command_buffer.mm
+++ b/iree/hal/metal/metal_command_buffer.mm

@@ -14,6 +14,7 @@
 
 #include "iree/hal/metal/metal_command_buffer.h"
 
+#include "iree/base/logging.h"
 #include "iree/base/status.h"
 #include "iree/base/tracing.h"
 
@@ -38,6 +39,58 @@
   [metal_handle_ release];
 }
 
+StatusOr<MetalBuffer*> MetalCommandBuffer::CastBuffer(Buffer* buffer) const {
+  // TODO(benvanik): assert that the buffer is from the right allocator and
+  // that it is compatible with our target queue family.
+  return static_cast<MetalBuffer*>(buffer->allocated_buffer());
+}
+
+id<MTLBlitCommandEncoder> MetalCommandBuffer::GetOrBeginBlitEncoder() {
+  IREE_TRACE_SCOPE0("MetalCommandBuffer::GetOrBeginBlitEncoder");
+
+  if (current_compute_encoder_) EndComputeEncoder();
+
+  @autoreleasepool {
+    if (!current_blit_encoder_) {
+      current_blit_encoder_ = [[metal_handle_ blitCommandEncoder] retain];
+    }
+  }
+
+  return current_blit_encoder_;
+}
+
+void MetalCommandBuffer::EndBlitEncoder() {
+  IREE_TRACE_SCOPE0("MetalCommandBuffer::EndBlitEncoder");
+  if (current_blit_encoder_) {
+    [current_blit_encoder_ endEncoding];
+    [current_blit_encoder_ release];
+    current_blit_encoder_ = nil;
+  }
+}
+
+id<MTLComputeCommandEncoder> MetalCommandBuffer::GetOrBeginComputeEncoder() {
+  IREE_TRACE_SCOPE0("MetalCommandBuffer::GetOrBeginComputeEncoder");
+
+  if (current_blit_encoder_) EndBlitEncoder();
+
+  @autoreleasepool {
+    if (!current_compute_encoder_) {
+      current_compute_encoder_ = [[metal_handle_ computeCommandEncoder] retain];
+    }
+  }
+
+  return current_compute_encoder_;
+}
+
+void MetalCommandBuffer::EndComputeEncoder() {
+  IREE_TRACE_SCOPE0("MetalCommandBuffer::EndComputeEncoder");
+  if (current_compute_encoder_) {
+    [current_compute_encoder_ endEncoding];
+    [current_compute_encoder_ release];
+    current_compute_encoder_ = nil;
+  }
+}
+
 Status MetalCommandBuffer::Begin() {
   IREE_TRACE_SCOPE0("MetalCommandBuffer::Begin");
   is_recording_ = true;
@@ -46,6 +99,8 @@
 
 Status MetalCommandBuffer::End() {
   IREE_TRACE_SCOPE0("MetalCommandBuffer::End");
+  EndBlitEncoder();
+  EndComputeEncoder();
   is_recording_ = false;
   return OkStatus();
 }
@@ -81,12 +136,39 @@
                                       device_size_t length, const void* pattern,
                                       size_t pattern_length) {
   IREE_TRACE_SCOPE0("MetalCommandBuffer::FillBuffer");
-  return UnimplementedErrorBuilder(IREE_LOC) << "MetalCommandBuffer::FillBuffer";
+  IREE_ASSIGN_OR_RETURN(auto* target_device_buffer, CastBuffer(target_buffer));
+
+  target_offset += target_buffer->byte_offset();
+
+  // Per the spec for fillBuffer:range:value: "The alignment and length of the range must both be a
+  // multiple of 4 bytes in macOS, and 1 byte in iOS and tvOS." Although iOS/tvOS is more relaxed on
+  // this front, we still require 4-byte alignment for uniformity across IREE.
+  if (target_offset % 4 != 0) {
+    return UnimplementedErrorBuilder(IREE_LOC)
+           << "MetalCommandBuffer::FillBuffer with offset that is not a multiple of 4 bytes";
+  }
+
+  // Note that fillBuffer:range:value: only accepts a single byte as the pattern but FillBuffer
+  // can accept 1/2/4 bytes. If the pattern itself contains repeated bytes, we can call into
+  // fillBuffer:range:value:. Otherwise we may need to find another way. Just implement the case
+  // where we have a single byte to fill for now.
+  if (pattern_length != 1) {
+    return UnimplementedErrorBuilder(IREE_LOC)
+           << "MetalCommandBuffer::FillBuffer with non-1-byte pattern";
+  }
+  uint8_t byte_pattern = *reinterpret_cast<const uint8_t*>(pattern);
+
+  [GetOrBeginBlitEncoder() fillBuffer:target_device_buffer->handle()
+                                range:NSMakeRange(target_offset, length)
+                                value:byte_pattern];
+
+  return OkStatus();
 }
 
 Status MetalCommandBuffer::DiscardBuffer(Buffer* buffer) {
   IREE_TRACE_SCOPE0("MetalCommandBuffer::DiscardBuffer");
-  return UnimplementedErrorBuilder(IREE_LOC) << "MetalCommandBuffer::DiscardBuffer";
+  // This is a hint. Nothing to do for Metal.
+  return OkStatus();
 }
 
 Status MetalCommandBuffer::UpdateBuffer(const void* source_buffer, device_size_t source_offset,
@@ -100,7 +182,28 @@
                                       Buffer* target_buffer, device_size_t target_offset,
                                       device_size_t length) {
   IREE_TRACE_SCOPE0("MetalCommandBuffer::CopyBuffer");
-  return UnimplementedErrorBuilder(IREE_LOC) << "MetalCommandBuffer::CopyBuffer";
+
+  IREE_ASSIGN_OR_RETURN(auto* source_device_buffer, CastBuffer(source_buffer));
+  IREE_ASSIGN_OR_RETURN(auto* target_device_buffer, CastBuffer(target_buffer));
+
+  source_offset += source_buffer->byte_offset();
+  target_offset += target_buffer->byte_offset();
+
+  // Per the spec for copyFromBuffer:sourceOffset:toBuffer:destinationOffset:size, the source/target
+  // offset must be a multiple of 4 bytes in macOS, and 1 byte in iOS and tvOS. Although iOS/tvOS
+  // is more relaxed on this front, we still require 4-byte alignment for uniformity across IREE.
+  if (source_offset % 4 != 0 || target_offset % 4 != 0) {
+    return UnimplementedErrorBuilder(IREE_LOC)
+           << "MetalCommandBuffer::CopyBuffer with offset that is not a multiple of 4 bytes";
+  }
+
+  [GetOrBeginBlitEncoder() copyFromBuffer:source_device_buffer->handle()
+                             sourceOffset:source_offset
+                                 toBuffer:target_device_buffer->handle()
+                        destinationOffset:target_offset
+                                     size:length];
+
+  return OkStatus();
 }
 
 Status MetalCommandBuffer::PushConstants(ExecutableLayout* executable_layout, size_t offset,

diff --git a/iree/hal/metal/metal_device.h b/iree/hal/metal/metal_device.h
index 7802664..729cbaf 100644
--- a/iree/hal/metal/metal_device.h
+++ b/iree/hal/metal/metal_device.h

@@ -17,6 +17,8 @@
 
 #import <Metal/Metal.h>
 
+#include <memory>
+
 #include "absl/types/span.h"
 #include "iree/base/memory.h"
 #include "iree/hal/allocator.h"
@@ -40,7 +42,7 @@
 
   std::string DebugString() const override;
 
-  Allocator* allocator() const override { return nullptr; }
+  Allocator* allocator() const override { return allocator_.get(); }
 
   absl::Span<CommandQueue*> dispatch_queues() const override {
     return absl::MakeSpan(&common_queue_, 1);
@@ -84,6 +86,8 @@
   ref_ptr<Driver> driver_;
   id<MTLDevice> metal_handle_;
 
+  std::unique_ptr<Allocator> allocator_;
+
   // Metal does not have clear graphics/dispatch/transfer queue distinction like
   // Vulkan; one just use the same newCommandQueue() API call on MTLDevice to
   // get command queues. Command encoders differ for different categories of

diff --git a/iree/hal/metal/metal_device.mm b/iree/hal/metal/metal_device.mm
index f293641..7f97422 100644
--- a/iree/hal/metal/metal_device.mm
+++ b/iree/hal/metal/metal_device.mm

@@ -19,10 +19,12 @@
 #include "iree/base/status.h"
 #include "iree/base/time.h"
 #include "iree/base/tracing.h"
+#include "iree/hal/allocator.h"
 #include "iree/hal/command_buffer_validation.h"
 #include "iree/hal/metal/dispatch_time_util.h"
 #include "iree/hal/metal/metal_command_buffer.h"
 #include "iree/hal/metal/metal_command_queue.h"
+#include "iree/hal/metal/metal_direct_allocator.h"
 #include "iree/hal/metal/metal_shared_event.h"
 
 namespace iree {
@@ -44,6 +46,9 @@
   // Grab one queue for dispatch and transfer.
   std::string name = absl::StrCat(device_info.name(), ":queue");
   id<MTLCommandQueue> metal_queue = [metal_handle_ newCommandQueue];  // retained
+
+  allocator_ = MetalDirectAllocator::Create(metal_handle_, metal_queue);
+
   command_queue_ = absl::make_unique<MetalCommandQueue>(
       name, CommandCategory::kDispatch | CommandCategory::kTransfer, metal_queue);
   common_queue_ = command_queue_.get();

diff --git a/iree/hal/metal/metal_direct_allocator.h b/iree/hal/metal/metal_direct_allocator.h
new file mode 100644
index 0000000..bf8dde8
--- /dev/null
+++ b/iree/hal/metal/metal_direct_allocator.h

@@ -0,0 +1,82 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef IREE_HAL_METAL_METAL_DIRECT_ALLOCATOR_H_
+#define IREE_HAL_METAL_METAL_DIRECT_ALLOCATOR_H_
+
+#import <Metal/Metal.h>
+
+#include <memory>
+
+#include "iree/base/status.h"
+#include "iree/hal/allocator.h"
+
+namespace iree {
+namespace hal {
+namespace metal {
+
+class MetalBuffer;
+
+// An allocator implementation for Metal that directly wraps a MTLDevice and
+// requests all allocations on the device. This is not of great performance,
+// but good for start.
+class MetalDirectAllocator final : public Allocator {
+ public:
+  static std::unique_ptr<MetalDirectAllocator> Create(
+      id<MTLDevice> device, id<MTLCommandQueue> transfer_queue);
+
+  ~MetalDirectAllocator() override;
+
+  bool CanUseBufferLike(Allocator* source_allocator,
+                        MemoryTypeBitfield memory_type,
+                        BufferUsageBitfield buffer_usage,
+                        BufferUsageBitfield intended_usage) const override;
+
+  bool CanAllocate(MemoryTypeBitfield memory_type,
+                   BufferUsageBitfield buffer_usage,
+                   size_t allocation_size) const override;
+
+  Status MakeCompatible(MemoryTypeBitfield* memory_type,
+                        BufferUsageBitfield* buffer_usage) const override;
+
+  StatusOr<ref_ptr<Buffer>> Allocate(MemoryTypeBitfield memory_type,
+                                     BufferUsageBitfield buffer_usage,
+                                     size_t allocation_size) override;
+
+  StatusOr<ref_ptr<Buffer>> AllocateConstant(
+      BufferUsageBitfield buffer_usage, ref_ptr<Buffer> source_buffer) override;
+
+  StatusOr<ref_ptr<Buffer>> WrapMutable(MemoryTypeBitfield memory_type,
+                                        MemoryAccessBitfield allowed_access,
+                                        BufferUsageBitfield buffer_usage,
+                                        void* data,
+                                        size_t data_length) override;
+
+ private:
+  explicit MetalDirectAllocator(id<MTLDevice> device,
+                                id<MTLCommandQueue> transfer_queue);
+
+  StatusOr<ref_ptr<MetalBuffer>> AllocateInternal(
+      MemoryTypeBitfield memory_type, BufferUsageBitfield buffer_usage,
+      MemoryAccessBitfield allowed_access, size_t allocation_size);
+
+  id<MTLDevice> metal_device_;
+  id<MTLCommandQueue> metal_transfer_queue_;
+};
+
+}  // namespace metal
+}  // namespace hal
+}  // namespace iree
+
+#endif  // IREE_HAL_METAL_METAL_DIRECT_ALLOCATOR_H_

diff --git a/iree/hal/metal/metal_direct_allocator.mm b/iree/hal/metal/metal_direct_allocator.mm
new file mode 100644
index 0000000..633bd41
--- /dev/null
+++ b/iree/hal/metal/metal_direct_allocator.mm

@@ -0,0 +1,157 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "iree/hal/metal/metal_direct_allocator.h"
+
+#include "absl/memory/memory.h"
+#include "iree/base/status.h"
+#include "iree/base/tracing.h"
+#include "iree/hal/metal/metal_buffer.h"
+
+namespace iree {
+namespace hal {
+namespace metal {
+
+namespace {
+
+// Returns the proper Metal resource storage mode given the specific MemoryType.
+MTLResourceOptions SelectMTLResourceStorageMode(MemoryType memory_type) {
+  // There are four MTLStorageMode:
+  // * Managed: The CPU and GPU may maintain separate copies of the resource, and any changes
+  //   must be explicitly synchronized.
+  // * Shared: The resource is stored in system memory and is accessible to both the CPU and
+  //   the GPU.
+  // * Private: The resource can be accessed only by the GPU.
+  // * Memoryless: The resource’s contents can be accessed only by the GPU and only exist
+  //   temporarily during a render pass.
+  // macOS has all of the above; MTLStorageModeManaged is not available on iOS.
+  //
+  // The IREE HAL is modeled after Vulkan so it's quite explicit. For buffers visible to both
+  // the host and the device, we would like to opt in with the explicit version
+  // (MTLStorageManaged) when possible because it should be more performant: "In macOS,
+  // there’s no difference in GPU performance between managed and private buffers." But for
+  // iOS, MTLStorageShared should be good given we have a unified memory model there.
+
+  if (AllBitsSet(memory_type, MemoryType::kDeviceLocal)) {
+    if (AllBitsSet(memory_type, MemoryType::kHostVisible)) {
+      // Device-local, host-visible.
+#ifdef IREE_PLATFORM_MACOS
+      return MTLResourceStorageModeManaged;
+#else
+      return MTLResourceStorageModeShared;
+#endif
+    } else {
+      // Device-local only.
+      return MTLResourceStorageModePrivate;
+    }
+  } else {
+    if (AllBitsSet(memory_type, MemoryType::kDeviceVisible)) {
+      // Host-local, device-visible.
+      return MTLResourceStorageModeShared;
+    } else {
+      // Host-local only.
+      // TODO(antiagainst): we probably want to just use HostBuffer here.
+      return MTLResourceStorageModeShared;
+    }
+  }
+}
+
+}  // namespace
+
+// static
+std::unique_ptr<MetalDirectAllocator> MetalDirectAllocator::Create(
+    id<MTLDevice> device, id<MTLCommandQueue> transfer_queue) {
+  IREE_TRACE_SCOPE0("MetalDirectAllocator::Create");
+  return absl::WrapUnique(new MetalDirectAllocator(device, transfer_queue));
+}
+
+MetalDirectAllocator::MetalDirectAllocator(id<MTLDevice> device, id<MTLCommandQueue> transfer_queue)
+    : metal_device_([device retain]), metal_transfer_queue_([transfer_queue retain]) {}
+
+MetalDirectAllocator::~MetalDirectAllocator() {
+  IREE_TRACE_SCOPE0("MetalDirectAllocator::dtor");
+  [metal_transfer_queue_ release];
+  [metal_device_ release];
+}
+
+bool MetalDirectAllocator::CanUseBufferLike(Allocator* source_allocator,
+                                            MemoryTypeBitfield memory_type,
+                                            BufferUsageBitfield buffer_usage,
+                                            BufferUsageBitfield intended_usage) const {
+  // TODO(benvanik): ensure there is a memory type that can satisfy the request.
+  return source_allocator == this;
+}
+
+bool MetalDirectAllocator::CanAllocate(MemoryTypeBitfield memory_type,
+                                       BufferUsageBitfield buffer_usage,
+                                       size_t allocation_size) const {
+  // TODO(benvanik): ensure there is a memory type that can satisfy the request.
+  return true;
+}
+
+Status MetalDirectAllocator::MakeCompatible(MemoryTypeBitfield* memory_type,
+                                            BufferUsageBitfield* buffer_usage) const {
+  // TODO(benvanik): mutate to match supported memory types.
+  return OkStatus();
+}
+
+StatusOr<ref_ptr<MetalBuffer>> MetalDirectAllocator::AllocateInternal(
+    MemoryTypeBitfield memory_type, BufferUsageBitfield buffer_usage,
+    MemoryAccessBitfield allowed_access, size_t allocation_size) {
+  IREE_TRACE_SCOPE0("MetalDirectAllocator::AllocateInternal");
+
+  MTLResourceOptions resource_options = SelectMTLResourceStorageMode(memory_type);
+
+  // IREE is more explicit than Metal: it tracks various state by itself. There is no need
+  // to incur Metal runtime overhead for hazard tracking.
+  resource_options |= MTLResourceHazardTrackingModeUntracked;
+
+  id<MTLBuffer> metal_buffer = [metal_device_ newBufferWithLength:allocation_size
+                                                          options:resource_options];  // retained
+
+  return MetalBuffer::CreateUnretained(
+      this, memory_type, allowed_access, buffer_usage, allocation_size, /*byte_offset=*/0,
+      /*byte_length=*/allocation_size, metal_buffer, metal_transfer_queue_);
+}
+
+StatusOr<ref_ptr<Buffer>> MetalDirectAllocator::Allocate(MemoryTypeBitfield memory_type,
+                                                         BufferUsageBitfield buffer_usage,
+                                                         size_t allocation_size) {
+  IREE_TRACE_SCOPE0("MetalDirectAllocator::Allocate");
+  return AllocateInternal(memory_type, buffer_usage, MemoryAccess::kAll, allocation_size);
+}
+
+StatusOr<ref_ptr<Buffer>> MetalDirectAllocator::AllocateConstant(BufferUsageBitfield buffer_usage,
+                                                                 ref_ptr<Buffer> source_buffer) {
+  IREE_TRACE_SCOPE0("MetalDirectAllocator::AllocateConstant");
+  // TODO(benvanik): import memory to avoid the copy.
+  IREE_ASSIGN_OR_RETURN(
+      auto buffer, AllocateInternal(MemoryType::kDeviceLocal | MemoryType::kHostVisible,
+                                    buffer_usage, MemoryAccess::kRead | MemoryAccess::kDiscardWrite,
+                                    source_buffer->byte_length()));
+  IREE_RETURN_IF_ERROR(buffer->CopyData(0, source_buffer.get(), 0, kWholeBuffer));
+  return buffer;
+}
+
+StatusOr<ref_ptr<Buffer>> MetalDirectAllocator::WrapMutable(MemoryTypeBitfield memory_type,
+                                                            MemoryAccessBitfield allowed_access,
+                                                            BufferUsageBitfield buffer_usage,
+                                                            void* data, size_t data_length) {
+  IREE_TRACE_SCOPE0("MetalDirectAllocator::WrapMutable");
+  return UnimplementedErrorBuilder(IREE_LOC) << "MetalDirectAllocator::WrapMutable";
+}
+
+}  // namespace metal
+}  // namespace hal
+}  // namespace iree

diff --git a/iree/hal/vulkan/vma_buffer.cc b/iree/hal/vulkan/vma_buffer.cc
index 51880d2..6496df9 100644
--- a/iree/hal/vulkan/vma_buffer.cc
+++ b/iree/hal/vulkan/vma_buffer.cc

@@ -55,21 +55,19 @@
     case 1: {
       uint8_t* data = static_cast<uint8_t*>(data_ptr);
       uint8_t value_bits = *static_cast<const uint8_t*>(pattern);
-      std::fill_n(data + byte_offset, byte_length, value_bits);
+      std::fill_n(data, byte_length, value_bits);
       break;
     }
     case 2: {
       uint16_t* data = static_cast<uint16_t*>(data_ptr);
       uint16_t value_bits = *static_cast<const uint16_t*>(pattern);
-      std::fill_n(data + byte_offset / sizeof(uint16_t),
-                  byte_length / sizeof(uint16_t), value_bits);
+      std::fill_n(data, byte_length / sizeof(uint16_t), value_bits);
       break;
     }
     case 4: {
       uint32_t* data = static_cast<uint32_t*>(data_ptr);
       uint32_t value_bits = *static_cast<const uint32_t*>(pattern);
-      std::fill_n(data + byte_offset / sizeof(uint32_t),
-                  byte_length / sizeof(uint32_t), value_bits);
+      std::fill_n(data, byte_length / sizeof(uint32_t), value_bits);
       break;
     }
     default:
@@ -111,8 +109,8 @@
                         MapMemory<uint8_t>(MemoryAccess::kDiscardWrite,
                                            target_offset, data_length));
   CHECK_EQ(data_length, target_mapping.size());
-  std::memcpy(target_mapping.mutable_data() + target_offset,
-              source_mapping.data(), data_length);
+  std::memcpy(target_mapping.mutable_data(), source_mapping.data(),
+              data_length);
   return OkStatus();
 }
 

diff --git a/iree/test/e2e/xla_ops/BUILD b/iree/test/e2e/xla_ops/BUILD
index 66a9fcf..fbcd1d2 100644
--- a/iree/test/e2e/xla_ops/BUILD
+++ b/iree/test/e2e/xla_ops/BUILD

@@ -40,23 +40,20 @@
         "abs.mlir",
         "add.mlir",
         "batch_norm_inference.mlir",
-        "broadcast_add.mlir",
         "broadcast.mlir",
+        "broadcast_add.mlir",
         "broadcast_in_dim.mlir",
         "clamp.mlir",
         "compare.mlir",
-        "convolution.mlir",
-        "convert.mlir",
         "concatenate.mlir",
         "constant.mlir",
+        "convert.mlir",
+        "convolution.mlir",
         "cosine.mlir",
         "divide.mlir",
         "dot.mlir",
         "exponential.mlir",
-
-        # TODO(#1696): Enable after standard dialect can support floor
-        # operation. Lowering from XLA -> linalg should be easy fix.
-        # "floor.mlir",
+        "floor.mlir",
         "gather.mlir",
         "gather_concat.mlir",
         "iota.mlir",
@@ -107,6 +104,7 @@
         "divide.mlir",
         "dot.mlir",
         "exponential.mlir",
+        "floor.mlir",
         "gather.mlir",
         "gather_concat.mlir",
         "iota.mlir",

diff --git a/iree/test/e2e/xla_ops/CMakeLists.txt b/iree/test/e2e/xla_ops/CMakeLists.txt
index 0303432..c44290a 100644
--- a/iree/test/e2e/xla_ops/CMakeLists.txt
+++ b/iree/test/e2e/xla_ops/CMakeLists.txt

@@ -46,6 +46,7 @@
     "divide.mlir"
     "dot.mlir"
     "exponential.mlir"
+    "floor.mlir"
     "gather.mlir"
     "gather_concat.mlir"
     "iota.mlir"
@@ -96,6 +97,7 @@
     "divide.mlir"
     "dot.mlir"
     "exponential.mlir"
+    "floor.mlir"
     "gather.mlir"
     "gather_concat.mlir"
     "iota.mlir"

diff --git a/iree/vm/native_module_test.h b/iree/vm/native_module_test.h
index ce08749..35de2d5 100644
--- a/iree/vm/native_module_test.h
+++ b/iree/vm/native_module_test.h

@@ -187,6 +187,7 @@
   IREE_RETURN_IF_ERROR(
       iree_allocator_malloc(allocator, sizeof(*state), (void**)&state));
   memset(state, 0, sizeof(*state));
+  state->allocator = allocator;
   *out_module_state = (iree_vm_module_state_t*)state;
   return iree_ok_status();
 }
@@ -281,6 +282,7 @@
   IREE_RETURN_IF_ERROR(
       iree_allocator_malloc(allocator, sizeof(*module), (void**)&module));
   memset(module, 0, sizeof(*module));
+  module->allocator = allocator;
 
   // Resolve types used by the module once so that we can share it across all
   // instances of the module.