mirror of
https://github.com/ROCm/jax.git
synced 2025-04-17 04:16:07 +00:00
Merge pull request #16491 from hawkinsp:iree
PiperOrigin-RevId: 547184602
This commit is contained in:
jax authors
commit
06a7ea91b0
10
jax/BUILD
10
jax/BUILD
@ -345,15 +345,6 @@ pytype_strict_library(
|
||||
],
|
||||
)
|
||||
|
||||
pytype_library(
|
||||
name = "iree",
|
||||
srcs = ["_src/iree.py"],
|
||||
deps = [
|
||||
":config",
|
||||
"//jax/_src/lib",
|
||||
] + py_deps("numpy"),
|
||||
)
|
||||
|
||||
pytype_library(
|
||||
name = "lax_reference",
|
||||
srcs = ["_src/lax_reference.py"],
|
||||
@ -612,7 +603,6 @@ pytype_strict_library(
|
||||
deps = [
|
||||
":cloud_tpu_init",
|
||||
":config",
|
||||
":iree",
|
||||
":traceback_util",
|
||||
":util",
|
||||
"//jax/_src/lib",
|
||||
|
||||
214
jax/_src/iree.py
214
jax/_src/iree.py
@ -1,214 +0,0 @@
|
||||
# Copyright 2021 The JAX Authors.
|
||||
#
|
||||
# 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.
|
||||
|
||||
"""Experimental IREE backend for JAX.
|
||||
|
||||
This backend is quite incomplete, but exists to allow experimenting with
|
||||
using IREE to compile and run JAX computations instead of XLA.
|
||||
"""
|
||||
|
||||
# pytype: skip-file
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import platform
|
||||
from typing import Any, Sequence, Optional
|
||||
|
||||
import iree.compiler
|
||||
import iree.runtime
|
||||
|
||||
from jax._src.config import flags
|
||||
from jax._src.lib import xla_client
|
||||
import numpy as np
|
||||
|
||||
FLAGS = flags.FLAGS
|
||||
|
||||
|
||||
flags.DEFINE_string(
|
||||
'jax_iree_backend', os.getenv('JAX_IREE_BACKEND', 'cpu'),
|
||||
'IREE compiler backend to use.')
|
||||
|
||||
iree_compiler_map = {
|
||||
"cpu" : "llvm-cpu",
|
||||
"cuda" : "cuda",
|
||||
"vmvx" : "vmvx",
|
||||
"vulkan" : "vulkan-spirv"
|
||||
}
|
||||
|
||||
iree_runtime_map = {
|
||||
"cpu" : "local-task",
|
||||
"cuda" : "cuda",
|
||||
"vmvx" : "local-task",
|
||||
"vulkan" : "vulkan"
|
||||
}
|
||||
|
||||
class IreeDevice:
|
||||
|
||||
def __init__(self, client):
|
||||
self.id = 0
|
||||
self.host_id = 0
|
||||
self.process_index = 0
|
||||
self.platform = "iree"
|
||||
self.device_kind = "IREE device"
|
||||
self.client = client
|
||||
|
||||
def __str__(self) -> str:
|
||||
return "IreeDevice"
|
||||
|
||||
def transfer_to_infeed(self, literal: Any):
|
||||
raise NotImplementedError("transfer_to_infeed")
|
||||
|
||||
def transfer_from_outfeed(self, shape: xla_client.Shape):
|
||||
raise NotImplementedError("transfer_to_outfeed")
|
||||
|
||||
def live_buffers(self) -> list[IreeBuffer]:
|
||||
raise NotImplementedError("live_buffers")
|
||||
|
||||
|
||||
class IreeBuffer:
|
||||
|
||||
def __init__(self, client, device, buffer):
|
||||
self.client = client
|
||||
self._device = device
|
||||
assert device is not None
|
||||
self._buffer = buffer
|
||||
|
||||
def copy_to_device(self, device):
|
||||
return self
|
||||
|
||||
def __array__(self, dtype=None, context=None):
|
||||
return np.asarray(self._buffer)
|
||||
|
||||
def to_iree(self):
|
||||
return self._buffer
|
||||
|
||||
def platform(self):
|
||||
return self.client.platform
|
||||
|
||||
def device(self):
|
||||
return self._device
|
||||
|
||||
def block_until_ready(self) -> IreeBuffer:
|
||||
return self # no async
|
||||
|
||||
# overrides repr on base class which expects _value and aval attributes
|
||||
def __repr__(self): return f'IreeBuffer({np.asarray(self)})'
|
||||
|
||||
@property
|
||||
def _value(self):
|
||||
return np.asarray(self)
|
||||
|
||||
def copy_to_host_async(self):
|
||||
return self
|
||||
|
||||
class IreeExecutable:
|
||||
|
||||
def __init__(self, client, devices, module_object, function_name):
|
||||
self.client = client
|
||||
self.traceback = None
|
||||
self.fingerprint = None
|
||||
self._devices = devices
|
||||
self.module_object = module_object
|
||||
self.function_name = function_name
|
||||
|
||||
def local_devices(self) -> list[IreeDevice]:
|
||||
return self._devices
|
||||
|
||||
def execute(self, arguments: Sequence[IreeBuffer]) -> list[IreeBuffer]:
|
||||
inputs = [arg.to_iree() for arg in arguments]
|
||||
outputs = self.module_object[self.function_name](*inputs)
|
||||
# TODO(phawkins): Have a way to just have it always return the list,
|
||||
# regardless of arity.
|
||||
if not isinstance(outputs, list) and not isinstance(outputs, tuple):
|
||||
outputs = [outputs]
|
||||
return [
|
||||
IreeBuffer(self.client, self._devices[0], output) for output in outputs
|
||||
]
|
||||
|
||||
|
||||
class IreeClient:
|
||||
|
||||
def __init__(self,
|
||||
*,
|
||||
iree_backend: Optional[str] = None):
|
||||
self.platform = "iree"
|
||||
self.platform_version = "0.0.1"
|
||||
self.runtime_type = "iree"
|
||||
self.iree_backend = (FLAGS.jax_iree_backend if iree_backend is None
|
||||
else iree_backend)
|
||||
self.compiler_driver = iree_compiler_map[self.iree_backend]
|
||||
self.runtime_driver = iree_runtime_map[self.iree_backend]
|
||||
self.iree_config = iree.runtime.system_api.Config(self.runtime_driver)
|
||||
self._devices = [IreeDevice(self)]
|
||||
|
||||
def process_index(self) -> int:
|
||||
return 0
|
||||
|
||||
def device_count(self) -> int:
|
||||
return len(self._devices)
|
||||
|
||||
def devices(self) -> list[IreeDevice]:
|
||||
return self._devices
|
||||
|
||||
def local_devices(self) -> list[IreeDevice]:
|
||||
return self._devices
|
||||
|
||||
def local_device_count(self) -> int:
|
||||
return len(self._devices)
|
||||
|
||||
def get_default_device_assignment(
|
||||
self,
|
||||
num_replicas: int) -> list[IreeDevice]:
|
||||
if num_replicas != 1:
|
||||
raise NotImplementedError("Only single-device computations implemented")
|
||||
return [self._devices[0]]
|
||||
|
||||
|
||||
def compile(self, computation: str,
|
||||
compile_options: xla_client.CompileOptions) -> IreeExecutable:
|
||||
del compile_options # Ignored.
|
||||
extra_args = []
|
||||
# extra_args=["--mlir-print-ir-after-all"]
|
||||
if platform.system() == "Darwin" and platform.machine() == "arm64":
|
||||
extra_args += ["--iree-llvm-target-triple=arm64-apple-darwin21.5.0"]
|
||||
iree_binary = iree.compiler.compile_str(
|
||||
computation, target_backends=[self.compiler_driver], input_type="mhlo",
|
||||
# extended_diagnostics=True,
|
||||
extra_args=extra_args,
|
||||
)
|
||||
# Load it into the runtime.
|
||||
vm_module = iree.runtime.VmModule.from_flatbuffer(
|
||||
self.iree_config.vm_instance, iree_binary)
|
||||
module_object = iree.runtime.load_vm_module(vm_module, self.iree_config)
|
||||
return IreeExecutable(self, self._devices, module_object, "main")
|
||||
|
||||
def buffer_from_pyval(
|
||||
self,
|
||||
argument: Any,
|
||||
device: Optional[IreeDevice],
|
||||
force_copy: bool = True,
|
||||
host_buffer_semantics: xla_client.HostBufferSemantics = xla_client
|
||||
.HostBufferSemantics.ZERO_COPY
|
||||
) -> IreeBuffer:
|
||||
# TODO(phawkins): IREE's python API will accept a numpy array directly but
|
||||
# may want to explicitly construct a lower level BufferView to avoid copies.
|
||||
if device is None:
|
||||
assert type(argument) is np.ndarray
|
||||
device = self._devices[0]
|
||||
return IreeBuffer(self, device, np.array(argument, copy=True))
|
||||
|
||||
|
||||
def iree_client_factory():
|
||||
return IreeClient()
|
||||
@ -43,13 +43,6 @@ from jax._src.lib import xla_extension_version
|
||||
from jax._src import traceback_util
|
||||
from jax._src import util
|
||||
|
||||
iree: Optional[Any]
|
||||
|
||||
try:
|
||||
import jax._src.iree as iree # type: ignore
|
||||
except (ModuleNotFoundError, ImportError):
|
||||
iree = None
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
jax_plugins: Optional[Any]
|
||||
@ -522,10 +515,6 @@ discover_pjrt_plugins()
|
||||
# windows).
|
||||
register_pjrt_plugin_factories_from_env()
|
||||
|
||||
if iree is not None:
|
||||
register_backend_factory("iree", iree.iree_client_factory, priority=-100,
|
||||
fail_quietly=True)
|
||||
|
||||
_platform_aliases = {
|
||||
"cuda": "gpu",
|
||||
"rocm": "gpu",
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user