Repository: tensorflow
Issue: Status of int8 dot/conv with XLA+CUDA
Labels: stat:awaiting tensorflower,type:bug,comp:gpu,comp:xla,type:performance

Are tensorcore-accelerated int8 dot/convs with XLA accessible via tensorflow APIs? https://github.com/tensorflow/tensorflow/pull/30771 suggests that int8 convs are supported, and https://github.com/tensorflow/tensorflow/issues/49140 also suggest that this functionality exists. However, the tensorflow convolution operations don't allow int8 arguments. The matmul operator allows int8 x int8 -> int32 via the output_type argument, and I've successfully compiled with XLA using this option. However, when doing this, I'm not getting the throughput that I'd expect with int8 tensorcores. I was expecting something approaching 2x throughput compared to fp16 for large matmuls ([8192, 8192] inputs). Admittedly my benchmarking is not very rigorous, but I'd be curious to know if my expectation of ~2x throughput is reasonable and whether XLA has the functionality to facilitate it.

Here's a minimal implementation of an XLA-compiled int8 matmul:
```python
@tf.function(jit_compile=True)
def int8_matmul(x, w):
    return tf.matmul(x, w, output_type=tf.int32)


x = tf.cast(tf.random.uniform([2048, 2048], minval=-127, maxval=127, dtype=tf.int32), tf.int8)
w = tf.cast(tf.random.uniform([2048, 2048], minval=-127, maxval=127, dtype=tf.int32), tf.int8)
y = int8_matmul(x, w)

print(int8_matmul.experimental_get_compiler_ir(x, w)(stage="optimized_hlo"))
# HloModule a_inference_int8_matmul_17__.10, alias_passthrough_params=true, entry_computation_layout={(s8[2048,2048]{1,0},s8[2048,2048]{1,0})->s32[2048,2048]{1,0}}

# ENTRY %a_inference_int8_matmul_17__.10 (arg0.1: s8[2048,2048], arg1.2: s8[2048,2048]) -> s32[2048,2048] {
#   %arg0.1 = s8[2048,2048]{1,0} parameter(0), parameter_replication={false}, metadata={op_name="XLA_Args"}
#   %arg1.2 = s8[2048,2048]{1,0} parameter(1), parameter_replication={false}, metadata={op_name="XLA_Args"}
#   %copy = s8[2048,2048]{0,1} copy(s8[2048,2048]{1,0} %arg1.2), metadata={op_name="XLA_Args"}
#   ROOT %cublas-gemm.1 = s32[2048,2048]{1,0} custom-call(s8[2048,2048]{1,0} %arg0.1, s8[2048,2048]{0,1} %copy), custom_call_target="__cublas$gemm", metadata={op_type="BatchMatMulV3" op_name="MatMul" source_file="int8_xla.py" source_line=22}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":0,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"1\"],\"rhs_contracting_dimensions\":[\"0\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"precision_config\":{\"operand_precision\":[\"DEFAULT\",\"DEFAULT\"]},\"epilogue\":\"DEFAULT\"}"
# }
```

And here's my attempt at implementing a 'canonical' quantized linear layer.  I think this is a fairly standard sequence of operations for a quantized layer, and I believe TensorRT would be able to fuse this entire operation.
```python
@tf.function(jit_compile=True)
def int8_linear_layer(x, w, b, s):
    # read in x and w as pre-quantized int8 tensors
    y = tf.matmul(x, w, output_type=tf.int32)

    # add bias and apply activation in fp32
    y = tf.cast(y, tf.float32)
    y = y + b
    y = tf.nn.relu(y)

    # quantize and store output as int8
    y = tf.round(y / s)
    y = tf.clip_by_value(y, -128, 127)
    y = tf.cast(y, tf.int8)
    return y


x = tf.cast(tf.random.uniform([2048, 2048], minval=-127, maxval=127, dtype=tf.int32), tf.int8)
w = tf.cast(tf.random.uniform([2048, 2048], minval=-127, maxval=127, dtype=tf.int32), tf.int8)
b = tf.random.normal([2048], dtype=tf.float32)  # bias
s = tf.random.normal([], dtype=tf.float32)  # per-tensor quantization scale for output activation
y = int8_linear_layer(x, w, b, s)

print(int8_linear_layer.experimental_get_compiler_ir(x, w, b, s)(stage="optimized_hlo"))
# HloModule a_inference_int8_linear_layer_65__.32, alias_passthrough_params=true, entry_computation_layout={(s8[2048,2048]{1,0},s8[2048,2048]{1,0},f32[2048]{0},f32[])->s8[2048,2048]{1,0}}

# %fused_computation (param_0.2: f32[], param_1.4: f32[2048], param_2.7: s32[2048,2048]) -> s8[2048,2048] {
#   %constant_2 = f32[] constant(-128), metadata={op_type="Maximum" op_name="clip_by_value" source_file="int8_xla.py" source_line=52}
#   %broadcast.4 = f32[2048,2048]{1,0} broadcast(f32[] %constant_2), dimensions={}, metadata={op_type="Maximum" op_name="clip_by_value" source_file="int8_xla.py" source_line=52}
#   %param_2.7 = s32[2048,2048]{1,0} parameter(2)
#   %convert.1 = f32[2048,2048]{1,0} convert(s32[2048,2048]{1,0} %param_2.7), metadata={op_type="Cast" op_name="Cast" source_file="int8_xla.py" source_line=46}
#   %param_1.4 = f32[2048]{0} parameter(1)
#   %broadcast.3 = f32[2048,2048]{1,0} broadcast(f32[2048]{0} %param_1.4), dimensions={1}, metadata={op_type="AddV2" op_name="add" source_file="int8_xla.py" source_line=47}
#   %add.0 = f32[2048,2048]{1,0} add(f32[2048,2048]{1,0} %convert.1, f32[2048,2048]{1,0} %broadcast.3), metadata={op_type="AddV2" op_name="add" source_file="int8_xla.py" source_line=47}
#   %constant_1 = f32[] constant(0), metadata={op_type="Relu" op_name="Relu" source_file="int8_xla.py" source_line=48}
#   %broadcast.2 = f32[2048,2048]{1,0} broadcast(f32[] %constant_1), dimensions={}, metadata={op_type="Relu" op_name="Relu"}
#   %maximum.0 = f32[2048,2048]{1,0} maximum(f32[2048,2048]{1,0} %add.0, f32[2048,2048]{1,0} %broadcast.2), metadata={op_type="Relu" op_name="Relu"}
#   %param_0.2 = f32[] parameter(0)
#   %broadcast.1 = f32[2048,2048]{1,0} broadcast(f32[] %param_0.2), dimensions={}, metadata={op_type="RealDiv" op_name="truediv" source_file="int8_xla.py" source_line=51}
#   %divide.0 = f32[2048,2048]{1,0} divide(f32[2048,2048]{1,0} %maximum.0, f32[2048,2048]{1,0} %broadcast.1), metadata={op_type="RealDiv" op_name="truediv" source_file="int8_xla.py" source_line=51}
#   %round-nearest-even.0 = f32[2048,2048]{1,0} round-nearest-even(f32[2048,2048]{1,0} %divide.0), metadata={op_type="Round" op_name="Round" source_file="int8_xla.py" source_line=51}
#   %constant_0 = f32[] constant(127), metadata={op_type="Minimum" op_name="clip_by_value/Minimum" source_file="int8_xla.py" source_line=52}
#   %broadcast.0 = f32[2048,2048]{1,0} broadcast(f32[] %constant_0), dimensions={}, metadata={op_type="Minimum" op_name="clip_by_value/Minimum" source_file="int8_xla.py" source_line=52}
#   %clamp.1 = f32[2048,2048]{1,0} clamp(f32[2048,2048]{1,0} %broadcast.4, f32[2048,2048]{1,0} %round-nearest-even.0, f32[2048,2048]{1,0} %broadcast.0), metadata={op_type="Maximum" op_name="clip_by_value" source_file="int8_xla.py" source_line=52}
#   ROOT %convert.0 = s8[2048,2048]{1,0} convert(f32[2048,2048]{1,0} %clamp.1), metadata={op_type="Cast" op_name="Cast_1" source_file="int8_xla.py" source_line=53}
# }

# ENTRY %a_inference_int8_linear_layer_65__.32 (arg0.1: s8[2048,2048], arg1.2: s8[2048,2048], arg2.3: f32[2048], arg3.4: f32[]) -> s8[2048,2048] {
#   %arg3.4 = f32[] parameter(3), parameter_replication={false}, metadata={op_name="XLA_Args"}
#   %arg2.3 = f32[2048]{0} parameter(2), parameter_replication={false}, metadata={op_name="XLA_Args"}
#   %arg0.1 = s8[2048,2048]{1,0} parameter(0), parameter_replication={false}, metadata={op_name="XLA_Args"}
#   %arg1.2 = s8[2048,2048]{1,0} parameter(1), parameter_replication={false}, metadata={op_name="XLA_Args"}
#   %copy = s8[2048,2048]{0,1} copy(s8[2048,2048]{1,0} %arg1.2), metadata={op_name="XLA_Args"}
#   %cublas-gemm.1 = s32[2048,2048]{1,0} custom-call(s8[2048,2048]{1,0} %arg0.1, s8[2048,2048]{0,1} %copy), custom_call_target="__cublas$gemm", metadata={op_type="BatchMatMulV3" op_name="MatMul" source_file="int8_xla.py" source_line=43}, backend_config="{\"alpha_real\":1,\"alpha_imag\":0,\"beta\":0,\"dot_dimension_numbers\":{\"lhs_contracting_dimensions\":[\"1\"],\"rhs_contracting_dimensions\":[\"0\"],\"lhs_batch_dimensions\":[],\"rhs_batch_dimensions\":[]},\"precision_config\":{\"operand_precision\":[\"DEFAULT\",\"DEFAULT\"]},\"epilogue\":\"DEFAULT\"}"
#   ROOT %fusion = s8[2048,2048]{1,0} fusion(f32[] %arg3.4, f32[2048]{0} %arg2.3, s32[2048,2048]{1,0} %cublas-gemm.1), kind=kLoop, calls=%fused_computation, metadata={op_type="Cast" op_name="Cast_1" source_file="int8_xla.py" source_line=53}
# }
```

System info: Ubuntu 20.04.5 LTS, TF 2.11.0 via pip, A100 GPU, CUDA Version 12.0