Abs node runs into error with bf16 tensor

[yuanyao-nv]

Describe the issue

I have the following simple ONNX model with a single Abs node operating on BF16 tensor.

ir_version: 10
producer_name: ""
producer_version: ""
graph {
  node {
    input: "X"
    output: "Y"
    name: "Abs0"
    op_type: "Abs"
  }
  name: "onnx_graphsurgeon_graph"
  input {
    name: "X"
    type {
      tensor_type {
        elem_type: 16
        shape {
          dim {
            dim_value: -1
          }
          dim {
            dim_value: 4
          }
          dim {
            dim_value: 5
          }
        }
      }
    }
  }
  output {
    name: "Y"
    type {
      tensor_type {
        elem_type: 16
      }
    }
  }
}
opset_import {
  domain: ""
  version: 13
}

And onnxruntime 1.20.1 gives the following error

onnxruntime.capi.onnxruntime_pybind11_state.NotImplemented: [ONNXRuntimeError] : 9 : NOT_IMPLEMENTED : Could not find an implementation for Abs(13) node with name 'Abs0'

The same model with dtype changed to FP16 works.

To reproduce

I used onnx-graphsurgeon to create the model and polygraphy to run onnxruntime.
The polygraphy command I used was polygraphy run --onnxrt Abs_bf16.onnx

Urgency

No response

Platform

Linux

OS Version

Ubuntu 24.02

ONNX Runtime Installation

Released Package

ONNX Runtime Version or Commit ID

1.20.1

ONNX Runtime API

Python

Architecture

X64

Execution Provider

Default CPU

Execution Provider Library Version

No response

[skottmckay]

The CPU EP doesn't implement bfloat16 support for Abs.

onnxruntime/onnxruntime/core/providers/cpu/math/element_wise_ops.cc

Lines 259 to 268 in 30c6825

	REG_ELEMENTWISE_TYPED_KERNEL(Abs, 13, float, Abs);
	REG_ELEMENTWISE_TYPED_KERNEL(Abs, 13, double, Abs);
	REG_ELEMENTWISE_TYPED_KERNEL(Abs, 13, int8_t, Abs);
	REG_ELEMENTWISE_TYPED_KERNEL(Abs, 13, int16_t, Abs);
	REG_ELEMENTWISE_TYPED_KERNEL(Abs, 13, int32_t, Abs);
	REG_ELEMENTWISE_TYPED_KERNEL(Abs, 13, int64_t, Abs);
	REG_ELEMENTWISE_TYPED_KERNEL(Abs, 13, uint8_t, Abs);
	REG_ELEMENTWISE_TYPED_KERNEL(Abs, 13, uint16_t, Abs);
	REG_ELEMENTWISE_TYPED_KERNEL(Abs, 13, uint32_t, Abs);
	REG_ELEMENTWISE_TYPED_KERNEL(Abs, 13, uint64_t, Abs);

A GPU based EP would be more likely to. e.g. the CUDA EP does.

[yuanyao-nv]

With onnxruntime-gpu I'm getting the same

onnxruntime.capi.onnxruntime_pybind11_state.NotImplemented: [ONNXRuntimeError] : 9 : NOT_IMPLEMENTED : Could not find an implementation for Abs(13) node with name 'Abs0'

[skottmckay]

Bug in the kernel registration. For CUDA there's a bfloat16 kernel created, but that wasn't added to the kernels the EP knows about.

Kernel creation

onnxruntime/onnxruntime/core/providers/cuda/math/unary_elementwise_ops.cc

Line 240 in 95225dd

UNARY_OP_BWUZCSILHFD(Abs, 13)

eventually gets to this line

onnxruntime/onnxruntime/core/providers/cuda/math/unary_elementwise_ops.cc

Line 214 in 95225dd

UNARY_OP_TYPED(name, ver, double)

But the CUDA EP setup is missing the bfloat16 kernel registration

onnxruntime/onnxruntime/core/providers/cuda/cuda_execution_provider.cc

Lines 1988 to 1998 in 95225dd

	BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, int8_t, Abs)>,
	BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, int16_t, Abs)>,
	BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, int32_t, Abs)>,
	BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, int64_t, Abs)>,
	BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, uint8_t, Abs)>,
	BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, uint16_t, Abs)>,
	BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, uint32_t, Abs)>,
	BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, uint64_t, Abs)>,
	BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, float, Abs)>,
	BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, double, Abs)>,
	BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, MLFloat16, Abs)>,

Is it just Abs you require bfloat16 support for? We seem to have only enabled bfloat16 for a small set of unary ops and I'm not sure why that is.

FWIW TensorRT seems to handle bfloat16 Abs.

[yuanyao-nv]

Thanks for the investigations. I encountered this issue not just for Abs but also for a few other ops.
Is there any reason the CPU EP doesn't enable them?