🏆 ยอดนิยม: openxla/stablehlo/pull/

sdasgup3 · 2023-05-25T17:22:13Z

The PR proposes the specification of quantized ReduceOp, ReduceWindowOp, and SelectAndScatterOp.

For ReduceOp (and the other proposed ops using ReduceOp in their specification ), the specification allows different quantized element types for inputs (or init_values) and the results . The reason is to allow the accumulation to happen in higher storage_type. However, the quantized element types for inputs and init_values are the same. This behavior is same as other reduction ops like dot_general.

The following shows an example usage of reduce op:

%reduce = "stablehlo.reduce"(%input, %init_value) ({
    ^bb0(%elem1: tensor<!quant.uniform<ui8:f32, input_scale:input_zp>>,
         %elem2: tensor<!quant.uniform<ui8:f32, input_scale:input_zp>>):

        %elem1_rescaled = "stablehlo.uniform_quantize"(%elem1)
            : (tensor<!quant.uniform<ui8:f32, input_scale:input_zp>>)
            -> tensor<!quant.uniform<i32:f32, output_scale:output_zp>>

         %elem1_rescaled = "stablehlo.uniform_quantize"(%elem2)
            : (tensor<!quant.uniform<ui8:f32, input_scale:input_zp>>)
            -> tensor<!quant.uniform<i32:f32, output_scale:output_zp>>


        %add = "stablehlo.add"(%elem1_rescaled, %%elem2_rescaled)
            : (tensor<!quant.uniform<i32:f32, output_scale:output_zp>>,
               tensor<!quant.uniform<i32:f32, output_scale:output_zp>>)
            -> tensor<!quant.uniform<i32:f32, output_scale:output_zp>>

        "stablehlo.return"(%add)
            : (tensor<!quant.uniform<i32:f32, output_scale:output_zp>>) -> ()
    }) {
        dimensions = dense<...> : tensor<1xi64>
    } : (tensor<... x !quant.uniform<ui8:f32, input_scale:input_zp>,
         tensor<!quant.uniform<ui8:f32, input_scale:input_zp>>)
    -> tensor<... x !quant.uniform<i32:f32, output_scale:outout_zp>>

Note that, in the example above, like the non-quantized version of reduce op, the type of blocks arguments are the same. That is because from the POV of the semantic, there are any two elements to be added. The bitwidth of the accumulation is captured via output quantized element type.

Please let me know your review comments.

loganchien · 2023-05-25T20:51:03Z

I only skimmed through the commit message, but I think only extending the return type of the compute region doesn't work because the input of the compute region can take larger size too.

For example: The schedule can be:

[init_val/8] [input/8]  [input/8]
    \        /         /
  [compute/32]        /
        \            /
         [compute/32]

sdasgup3 · 2023-05-26T19:17:00Z

@loganchien Thanks for the feedback.

Yes, I can see your point and agree with it.

Upon further thinking into this: Here is an high level overview of an option in order to accommodate higher (than input and init_value) bit-width accumulation. I am borrowing the symbol names from the example mentioned above:

input and init_value to have the same storage element type. This is motivated by the fact that the operand types for input and init_value represent the value types and accumulation semantics is expressed in the reduced block.
The block arguments of reduce can have same storage type (which is the case when the reducer block is performing non-accumulative operations like min/max) OR different storage types (for reducer block with accumulation). From the POV of specification of constraints, this means that there will be no constraint between the storage type of block arguments.
- In the latter case, either (A) both be equal to input's storage type or (B) both be equal to accumulation storage type or (C) a mix of storage types of input and accumulator. For (C), the order is not important, as the op semantics does not require any specific order to fetch the binary tree nodes for computation.
The result storage type of the operation should match the block-return storage type.
If the block_arguments have the same storage_type , then storage_type of block-result is equal to the storage type of one of the block arguments. Else, storage_type of block-result is equal to the one with large storage bitwidth.
The expressed_type, quantization parameters of input, init_value, block_arguments, block_result and result will be the same.

Notes:

(4) (5) ensures that the op return type can be inferred.
IMO, we should not impose any restriction of the storage_type of block arguments. Th accumulation bitwidth will be defined by the user.

Please, let me know your feedback!

wecing · 2023-05-26T20:44:25Z

Thank you @sdasgup3 for the update.

For (C), the order is not important, as the op semantics does not require any specific order to fetch the binary tree nodes for computation.

Does that imply that for reduce(+, [1, 2, 3, 4]), the spec allows ((a + b) + (c + d)) as well? If that's the case then maybe we should prohibit the hybrid option C.

Also what do you think about explicitly defining the conversion behavior between input/output type and accumulator type, e.g. "converted by zero-extending storage value"?

sdasgup3 · 2023-05-26T22:32:23Z

Hi @loganchien (upd: sorry I should have referred my reply to refer to @wecing )

Does that imply that for reduce(+, [1, 2, 3, 4]), the spec allows ((a + b) + (c + d)) as well? If that's the case then maybe we should prohibit the hybrid option C.

Yes, the spec allows that.

Let say V = value type, A = accumulation type. There are four possible signatures for a block, which represents computations with specific types.
(1) block(V, V): The computation involves up casting both block arguments to accum type and doing the computation.
(2) block(A, A): The computation involves just doing the computation (no up casting to accum type needed)
(3) block(A, V) // hybrid: The computation invove up-casting one of the arg with value type to accum type and then doing the computation.
(4) block(V, A) // hybrid: Same as the above

The specific schedule ((a + b) + (c + d)) allows computations (1) and (2).
There could be other schedules like (((a+b)+c) +d) which involve computation (1) and (3/4)
Also, for schedule ((a+b) + c) + (d +e) involve all the four computations.

The open question, that I was pondering about, is how a block, with a fixed static type signature, like any one of 1-4, can be used for any of these schedules. One solution could be use only (2) and update the specification of schedule execution to always up cast the binary tree nodes to accum type before before applying (2).

Also what do you think about explicitly defining the conversion behavior between input/output type and accumulator type, e.g. "converted by zero-extending storage value"?

May be I am not sure about what you meant by "defining the conversion behavior between input/output type and accumulator type". Can you suggest an example?

On a related note how about we explicitly zero-extending the input and init_value tensors to accumulation type before doing the reduction operation. This need minimalist changes to the spec (same type constraints as we currently have). WDYT?

loganchien · 2023-05-26T23:40:21Z

@wecing wrote:

defining the conversion behavior between input/output type and accumulator type

Also what do you think about explicitly defining the conversion behavior between input/output type and accumulator type, e.g. "converted by zero-extending storage value"?

@sdasgup3 wrote:

On a related note how about we explicitly zero-extending the input and init_value tensors to accumulation type before doing the reduction operation

This is another idea that circulates in our team. Keep the reduce computation with symmetric semantics but allow users to define input->intermediate or intermediate->output computation:

%result = "stablehlo.reduce"(%input, %init_value) ({
    ^input_conversion(
            %input: tensor<!quant.uniform<ui8:f32, input_scale:input_zp>>):
        %input_rescaled = "stablehlo.uniform_quantize"(%input)
            : (tensor<!quant.uniform<ui8:f32, input_scale:input_zp>>)
            -> tensor<!quant.uniform<i32:f32, input_scale:0>>
        "stablehlo.return"(%input_rescaled)
            : (tensor<!quant.uniform<i32:f32, input_scale:0>>) -> ()

    }, {
    ^reduce_computation(
            %lhs: tensor<!quant.uniform<i32:f32, input_scale:0>>,
            %rhs: tensor<!quant.uniform<i32:f32, input_scale:0>>):
        %add = "stablehlo.add"(%lhs, %rhs)
            : (tensor<!quant.uniform<i32:f32, input_scale:0>>,
               tensor<!quant.uniform<i32:f32, input_scale:0>>)
            -> tensor<!quant.uniform<i32:f32, input_scale:0>>
        "stablehlo.return"(%add)
            : (tensor<!quant.uniform<i32:f32, input_scale:0>>) -> ()
    }, {
    ^output_conversion(
            %intermediate_result: tensor<!quant.uniform<i32:f32, input_scale:0>>):
        %output_rescaled = "stablehlo.uniform_quantize"(%intermediate_result)
            : (tensor<!quant.uniform<i32:f32, input_scale:0>>)
            -> tensor<!quant.uniform<ui8:f32, output_scale:output_zp>>
        "stablehlo.return"(%output_rescaled)
            : (tensor<!quant.uniform<ui8:f32, output_scale:output_zp>>) -> ()
    }) {
        dimensions = dense<1> : tensor<1xi64>
    } : (tensor<!quant.uniform<ui8:f32, input_scale:input_zp>>,
         tensor<!quant.uniform<ui8:f32, input_scale:input_zp>>)
    -> tensor<!quant.uniform<ui8:f32, output_scale:output_zp>>

sdasgup3 · 2023-05-27T00:55:00Z

On a related note how about we explicitly zero-extending the input and init_value tensors to accumulation type before doing the reduction operation

Just for the completeness of my thoughts, the following is what i meant:

%input_rescaled = "stablehlo.uniform_quantize"(%input)
            : (tensor<!quant.uniform<ui8:f32, input_scale:input_zp>>)
            -> tensor<!quant.uniform<i32:f32, input_scale:0>>

%init_value_rescaled = "stablehlo.uniform_quantize"(%init_value)
            : (tensor<!quant.uniform<ui8:f32, input_scale:input_zp>>)
            -> tensor<!quant.uniform<i32:f32, input_scale:0>>

%result = "stablehlo.reduce"(%input, %init_value) ({
    ^reduce_computation(
            %lhs: tensor<!quant.uniform<i32:f32, input_scale:0>>,
            %rhs: tensor<!quant.uniform<i32:f32, input_scale:0>>):
        %add = "stablehlo.add"(%lhs, %rhs)
            : (tensor<!quant.uniform<i32:f32, input_scale:0>>,
               tensor<!quant.uniform<i32:f32, input_scale:0>>)
            -> tensor<!quant.uniform<i32:f32, input_scale:0>>
        "stablehlo.return"(%add)
            : (tensor<!quant.uniform<i32:f32, input_scale:0>>) -> ()
    }) {
        dimensions = dense<1> : tensor<1xi64>
    } : (tensor<!quant.uniform<i32:f32, input_scale:0>>,
         tensor<!quant.uniform<i32:f32, input_scale:0>>)
    -> tensor<!quant.uniform<i32:f32, output_scale:0>>

With that, spec for reduce can will be pretty much the same (and simple) in terms of type constraints. Also, I note that to accommodate multiple input operands in the stablehlo.reduce, we might need to support stablehlo.uniform_quantize with variadic-input support. Do you see any problem with that?

This is another idea that circulates in our team. Keep the reduce computation with symmetric semantics but allow users to define input->intermediate or intermediate->output computation:

Thanks @loganchien for the example. That sounds like a valid option.

A few observations, along with some questions, that I have here:

I can understand that output_conversion is introduced to achieve symmetry in computation enabling the reduce_computation to have a fixed type signature. IMO, this is at the cost of redundancy in application of output_conversion after every reduce_computation, right?
There could be scenarios when the application input_conversion should be skipped. For example, when the reduce_computation is not accumulative in nature (like for min/max computations) or for non-quantized scenario. It seems that you meant the input_conversion and output_conversion blocks to be optional and to be provided only when the the block argument type of reducer_computation in different from %input (or %init_value). From the POV of specification, there should be verification constraints to ensure that.
There is no update needed in the specification in schedule execution, which looks pretty clean.

Let me know if I miss out anything in my observations.

loganchien · 2023-05-27T01:31:12Z

IMO, this is at the cost of redundancy in application of output_conversion after every reduce_computation, right?

Not every reduce_computation. It is only applied once after all reduce computation are done for a particular output element.

There could be scenarios when the application input_conversion should be skipped.

Yes. In that case, it can be optional and the input element is copied as-is. The verification constraint should be: if the types are different, there must be a conversion block. If types are the same, it is optional (people can do some preprocessing if they want) but can be skipped for brevity.

sdasgup3 · 2023-05-27T01:52:12Z

Not every reduce_computation. It is only applied once after all reduce computation are done for a particular output element.

IIUC, for all the input and init_value s contributing to a particular output element, we first use the input_conversion to scale all of them up to accum type and then use the reduce_computation pairwise on the scaled up elements and at the end scale down the final result, for that output element. I think you meant this to be one possible implementation in order to avoid redundancy. right?

sdasgup3 · 2023-05-27T01:55:08Z

Also looking forward to some feedback on

On a related note how about we explicitly zero-extending the input and init_value tensors to accumulation type before doing the reduction operation

loganchien · 2023-05-27T02:27:43Z

IIUC, for all the input and init_value s contributing to a particular output element, we first use the input_conversion to scale all of them up to accum type and then use the reduce_computation pairwise on the scaled up elements and at the end scale down the final result, for that output element. I think you meant this to be one possible implementation in order to avoid redundancy. right?

It is not about redundancy. It is about not rounding/saturating value too early.

input_convert(qui8) -> qi32
reduce(qi32, qi32) -> qi32
output_convert(qi32) -> qui8

For a 4-element qi8 input, we only want to do:

init_value_converted = input_convert(init_value)

input_converted[0] = input_convert(input[0])
input_converted[1] = input_convert(input[1])
input_converted[2] = input_convert(input[2])
input_converted[3] = input_convert(input[3])

accumulated = reduce(
    reduce(input_converted[0], input_converted[1]),
    reduce(input_converted[2], input_converted[3]))

output = output_convert(accumulated)

We don't want to round off the bits or saturate at qui8 range when we are doing the reduce, so that the inputs get the chance to cancel each other. Allowing them to be cancelled is the motivating reason to do this kind of extension.

For example, if the zero-point adjusted input is [64, 64, -128, -128], we want the final result to be -128 instead of saturate_add_i8(saturate_add_i8(64, 64), saturate_add_i8(-128, -128)) = saturate_add_i8(127, -128) = -1.

On a related note how about we explicitly zero-extending the input and init_value tensors to accumulation type before doing the reduction operation.

That is the third option being circulated in the internal doc. But I don't like it because doing tensor-level type conversion means 4x (run-time) memory usage. All serious compilers/interpreters must do a graph-level rewriting to remove the inefficiency before the they compile or execute the model.

loganchien · 2023-05-27T07:12:02Z

Also, I note that to accommodate multiple input operands in the stablehlo.reduce, we might need to support stablehlo.uniform_quantize with variadic-input support. Do you see any problem with that?

No. I don't see problem with multiple-operand reductions. Multiple-operand reductions only need multiple-operand input_conversion. They don't need multiple-operand stablehlo.uniform_quantize:

%result = "stablehlo.reduce"(%input0, %init_value0, %input1, %init_value1) ({
    ^input_conversion(
            %input0: tensor<!quant.uniform<ui8:f32, input0_scale:input0_zp>>,
            %input1: tensor<!quant.uniform<ui8:f32, input0_scale:input0_zp>>):
        %input0_rescaled = "stablehlo.uniform_quantize"(%input0)
            : (tensor<!quant.uniform<ui8:f32, input0_scale:input0_zp>>)
            -> tensor<!quant.uniform<i32:f32, input0_scale:0>>
        %input1_rescaled = "stablehlo.uniform_quantize"(%input1)
            : (tensor<!quant.uniform<ui8:f32, input1_scale:input1_zp>>)
            -> tensor<!quant.uniform<i32:f32, input1_scale:0>>
        "stablehlo.return"(%input0_rescaled, %input1_rescaled)
            : (tensor<!quant.uniform<i32:f32, input0_scale:0>>,
               tensor<!quant.uniform<i32:f32, input1_scale:0>>) -> ()
    }, {
    ^reduce_computation(
            %lhs: tensor<!quant.uniform<i32:f32, input0_scale:0>>,
            %rhs: tensor<!quant.uniform<i32:f32, input1_scale:0>>):
        // ....
    }, {
    ^output_conversion(
            %intermediate0_result: tensor<!quant.uniform<i32:f32, input0_scale:0>>
            %intermediate1_result: tensor<!quant.uniform<i32:f32, input1_scale:0>>):
       // ....
    }) {
        dimensions = dense<1> : tensor<1xi64>
    } : (tensor<... x !quant.uniform<ui8:f32, input0_scale:input_zp>>,
         tensor<... x !quant.uniform<ui8:f32, input0_scale:input_zp>>,
         tensor<... x !quant.uniform<ui8:f32, input1_scale:input_zp>>,
         tensor<... x !quant.uniform<ui8:f32, input1_scale:input_zp>>)
    -> (tensor<... x !quant.uniform<ui8:f32, output0_scale:output0_zp>>,
        tensor<... x !quant.uniform<ui8:f32, output1_scale:output1_zp>>)

wecing · 2023-06-02T21:49:17Z

So, if I understand it correctly, to summarize the discussions so far (please let me know if I missed anything):

(Let V = input type, A = accumulator type)

@loganchien pointed out that the original proposal reduce_region(V, V) -> A doesn't work since it implies A -> V truncation in the middle of reduction;
@wecing suggested since stablehlo.reduce(+, [1, 2, 3, 4]) allows (1 + 2) + (3 + 4), we may have to use reduce_region(A, A) -> A, which brings up the issue of how/where to do the initial V -> A conversion;
@loganchien brought up the idea of adding input_conversion and output_conversion regions to stablehlo.reduce;
@sdasgup3 proposeed using separate casting ops outside of stablehlo.reduce (before and after it) to do the qi8 <=> qi32 conversion, so that stablehlo.reduce only needs to deal with qi32;
@sdasgup3 mentioned varargs stablehlo.uniform_quantize but @loganchien thought it's unnecessary.

There were also some discussions between @loganchien and @sdasgup3 where IIUC @sdasgup3 was asking about from_accum(reduce(to_accum(input))) and @loganchien was also saying reduction should happen on accumulator type, i.e. we shouldn't truncate too early.

Here are my thoughts:

StableHLO needs to support the reduce_region(A, A) -> A use case (but does not require it, so people could still write reduce_region(V, V) -> V if they really want to). The question is whether to do V <=> A conversion inside (via input_conversion and output_conversion regions) or outside (via casting ops) of stablehlo.reduce. I think either way works, but personally I think doing that inside stablehlo.reduce is slightly easier for my case.
We could add restrictions on the V and A type, e.g. we could require them to have the same quant params (scale and zp), and only (optionally) differ in the storage type. In this way we could drop the requirement for input_conversion and output_conversion regions and specify that stablehlo.reduce will zero-extend / truncate the values for us. So I am envisioning something like this:

%result = "stablehlo.reduce"(%input, %init_value) ({
    ^reduce_computation(
            %lhs: tensor<!quant.uniform<i32:f32, input_scale:0>>,
            %rhs: tensor<!quant.uniform<i32:f32, input_scale:0>>):
        %add = "stablehlo.add"(%lhs, %rhs)
            : (tensor<!quant.uniform<i32:f32, input_scale:0>>,
               tensor<!quant.uniform<i32:f32, input_scale:0>>)
            -> tensor<!quant.uniform<i32:f32, input_scale:0>>
        "stablehlo.return"(%add)
            : (tensor<!quant.uniform<i32:f32, input_scale:0>>) -> ()
    }) {
        dimensions = dense<1> : tensor<1xi64>
    } : (tensor<!quant.uniform<i8:f32, input_scale:0>>,
         tensor<!quant.uniform<i8:f32, input_scale:0>>)
    -> tensor<!quant.uniform<i8:f32, output_scale:0>>

What do you think?

sdasgup3 · 2023-06-09T00:12:13Z

Hi @wecing
Thanks for summarizing the proposals.

Regarding your proposal:
It proposes to use a reducer block which can work only on the accumulator type. One thing that stuck me initailly that how this reducer block be used to compute something like (+ 1 2), which expects the reducer block to take as argument two value of type tensor<!quant.uniform<i8:f32, input_scale:0>>. Now it is clarified (with offline discussion) that what you meant in that the leaf node of the schedule binary tree has to be implicitly converted. I believe @loganchien proposal makes that conversion step explicit by using input_conversion and output_conversion blocks. Also, I note that, per @loganchien proposal, for reduce op with quantized types

The input_conversion will be applied to convert just the leaf nodes of the binary tree.
The output_conversion is applied once after all reduce computation are done for a particular output element.

@jingpu

Re:
Logan>> That is the third option being circulated in the internal doc. But I don't like it because doing tensor-level type conversion means 4x (run-time) memory usage. All serious compilers/interpreters must do a graph-level rewriting to remove the inefficiency before the they compile or execute the model.

Just for context, @loganchien is referring to A where the proposal is about using separate quantize and dequantize ops before and after stablehlo.reduce to do the qi8 <=> qi32 conversion, so that stablehlo.reduce only needs to deal with qi32.

Can you please elaborate on the problem @loganchien is referring to? (sorry that I have to delegate)

If we have one single op, as proposed in B, then I believe it can be readily mapped an efficient reduce kernel. For (A), pattern matching could be an additional concern and I am trying to understand how much is that an overhead.

rewu93

One dumb question. How does the user define the accumulator type A if it is inside the reduce block?

sdasgup3 · 2023-06-12T18:02:14Z

One dumb question. How does the user define the accumulator type A if it is inside the reduce block?

Per @loganchien proposal, the input conversion block with specify the accumulation type.
Per my proposal, the stablehlo.uniform_quantize will specify the accumulation type.

wecing · 2023-06-12T23:21:28Z

Hey @sdasgup3 -- I synced with @jingpu offline, and unfortunately neither of us know what @loganchien's concern on your proposal A was. But, as I mentioned above in the summary comment,

There were also some discussions between @loganchien and @sdasgup3 where IIUC @sdasgup3 was asking about from_accum(reduce(to_accum(input))) and @loganchien was also saying reduction should happen on accumulator type, i.e. we shouldn't truncate too early.

My guess is that:

It was partially a mis-communication issue and @loganchien was worried about potential early truncation, but your proposal doesn't really have this issue (FYI: I think @loganchien will be back next week; hopefully we could get more clarification by then);
As you quoted, All serious compilers/interpreters must do a graph-level rewriting to remove the inefficiency before the they compile or execute the model, the concern was that this would make certain things inefficient, e.g. interpreters, and some debug tools, would have to do the q8->q32 quantization, which may be impractical for very large models, e.g. LLMs.

loganchien · 2023-06-20T20:51:23Z

Let me elaborate my thinking on two different ways to express "accumulator" type in the reduction:

The first one (my favorite one for now) is to introduce input_conversion and output_conversion as described in the earlier comment:

%result = "stablehlo.reduce"(%input, %init_value) ({
    ^input_conversion(
            %input: tensor<!quant.uniform<ui8:f32, input_scale:input_zp>>):
        %input_rescaled = "stablehlo.uniform_quantize"(%input)
            : (tensor<!quant.uniform<ui8:f32, input_scale:input_zp>>)
            -> tensor<!quant.uniform<i32:f32, input_scale:0>>
        "stablehlo.return"(%input_rescaled)
            : (tensor<!quant.uniform<i32:f32, input_scale:0>>) -> ()

    }, {
    ^reduce_computation(
            %lhs: tensor<!quant.uniform<i32:f32, input_scale:0>>,
            %rhs: tensor<!quant.uniform<i32:f32, input_scale:0>>):
        %add = "stablehlo.add"(%lhs, %rhs)
            : (tensor<!quant.uniform<i32:f32, input_scale:0>>,
               tensor<!quant.uniform<i32:f32, input_scale:0>>)
            -> tensor<!quant.uniform<i32:f32, input_scale:0>>
        "stablehlo.return"(%add)
            : (tensor<!quant.uniform<i32:f32, input_scale:0>>) -> ()
    }, {
    ^output_conversion(
            %intermediate_result: tensor<!quant.uniform<i32:f32, input_scale:0>>):
        %output_rescaled = "stablehlo.uniform_quantize"(%intermediate_result)
            : (tensor<!quant.uniform<i32:f32, input_scale:0>>)
            -> tensor<!quant.uniform<ui8:f32, output_scale:output_zp>>
        "stablehlo.return"(%output_rescaled)
            : (tensor<!quant.uniform<ui8:f32, output_scale:output_zp>>) -> ()
    }) {
        dimensions = dense<1> : tensor<1xi64>
    } : (tensor<!quant.uniform<ui8:f32, input_scale:input_zp>>,
         tensor<!quant.uniform<ui8:f32, input_scale:input_zp>>)
    -> tensor<!quant.uniform<ui8:f32, output_scale:output_zp>>

Its execution looks like:

[In0:i8]           [In1:i8]           [In2:i8]           [In3:i8]
|                  |                  |                  |
v                  v                  v                  v
[InputConversion]  [InputConversion]  [InputConversion]  [InputConversion]
(i32)              (i32)              (i32)              (i32)
|                  |                  |                  |
v                  v                  v                  v                  \
[ReduceComputation_]                  [ReduceComputation_]                  |
(i32)                                 (i32)                                 |
|                                     |                                     | Reduce schedule can be arbitrary
v                                     v                                     |
[ReduceComputation____________________]                                     /
(i32)
|
v
[OutputConversion]
(i8)

The second one is proposed by @sdasgup3, which simply suggest adding stablehlo.convert/stablehlo.uniform_quantize ops before and after the reduce op, which operates on the "accumulator" type:

%widen = "stablehlo.convert"(%input)
    : (tensor<... x !quant.uniform<ui8:f32, ...>>) -> tensor<... x !quant.uniform<i32:f32, ...>>
%reduce = "stablehlo.reduce"(%widen) {
    ^reduce_computation(%lhs: !quant.uniform<i32:f32, ...>, %rhs: !qunat.uniform<i32:f32, ...>):
        // reduce_computation_block
    }
    : (tensor<... x !quant.uniform<i32:f32, ...>>) -> tensor<... x !quant.uniform<i32:f32, ...>>
%narrowed = "stablehlo.convert"(%reduce)
    : (tensor<... x !quant.uniform<i32:f32, ...>>) -> tensor<... x !quant.uniform<ui8:f32, ...>>

Both options have pros and cons:

The first option enables programmers to program at (almost) baremetal (*argument1). If the HW can support reduction computation in narrow type (e.g. in the SIMD instruction set, we typically do widening/compute/narrowing within the kernel to save the memory bandwidth), the programmer can explicitly request for that. The second option can't express the intent. It will rely on the compiler to do pattern matching and optimization to avoid the 4x memory bandwidth increase (*argument2). The compiler (or even interpreter) must pattern match 3 operations and map them into some internal representation (that is not expressible in StableHLO) before their compilation or execution (*argument3).

In the second option, there is also a risk that the consecutive stablehlo.convert operations might be folded or converted to something else. And then, the best-effort approach of the backend may fail. For example, in the bf16 reduction:

%widen0 = "stablehlo.convert"(%input0)
    : (tensor<... x bf16>) -> tensor<... x f32>
%reduce0 = "stablehlo.reduce"(%widen0) {
    ^reduce_computation(%lhs: f32, %rhs: f32):
        // reduce_computation_block
    }
    : (tensor<... x f32>) -> tensor<... x f32>
%narrowed0 = "stablehlo.convert"(%reduce0)        //
    : (tensor<... x f32>) -> tensor<... x bf16>   //
                                                  // potential target of so-called target-independent optimization
%widen1 = "stablehlo.convert"(%narrowed0)         //
    : (tensor<... x bf16>) -> tensor<... x f32>   //
%reduce1 = "stablehlo.reduce"(%widen1) {
    ^reduce_computation(%lhs: f32, %rhs: f32):
        // reduce_computation_block
    }
    : (tensor<... x f32>) -> tensor<... x f32>
%narrowed1 = "stablehlo.convert"(%reduce1)
    : (tensor<... x f32>) -> tensor<... x bf16>

Some optimizations or folders may transform the code above into:

%widen0 = "stablehlo.convert"(%input0)
    : (tensor<... x bf16>) -> tensor<... x f32>
%reduce0 = "stablehlo.reduce"(%widen0) {
    ^reduce_computation(%lhs: f32, %rhs: f32):
        // reduce_computation_block
    }
    : (tensor<... x f32>) -> tensor<... x f32>

%reduce1 = "stablehlo.reduce"(%reduce0) {
    ^reduce_computation(%lhs: f32, %rhs: f32):
        // reduce_computation_block
    }
    : (tensor<... x f32>) -> tensor<... x f32>
%narrowed1 = "stablehlo.convert"(%reduce1)
    : (tensor<... x f32>) -> tensor<... x bf16>

and the best-effort approach in the backend can't kick in and the compiler must allocate an f32 buffer for %reduce0 (*argument4).

The second option still have some pros though. It doesn't require significant changes to existing spec (ReduceOp in particular). However, it puts the burden to the backend writer.

From my own perspective, since my backend must do semantic graph rewrite anyway, the differences between two options are only about (1) the ease of the pattern matching implementation and (2) keep early transformations away from breaking the best-effort patterns. However, I believe the extra expressiveness might find some use cases in other more general architectures (e.g. CPU/GPU).

Regarding the point on extending stablehlo.uniform_quantize for variadic operands, I think it is not necessary even if we choose the first option. For example, if a stablohlo.reduce takes multiple input operands, we can expand the input/output conversion into multiple stablehlo.uniform_quantize operations in the conversion block if needed:

%result = "stablehlo.reduce"(%input0, %init_value0, %input1, %init_value1) ({
    ^input_conversion(
            %input0: tensor<!quant.uniform<ui8:f32, input0_scale:input0_zp>>,
            %input1: tensor<!quant.uniform<ui8:f32, input1_scale:input1_zp>>):
        %input0_rescaled = "stablehlo.uniform_quantize"(%input0)
            : (tensor<!quant.uniform<ui8:f32, input0_scale:input0_zp>>)
            -> tensor<!quant.uniform<i32:f32, input0_scale:0>>
        %input1_rescaled = "stablehlo.uniform_quantize"(%input1)
            : (tensor<!quant.uniform<ui8:f32, input1_scale:input1_zp>>)
            -> tensor<!quant.uniform<i32:f32, input1_scale:0>>
        "stablehlo.return"(%input0_rescaled, %input1_rescaled)
            : (tensor<!quant.uniform<i32:f32, input0_scale:0>>,
               tensor<!quant.uniform<i32:f32, input1_scale:0>>) -> ()
    }, {
    ^reduce_computation(
            %lhs0: tensor<!quant.uniform<i32:f32, input0_scale:0>>,
            %rhs0: tensor<!quant.uniform<i32:f32, input0_scale:0>>,
            %lhs1: tensor<!quant.uniform<i32:f32, input1_scale:0>>,
            %rhs1: tensor<!quant.uniform<i32:f32, input1_scale:0>>):
        // ... skipped ...
        "stablehlo.return"(..., ...)
            : (tensor<!quant.uniform<i32:f32, input0_scale:0>>,
               tensor<!quant.uniform<i32:f32, input1_scale:0>>) -> ()
    }, {
    ^output_conversion(
            %intermediate0_result: tensor<!quant.uniform<i32:f32, input0_scale:0>>
            %intermediate1_result: tensor<!quant.uniform<i32:f32, input1_scale:0>>):
        %output0_rescaled = "stablehlo.uniform_quantize"(%intermediate0_result)
            : (tensor<!quant.uniform<i32:f32, input0_scale:0>>)
            -> tensor<!quant.uniform<ui8:f32, output0_scale:output0_zp>>
        %output1_rescaled = "stablehlo.uniform_quantize"(%intermediate1_result)
            : (tensor<!quant.uniform<i32:f32, input1_scale:0>>)
            -> tensor<!quant.uniform<ui8:f32, output1_scale:output1_zp>>
        "stablehlo.return"(%output0_rescaled, %output1_rescaled)
            : (tensor<!quant.uniform<ui8:f32, output0_scale:output0_zp>>,
               tensor<!quant.uniform<ui8:f32, output1_scale:output1_zp>>) -> ()
    }) {
        dimensions = dense<1> : tensor<1xi64>
    } : (tensor<!quant.uniform<ui8:f32, input0_scale:input0_zp>>,
         tensor<!quant.uniform<ui8:f32, input0_scale:input0_zp>>,
         tensor<!quant.uniform<ui8:f32, input0_scale:input1_zp>>,
         tensor<!quant.uniform<ui8:f32, input0_scale:input1_zp>>)
    -> (tensor<!quant.uniform<ui8:f32, output0_scale:output0_zp>>,
        tensor<!quant.uniform<ui8:f32, output1_scale:output1_zp>>)

sdasgup3 · 2023-06-22T18:33:33Z

@loganchien Thanks for the elaborate explanation and I agree with your point.
As the proposal needs some structural changes in the op (adding additional blocks), it should go an RFC process.
Let me document that which we all can further iterate over.

…Extensibility`, and `Miscellaneous` categories. (#1647) ### Summary This PRs provides the specification of the ops from the following categories: * **Control-flow ops**: `after_all, case, if, optimization_barrier, while` - `afer_all`: Not applicable for quantization as the ops deals with setting data-flow ordering using tokens. - `case, if`: These ops implicitly captures operands from the parent scope. The result of these operations and the corresponding blocks can be quantized types. - `optimization_barrier`: The `operand` and `result` can both be quantized tensors. - `while`: The `operand` and `result` can both be quantized tensors and of same type. Per the semantics of `while` op, the output of the associated block is used for the next iteration with the `operand` used for the very first iteration. With that, I do not think we can use the `operand` and `result` with different quantization types. * **Distribution**: `all_gather, all_reduce, all_to_all, collective_permute, infeed, outfeed, partition_id, recv, reduce_scatter, replica_id, send` All the ops in this category except for the following ones will support quantized types: - `all_reduce, reduce_scatter`: These involve quantization of reduction operation #1538. I would propose to address these two after the quantization of reduction ops are addressed. - `partition_id, replica_id`: Not applicable * **Extensibility**: `custom_call, get_tuple_element, tuple` All the ops in this category supports all the `ValueType`s including `QuantizedTensorType`. No change is needed in the existing spec. * **Miscellaneous**: `batch_norm_grad, batch_norm_inference, batch_norm_training, cholesky, constant, fft, iota, rng, rng_bit_generator, triangular_solve` - The specification of `batch_norm_inference`, `cholesky`, `triangular_solve` ops follow the same `dequantize-op-quantize` strategy as we followed for the elementwise ops. - `batch_norm_grad, batch_norm_training`: These involve quantization of reduction operation #1538. I would propose to address these two after the quantization of reduction ops are addressed. - `constant, iota`: The output type of these ops can be a quantized type of user-defined quantization parameters. - `fft`: This op supports `element_type(operand) = element_type(result) = complex type`, when the `fft_type` is `FFT` or `IFFT`. For `fft_type = RFFT`, the op supports floating-point `operand` and complex type `result` and for `fft_type = IRFFT`, it supports complex type `operand` and floating-point `result`. IMO, only for `fft_type = RFFT`, the operand can support quantized type and the semantics could be to quantize the operand to floating-point before doing the computation. - `rng`: The `operand` and `result` can be of quantized type only when the `rng_distribution` is `NORMAL` in which case the semantics follows the `dequantize_op_quantize` strategy. The motivation for this choice comes from the fact the op supports floating point type only when `rng_distribution` is `NORMAL` and we propose to treat this op, on floating-point tensors, as an op on quantized tensor, along the lines of dequantize -> float computation -> quantize. - `rng_bit_generator`: To be discussed in "open questions" section below. ### Open questions - `rng_bit_generator`: This op generates a randomized output tensor based on a `rng_algorithm` algorithm (`DEFAULT`, `THREE_FRY`, and `PHILOX`). It is not clear how there algorithm should handle generating quantized output. My proposal here is to skip quantization of this op unless we have a use case for it. Please let me know your opinion on this. Please review and let me know your feedback. **update 06/26/2023** Scoping `rng` out from quantization based on #1647 (comment) **update 07/13/2023** Per #1647 (comment), I proposed not to fuse de-quantization of operand and quantization of result for `fft_type = RFFT` and `fft_type = IRFFT` respectively for fft op. With that a user can explicit de-quantize of operand or quantize of result as needed. In the current proposal, the `batch_norm_inference` op, do not consider the constant operand `epsilon` to be of quantized type. IMO, considering that operand with quantized type unnecessary complicate the expressiveness of the op.

burmako · 2023-08-18T20:47:34Z

Superseded by #1664 - which has summarized the discussion here into an RFC, which is under discussion and will result in spec updates once approved.

From the discussion emerged in #1538 about the specification of quantized `reduce` op, we zeroed in to a proposal which need some structural changes in the op (adding additional blocks). This proposed RFC is to further explore the proposal in detail. Please let me know your feedback. See also the [OpenXLA Discuss post](https://groups.google.com/a/openxla.org/g/openxla-discuss/c/iwE9is49SS4).

From the discussion emerged in openxla#1538 about the specification of quantized `reduce` op, we zeroed in to a proposal which need some structural changes in the op (adding additional blocks). This proposed RFC is to further explore the proposal in detail. Please let me know your feedback. See also the [OpenXLA Discuss post](https://groups.google.com/a/openxla.org/g/openxla-discuss/c/iwE9is49SS4).

sdasgup3 added 2 commits May 25, 2023 16:29

Specification for quantized ReduceOp ReduceWindowOp

28041ff

Specification for quantized SelectAndScatterOp

795a48c

sdasgup3 added the Spec label May 25, 2023

sdasgup3 requested a review from burmako May 25, 2023 17:22

sdasgup3 self-assigned this May 25, 2023

rewu93 reviewed Jun 12, 2023

View reviewed changes

This was referenced Jun 25, 2023

Specification of quantized ops from Control-flow, Distribution, Extensibility, and Miscellaneous categories. #1647

Merged

[RFC] Specification of quantized reduction operation #1664

Merged

burmako mentioned this pull request Jul 3, 2023

Finalize the spec for quantized AllReduceOp, BatchNormTrainingOp, BatchNormGradOp and ReduceScatterOp #1666

Closed

burmako closed this Aug 18, 2023

sdasgup3 added the Quantization label Sep 6, 2024

Conversation

sdasgup3 commented May 25, 2023

Uh oh!

loganchien commented May 25, 2023

Uh oh!

sdasgup3 commented May 26, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

wecing commented May 26, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sdasgup3 commented May 26, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

loganchien commented May 26, 2023

Uh oh!

sdasgup3 commented May 27, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

loganchien commented May 27, 2023

Uh oh!

sdasgup3 commented May 27, 2023

Uh oh!

sdasgup3 commented May 27, 2023

Uh oh!

loganchien commented May 27, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

loganchien commented May 27, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

wecing commented Jun 2, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sdasgup3 commented Jun 9, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

rewu93 left a comment

Choose a reason for hiding this comment

Uh oh!

sdasgup3 commented Jun 12, 2023

Uh oh!

wecing commented Jun 12, 2023

Uh oh!

loganchien commented Jun 20, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sdasgup3 commented Jun 22, 2023

Uh oh!

burmako commented Aug 18, 2023

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

sdasgup3 commented May 26, 2023 •

edited

Loading

wecing commented May 26, 2023 •

edited

Loading

sdasgup3 commented May 26, 2023 •

edited

Loading

sdasgup3 commented May 27, 2023 •

edited

Loading

loganchien commented May 27, 2023 •

edited

Loading

loganchien commented May 27, 2023 •

edited

Loading

wecing commented Jun 2, 2023 •

edited

Loading

sdasgup3 commented Jun 9, 2023 •

edited

Loading

loganchien commented Jun 20, 2023 •

edited

Loading