Status of int8 dot/conv with XLA+CUDA
Hi,
I would like to contribute to this organization. But, I don't know how to. Can you suggest me what to learn ?. I have learnt about deep learning. And I have tried reading as many as issue's. Even though I have not been able to understand how to solve these issue's. So please let me know what I should learn first.   Hi @SuryanarayanaY, can you please look into this? Thanks. Hi @jmc128 ,

The issue [49140](https://github.com/tensorflow/tensorflow/issues/49140) still open in our internal ticket also.I will check and let you know if there is any update.

I would like to know how you have implemented the dot product with XLA. Is it annotated with `@tf.function(jit_compile=True) `? 

Also request you to please submit your issue in the standard format attached [here](https://github.com/tensorflow/tensorflow/issues/new?assignees=&labels=&template=tensorflow_issue_template.yaml).

Thanks! Hi @SuryanarayanaY, I've updated the post with a code example and system information. @jmc128 , 

Thanks for the implementation code. I have tested the same in colab and it seems the dot(tf.matmul) is success but tf.nn.Conv2D still not supported as per [gist](https://colab.research.google.com/gist/SuryanarayanaY/5431db781950603aa2995ad8612a1b2b/59530.ipynb).

Could you please submit minimal code snippet for performance comparison as you observed on tf.matmul? I've attached a script which I've been using to benchmark the throughput of linear layers using int8 and fp16 precision.  

On an A100, it produces the following results :
```
TOPS/TFLOPS
==========================
dim    |     int8     fp16
--------------------------
256    |     0.29     0.34
512    |     2.33     2.79
1024   |    16.51    19.47
2048   |    96.75   102.14
4096   |   261.43   205.79
8192   |   354.98   263.79
16384  |   367.82   272.62
==========================
```
My issue regarding performance is that int8 throughput should be approaching 2x the throughput of fp16.  In my experiments with TensorRT, int8 attains 2x throughput vs fp16 for dims >= 2048, and for dims between 256 and 2048 throughput is still at least 1.7x.  I was hoping someone familiar with the internals of XLA could comment on whether this level of performance can be attained with XLA.  Is there a particular way to express this computation in TensorFlow which enables XLA to output better-optimized kernels, or is this simply not possible at all with XLA at the moment?

```python
import time
import tensorflow as tf


def int8_linear_layer(x, params):
    # read in x and w as pre-quantized int8 tensors
    y = tf.matmul(x, params["w"], output_type=tf.int32)

    # add bias and apply activation in fp32
    y = tf.cast(y, tf.float32)
    y = y + params["b"]
    y = tf.nn.relu(y)

    # quantize and store output as int8
    y = tf.round(y / params["s"])
    y = tf.clip_by_value(y, -128, 127)
    y = tf.cast(y, tf.int8)
    return y


def fp_linear_layer(x, params):
    return tf.nn.relu(tf.matmul(x, params["w"]) + params["b"])


def linear_layer(x, params):
    if params["w"].dtype == tf.int8:
        return int8_linear_layer(x, params)
    else:
        return fp_linear_layer(x, params)


@tf.function(jit_compile=True)
def network(x, network_params):
    for layer_params in network_params:
        x = linear_layer(x, layer_params)
    return x[0, 0]


def benchmark(dim, dtype, n_layers=32, iters=200):
    """
    Benchmarks the throughput of a linear layer (matmul + bias + act (+ quantization))

    To drown out data transfer costs and other overheads, each run performs a
    forward pass of a network consisting of `n_layers` of linear layers in series

    The full network is run for `iters` iterations, and the last 10% of iterations are
    used to compute the throughput in TFLOPS/TOPS
    """

    # create input tensor for network
    if dtype == tf.int8:
        x = tf.random.uniform([dim, dim], minval=-127, maxval=127, dtype=tf.int32)
        x = tf.cast(x, tf.int8)
    else:
        x = tf.random.uniform([dim, dim], dtype=dtype)

    # create weight + bias (+ quant scale) tensors for each layer
    network_params = []
    for i in range(n_layers):
        if dtype == tf.int8:
            w = tf.random.uniform([dim, dim], minval=-127, maxval=127, dtype=tf.int32)
            w = tf.cast(w, tf.int8)
            b = tf.random.normal([dim], dtype=tf.float32)
            s = tf.random.normal([], dtype=tf.float32)
            layer_params = {"w": w, "b": b, "s": s}
        else:
            w = tf.random.uniform([dim, dim], dtype=dtype)
            b = tf.random.uniform([dim], dtype=dtype)
            layer_params = {"w": w, "b": b}
        network_params.append(layer_params)

    times = []
    for i in range(iters):
        t0 = time.time()
        y = network(x, network_params).numpy()
        elapsed_time = time.time() - t0  # in ms
        times.append(elapsed_time)

    times = times[-(iters // 10) :]  # discard warmup iters
    avg_time = sum(times) / len(times)
    tflops = (n_layers * 2 * dim ** 3) / avg_time * 1e-12
    return tflops


if __name__ == "__main__":
    benchmark(256, tf.float16)  # not used, just triggers tensorflow/xla initial log dump

    print("\nTOPS/TFLOPS")
    cols = f"{'dim':<6} | {'int8':>8} {'fp16':>8}"
    print("=" * len(cols) + "\n" + cols + "\n" + "-" * len(cols))
    for dim in [256, 512, 1024, 2048, 4096, 8192]:
        int8_time = benchmark(dim, tf.int8)
        fp16_time = benchmark(dim, tf.float16)
        print(f"{dim:<6} | {int8_time:>8.2f} {fp16_time:>8.2f}")
    print("=" * len(cols))
``` I have replicated the reported behaviour and attached the [gist](https://colab.research.google.com/gist/SuryanarayanaY/d8c1fdccc447aed3eac72a8af37c47cb/59530_r1.ipynb) here.

@sachinprasadhs Could you please have a look into this? Thanks! @Reed Wanderman-Milne ***@***.***> could you take a look?

On Thu, Feb 16, 2023 at 1:01 AM Sachin Prasad ***@***.***>
wrote:

> Assigned #59530 <https://github.com/tensorflow/tensorflow/issues/59530>
> to @cheshire <https://github.com/cheshire>.
>
> —
> Reply to this email directly, view it on GitHub
> <https://github.com/tensorflow/tensorflow/issues/59530#event-8531672393>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AACVGH4Y6JCZ6EMSU3XNYZTWXVU4TANCNFSM6AAAAAAUPOBSFM>
> .
> You are receiving this because you were assigned.Message ID:
> ***@***.***>
>

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