mirror of
https://github.com/ROCm/jax.git
synced 2025-04-17 04:16:07 +00:00
a1480c454e
As discussed over the last few months, it is desirable to migrate JAX from producing MHLO to producing StableHLO, and this CL makes this happen. More specifically: 1) MLIR lowerings now produce StableHLO ops instead of MHLO ops. 2) Fallback lowerings now produce StableHLO ops as well. 3) Occurrences of "MHLO" in prose have been changed to "StableHLO", unless the documents are immutable (changelog, JEPs). From time to time, it might be useful to produce MHLO directly, so MHLO is not going away and is still within arm's reach (although compatibility guarantees will only be provided for StableHLO and not for MHLO): a) `from jax._src.lib.mlir.dialects import mhlo` still does the same thing. b) `XlaLowering.mhlo()` is available as well, but its implementation has changed - it calls `stablehlo-legalize-to-hlo` underneath. c) `Lowering.as_text()/compiler_ir()` still support `dialect="mhlo"`, but the default has changed to "stablehlo". d) We're still using `mhlo.is_same_data_across_replicas` and `mhlo.sharding` because StableHLO currently lacks comparable functionality. https://github.com/openxla/stablehlo/issues/744 tracks the corresponding work, but it is not a blocker - we can use these attributes with StableHLO without any issues. PiperOrigin-RevId: 497978733
126 lines
4.7 KiB
Python
126 lines
4.7 KiB
Python
# Copyright 2022 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.
|
|
|
|
import jaxlib.mlir.ir as ir
|
|
import jaxlib.mlir.dialects.stablehlo as hlo
|
|
|
|
import numpy as np
|
|
|
|
from jaxlib import xla_client
|
|
|
|
try:
|
|
from .cuda import _rnn as _rnn
|
|
for _name, _value in _rnn.registrations().items():
|
|
xla_client.register_custom_call_target(_name, _value, platform='CUDA')
|
|
except ImportError:
|
|
_rnn = None
|
|
|
|
if _rnn:
|
|
compute_rnn_workspace_reserve_space_sizes = _rnn.compute_rnn_workspace_reserve_space_sizes
|
|
|
|
|
|
def cudnn_rnn_lowering(ctx, input, h_0, c_0, weights, seq_lengths, *,
|
|
input_size: int, hidden_size: int, num_layers: int,
|
|
dropout: bool, bidirectional: bool):
|
|
"""CuDnn RNN."""
|
|
out_dtype = ctx.avals_out[0].dtype
|
|
if out_dtype == np.float32:
|
|
out_type = ir.F32Type.get()
|
|
elif out_dtype == np.float64:
|
|
out_type = ir.F64Type.get()
|
|
elif out_dtype == np.complex64:
|
|
out_type = ir.ComplexType.get(ir.F32Type.get())
|
|
elif out_dtype == np.complex128:
|
|
out_type = ir.ComplexType.get(ir.F64Type.get())
|
|
else:
|
|
raise ValueError(f'Unknown output type {out_dtype}')
|
|
|
|
output_type = ir.RankedTensorType.get(ctx.avals_out[0].shape, out_type)
|
|
batch_size = ctx.avals_in[0].shape[0]
|
|
max_seq_length = ctx.avals_in[0].shape[1]
|
|
workspace_shape = ctx.avals_out[3].shape
|
|
reserve_space_shape = ctx.avals_out[4].shape
|
|
workspace_type = ir.RankedTensorType.get(workspace_shape, ir.F32Type.get())
|
|
reserve_space_type = ir.RankedTensorType.get(reserve_space_shape,
|
|
ir.F32Type.get())
|
|
opaque = _rnn.build_rnn_descriptor(input_size, hidden_size, num_layers,
|
|
batch_size, max_seq_length, dropout,
|
|
bidirectional, workspace_shape[0],
|
|
reserve_space_shape[0])
|
|
|
|
i32_type = ir.IntegerType.get_signless(32)
|
|
|
|
out = hlo.CustomCallOp(
|
|
[
|
|
ir.TupleType.get_tuple([
|
|
output_type, h_0.type, c_0.type, workspace_type,
|
|
reserve_space_type
|
|
])
|
|
],
|
|
[input, h_0, c_0, weights, seq_lengths],
|
|
call_target_name=ir.StringAttr.get('cudnn_rnn'),
|
|
has_side_effect=ir.BoolAttr.get(False),
|
|
backend_config=ir.StringAttr.get(opaque),
|
|
api_version=ir.IntegerAttr.get(i32_type, 2),
|
|
called_computations=ir.ArrayAttr.get([]),
|
|
)
|
|
return [
|
|
hlo.GetTupleElementOp(out, ir.IntegerAttr.get(i32_type, i)).result
|
|
for i in range(5)
|
|
]
|
|
|
|
|
|
def _hlo_zeros_f32(shape):
|
|
return hlo.ConstantOp(
|
|
ir.DenseElementsAttr.get(
|
|
np.zeros(shape, dtype=np.float32), type=ir.F32Type.get())).result
|
|
|
|
|
|
def cudnn_rnn_bwd_lowering(ctx, dy, dhn, dcn, x, h0, c0, w, y, workspace,
|
|
reserve_space, seq_lengths, *, input_size: int,
|
|
hidden_size: int, num_layers: int, dropout: bool,
|
|
bidirectional: bool):
|
|
"""CuDnn RNN Backward pass."""
|
|
batch_size = ctx.avals_in[3].shape[0]
|
|
max_seq_length = ctx.avals_in[3].shape[1]
|
|
workspace_shape = ctx.avals_in[8].shape
|
|
reserve_space_shape = ctx.avals_in[9].shape
|
|
opaque = _rnn.build_rnn_descriptor(input_size, hidden_size, num_layers,
|
|
batch_size, max_seq_length, dropout,
|
|
bidirectional, workspace_shape[0],
|
|
reserve_space_shape[0])
|
|
|
|
i32_type = ir.IntegerType.get_signless(32)
|
|
zeroed_dw = _hlo_zeros_f32(ctx.avals_out[3].shape)
|
|
out = hlo.CustomCallOp(
|
|
[ir.TupleType.get_tuple([x.type, h0.type, c0.type, w.type])], [
|
|
dy, dhn, dcn, x, h0, c0, w, y, workspace, reserve_space, zeroed_dw,
|
|
seq_lengths
|
|
],
|
|
call_target_name=ir.StringAttr.get('cudnn_rnn_bwd'),
|
|
has_side_effect=ir.BoolAttr.get(False),
|
|
backend_config=ir.StringAttr.get(opaque),
|
|
api_version=ir.IntegerAttr.get(i32_type, 2),
|
|
called_computations=ir.ArrayAttr.get([]),
|
|
output_operand_aliases=ir.ArrayAttr.get([
|
|
hlo.OutputOperandAlias.get(
|
|
output_tuple_indices=[3],
|
|
operand_index=10,
|
|
operand_tuple_indices=[])
|
|
]))
|
|
return [
|
|
hlo.GetTupleElementOp(out, ir.IntegerAttr.get(i32_type, i)).result
|
|
for i in range(4)
|
|
]
|