Bugfix: Add missing 64-bit integers support for some reduction operators

[huningxin]

reduceL1, reduceProduct, reduceSum and reduceSumSquare already support 32-bit integers. 64-bit integers should also be supported.

Fix #283, #694

---

Preview | Diff

[Honry]

@fdwr, @huningxin, starting from ONNX Opset 18, all Reduce* ops support int64 and uint64 in DML EP, see

https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/core/providers/dml/DmlExecutionProvider/src/Operators/OperatorRegistration.cpp#L937-L979

While from https://learn.microsoft.com/en-us/windows/ai/directml/dml-feature-level-history#dml_feature_level_5_0

It mentions the data type expanding for

• DML_REDUCE_FUNCTION_L1
• DML_REDUCE_FUNCTION_MAX
• DML_REDUCE_FUNCTION_MIN
• DML_REDUCE_FUNCTION_MULTIPLY
• DML_REDUCE_FUNCTION_SUM
• DML_REDUCE_FUNCTION_SUM_SQUARE

But it doesn't mention floating point values for reduceMean and others. Is it a documentation issue?

[inexorabletash]

LGTM

[a-sully]

#654 proposes removing 64-bit int support for some operators. What's the rationale for adding 64-bit int support for these operators?

This may be a discussion for another issue, but does WebNN need to support 64-bit integer types at all given trends towards smaller data types (e.g. int4 and int8) for on-device ML and that some backends have little/no support for 64-bit values in the first place? (e.g. CoreML has ~no support, and most GPUs emulate support for int64)

[huningxin]

@a-sully

#654 proposes removing 64-bit int support for some operators.

IIUC, #654 mentioned lower feature level DirectML (before FL 4.1) doesn't support 64-bit integers for some operators. Higher feature DirectML doesn't have that issue. As @inexorabletash mentioned we may assume the browser always carries a copy of library that ensures the highest feature level. If that's the case, we may close that issue.

What's the rationale for adding 64-bit int support for these operators?

As #694 mentioned, the safety checker model for stable diffusion turbo demo uses reduceSum for int64 input and DirectML DML_REDUCE_FUNCTION_SUM supports that 64-bit integers.

some backends have little/no support for 64-bit values in the first place? (e.g. CoreML has ~no support, and most GPUs emulate support for int64)

We may want to follow the backend difference through #463 .

[a-sully]

What's the rationale for adding 64-bit int support for these operators?

As #694 mentioned, the safety checker model for stable diffusion turbo demo uses reduceSum for int64 input and DirectML DML_REDUCE_FUNCTION_SUM supports that 64-bit integers.

I know this is already implemented in Chromium so I don't mean to quibble too much over this, but I think it's worth questioning whether there should be a process for changing supported data types similar to the existing process for adding new operators

[huningxin]

@a-sully

but I think it's worth questioning whether there should be a process for changing supported data types similar to the existing process for adding new operators

I'd support adding a process for updating the operators. A PR is proposed: #705. PTAL.

This PR is a follow-up bug fix for #283. In that issue, @fdwr confirmed L1, SUM_SQUARE, MULTIPLY, SUM reduce functions support 32 and 64 bit integers. However, I forgot to add 64-bit integers into the table of #283 which causes @inexorabletash 's PR #646 missing the 64-bit integers support for those reduce operators.

As mentioned, this issue causes the safety checker model for stable diffusion turbo demo failing because it uses reduceSum with int64 input. Chromium already fixed it and I also updated the table of #283.

[fdwr]

But it doesn't mention floating point values for reduceMean and others. Is it a documentation issue?

@Honry : The doc is correct. DML_OPERATOR_REDUCE + REDUCE_FUNCTION_AVERAGE https://learn.microsoft.com/en-us/windows/win32/api/directml/ns-directml-dml_reduce_operator_desc only supports float. We didn't support int64 for mean/average reduction because:

• (1) the averaging division most likely would produce a floating-point result anyway
• (2) it saves shader space if not actually used (and there weren't any clients for int64 average at the time)
• (3) it raises policy questions for fractional values of whether to truncate, floor, ceil, round to nearest evens...
• (4) it's trivial to implement with an explicit REDUCE_SUM followed by DIVIDE.

Notice that similarly ReduceL1 was supported, but ReduceL2 (which has a square root that yields fractional values) does not have an int64 version (and also LogSum, LogSumExp...).

{
  "groupFieldValues": ["ARGMIN", "ARGMAX"],
  "capabilities": [
    {"featureLevel": "1.0", "InputDataType": ["float16", "float32"], "OutputDataType": ["uint32"], "DefaultRank": {"Min": 4, "Max": 4}},
    {"featureLevel": "2.1", "InputDataType": ["AllTensorDataTypes8To32"], "OutputDataType": ["uint32"], "DefaultRank": {"Min": 4, "Max": 4}},
    {"featureLevel": "3.0", "InputDataType": ["AllTensorDataTypes8To32"], "OutputDataType": ["AllNonFloatTensorDataTypes32To64"], "DefaultRank": {"Min": 1, "Max": 8}},
    {"featureLevel": "4.1", "InputDataType": ["AllTensorDataTypes8To64ExceptFloat64"], "OutputDataType": ["AllNonFloatTensorDataTypes32To64"], "DefaultRank": {"Min": 1, "Max": 8}}
  ]
},{
  "groupFieldValues": ["AVERAGE", "L2", "LOG_SUM", "LOG_SUM_EXP"],
  "capabilities": [
    {"featureLevel": "1.0", "InputDataType": ["float16", "float32"], "OutputDataType": ["float16", "float32"], "DefaultRank": {"Min": 4, "Max": 4}},
    {"featureLevel": "3.0", "InputDataType": ["float16", "float32"], "OutputDataType": ["float16", "float32"], "DefaultRank": {"Min": 1, "Max": 8}}
  ]
},{
  "groupFieldValues": ["L1", "SUM_SQUARE"],
  "capabilities": [
    {"featureLevel": "1.0", "InputDataType": ["float16", "float32"], "OutputDataType": ["float16", "float32"], "DefaultRank": {"Min": 4, "Max": 4}},
    {"featureLevel": "3.0", "InputDataType": ["float16", "float32"], "OutputDataType": ["float16", "float32"], "DefaultRank": {"Min": 1, "Max": 8}},
    {"featureLevel": "5.0", "InputDataType": ["AllTensorDataTypes32To64ExceptFloat64"], "OutputDataType": ["AllTensorDataTypes32To64ExceptFloat64"], "DefaultRank": {"Min": 1, "Max": 8}}
  ]
},{
  "groupFieldValues": ["MIN", "MAX"],
  "capabilities": [
    {"featureLevel": "1.0", "InputDataType": ["float16", "float32"], "OutputDataType": ["float16", "float32"], "DefaultRank": {"Min": 4, "Max": 4}},
    {"featureLevel": "2.1", "InputDataType": ["float16", "float32", "uint32", "int32"], "OutputDataType": ["float16", "float32", "uint32", "int32"], "DefaultRank": {"Min": 4, "Max": 4}},
    {"featureLevel": "3.0", "InputDataType": ["AllTensorDataTypes8To32"], "OutputDataType": ["AllTensorDataTypes8To32"], "DefaultRank": {"Min": 1, "Max": 8}},
    {"featureLevel": "5.0", "InputDataType": ["AllTensorDataTypes8To64ExceptFloat64"], "OutputDataType": ["AllTensorDataTypes8To64ExceptFloat64"], "DefaultRank": {"Min": 1, "Max": 8}}
  ]
},{
  "groupFieldValues": ["MULTIPLY", "SUM"],
  "capabilities": [
    {"featureLevel": "1.0", "InputDataType": ["float16", "float32"], "OutputDataType": ["float16", "float32"], "DefaultRank": {"Min": 4, "Max": 4}},
    {"featureLevel": "2.1", "InputDataType": ["float16", "float32", "uint32", "int32"], "OutputDataType": ["float16", "float32", "uint32", "int32"], "DefaultRank": {"Min": 4, "Max": 4}},
    {"featureLevel": "3.0", "InputDataType": ["float16", "float32", "uint32", "int32"], "OutputDataType": ["float16", "float32", "uint32", "int32"], "DefaultRank": {"Min": 1, "Max": 8}},
    {"featureLevel": "5.0", "InputDataType": ["AllTensorDataTypes32To64ExceptFloat64"], "OutputDataType": ["AllTensorDataTypes32To64ExceptFloat64"], "DefaultRank": {"Min": 1, "Max": 8}}
  ]
}

[fdwr]

👍 Thanks Ningxin.

[philloooo]

As mentioned, this issue causes the safety checker model for stable diffusion turbo demo failing because it uses reduceSum with int64 input.

@huningxin given that int64 is not supported on CoreML, how hard is it for the stable diffusion turbo demo to convert the model from int64 to int32?

[huningxin]

@philloooo

As mentioned, this issue causes the safety checker model for stable diffusion turbo demo failing because it uses reduceSum with int64 input.

@huningxin given that int64 is not supported on CoreML, how hard is it for the stable diffusion turbo demo to convert the model from int64 to int32?

The original model casts to int64 before reduceSum. We'll try to change it to int32 and see whether it still works. Will keep this thread posted.

[huningxin]

@philloooo

As mentioned, this issue causes the safety checker model for stable diffusion turbo demo failing because it uses reduceSum with int64 input.

@huningxin given that int64 is not supported on CoreML, how hard is it for the stable diffusion turbo demo to convert the model from int64 to int32?

The original model casts to int64 before reduceSum. We'll try to change it to int32 and see whether it still works. Will keep this thread posted.

After the investigation, the two reduceSum operators of safety checker model are able to take int32 input. @Honry helped create an int32 version and host it at https://huggingface.co/lwanming/sd-turbo-ort-web/blob/main/safety_checker_int32_reduceSum.onnx, feel free to check it out.

More details:

1. The first reduceSum (name "/ReduceSum") takes a 2-D input tensor in shape [1, 3] and reduce along axis 1. The tensor value is either 0 or 1 because it is the output of preceding greater operator. It's safe to cast the input tensor to int32 because it won't overflow when summing three 1s together.
2. Similarly, the second reduceSum (name "/ReduceSum_1") takes a 2-D input tensor as output of preceding greater operator in shape [1, 17] and reduce along axis 1. It's also safe to cast the input tensor to int32 for the same reason.

[fdwr]

@mwyrzykowski Do you know any plans to update CoreML to support int64? It's oddly inconsistent that all the other Apple ML API's (BNNS, MPS, MLX) support int64, but CoreML does not 🤔. Is CoreML still the right API these days to implement a WebNN backend, or is it left behind by newer ones? Thanks for any information or redirections.

Type	BNNS	MPS	MLX	CoreML
uint8	BNNSDataTypeUInt8	MPSDataType.uInt8	mx.uint8	x
uint16	BNNSDataTypeUInt16	MPSDataType.uInt16	mx.uint16	x
uint32	BNNSDataTypeUInt32	MPSDataType.uInt32	mx.uint32	x
uint64	BNNSDataTypeUInt64	MPSDataType.uInt64	mx.uint64	x 🤔
int8	BNNSDataTypeInt8	MPSDataType.int8	mx.int8	x
int16	BNNSDataTypeInt16	MPSDataType.int16	mx.int16	x
int32	BNNSDataTypeInt32	MPSDataType.int32	mx.int32	ArrayFeatureType.ArrayDataType.INT32
int64	BNNSDataTypeInt64	MPSDataType.int64	mx.int64	x 🤔
float16f10e5s1 IEEE	BNNSDataTypeFloat16	MPSDataType.float16	mx.float16	ArrayFeatureType.ArrayDataType.FLOAT16
float16f7e8s1 Brain	BNNSDataTypeBFloat16	MPSDataType.bFloat16	x	x
float32f23e8s1 IEEE	BNNSDataTypeFloat32	MPSDataType.float32	mx.float32	ArrayFeatureType.ArrayDataType.FLOAT32
float64f52e11s1 IEEE	x	x	x	ArrayFeatureType.ArrayDataType.DOUBLE
float16 x 2	x	MPSDataType.complexFloat16	x	x
float32 x 2	x	MPSDataType.complexFloat32	x	x
float64 x 2	x	x	x	x
bool8	BNNSDataTypeBoolean	MPSDataType.bool	bool_	x

[mwyrzykowski]

@fdwr BNNS only runs on the CPU and MLX being open source can not use the ANE / NPU. MPS being backed by Metal only runs on the GPU.

If running on the ANE is a goal, CoreML is necessary.

[a-sully]

If running on the ANE is a goal, CoreML is necessary.

Can confirm this is an important goal :)

[fdwr]

If running on the ANE is a goal, CoreML is necessary.

Can confirm this is an important goal :)

Concur, this is important (at least for those models that can fully run on ANN or those parts of a model that are viable to run on it).