Repository: microsoft/onnxruntime
File: onnxruntime/core/providers/cpu/tensor/cast_op.cc (C++)

Code before:
  }
};

// tensor float -> MLFloat16
template <>
struct TensorCaster<float, MLFloat16> {
  void Cast(const OpKernelContext& ctx, const TensorShape& shape, const Tensor& in, Tensor& out) const {
    auto out_data = out.MutableData<MLFloat16>();
    auto in_data = in.Data<float>();
    const size_t shape_size = narrow<size_t>(shape.Size());

    // Check if the tensor is long enough to use threads
    if (shape_size <= 128000) {
      MlasConvertFloatToHalfBuffer(in_data, out_data, shape_size);
      return;
    }
    // Calculate the number of compute cyles per implementation
    auto cpu_info = CPUIDInfo::GetCPUIDInfo();
    double num_compute_cycles;
    if (cpu_info.HasSSE3()) {
      num_compute_cycles = static_cast<double>(shape_size >> 1);
    } else if (cpu_info.HasAVX2()) {
      num_compute_cycles = static_cast<double>(shape_size >> 2);
    } else {
      num_compute_cycles = static_cast<double>(shape_size * 10);
    }

    concurrency::ThreadPool::TryParallelFor(ctx.GetOperatorThreadPool(), shape_size,
                                            {shape_size * 2.f, shape_size * 4.f, num_compute_cycles},
                                            [in_data, out_data](std::ptrdiff_t first_span, std::ptrdiff_t last_span) {
                                              MlasConvertFloatToHalfBuffer(in_data + first_span, out_data + first_span, static_cast<size_t>(last_span - first_span));
                                            });
  }
};

// (U)Int4x2 -> string or numeric types
template <typename SrcType, typename DstType>
struct TensorCaster<SrcType, DstType,
                    std::enable_if_t<IsOrtInt4Type<SrcType>::value && IsOrtInt4ConversionType<DstType>::value>> {
  void Cast(const OpKernelContext&, const TensorShape& shape, const Tensor& in, Tensor& out) const {
    const ptrdiff_t shape_size = narrow<ptrdiff_t>(shape.Size());
    const auto* in_data = in.Data<SrcType>();
    auto* out_data = out.MutableData<DstType>();

    for (ptrdiff_t i = 0; i < shape_size; ++i) {
      // elem 0 is the low nibble, 1 the high nibble
      auto val = in_data[i >> 1].GetElem(i & 0x1);
      out_data[i] = FromInt4Converter<SrcType, DstType>::Convert(val);
    }
  }
};

Review comment (suggestion):
I think these platform-specific details should be within the MLAS code instead of an individual op implementation. in the other direction, we have `MlasConvertHalfToFloatBufferInParallel()`. maybe we should create an analogous `MlasConvertFloatToHalfBufferInParallel()`.

https://github.com/microsoft/onnxruntime/blob/af4bf436d70f62afe486bb5b9d9cc7c8f3b5b958/onnxruntime/core/mlas/lib/cast.cpp#L37-L88

Code after:
  }
};

// tensor float -> MLFloat16
template <>
struct TensorCaster<float, MLFloat16> {
  void Cast(const OpKernelContext& ctx, const TensorShape& shape, const Tensor& in, Tensor& out) const {
    auto out_data = out.MutableData<MLFloat16>();
    auto in_data = in.Data<float>();
    const size_t shape_size = narrow<size_t>(shape.Size());
    MlasConvertFloatToHalfBufferInParallel(in_data, out_data, shape_size, ctx.GetOperatorThreadPool());
  }
};

// (U)Int4x2 -> string or numeric types
template <typename SrcType, typename DstType>
struct TensorCaster<SrcType, DstType,
                    std::enable_if_t<IsOrtInt4Type<SrcType>::value && IsOrtInt4ConversionType<DstType>::value>> {
  void Cast(const OpKernelContext&, const TensorShape& shape, const Tensor& in, Tensor& out) const {
    const ptrdiff_t shape_size = narrow<ptrdiff_t>(shape.Size());
    const auto* in_data = in.Data<SrcType>();
    auto* out_data = out.MutableData<DstType>();

    for (ptrdiff_t i = 0; i < shape_size; ++i) {
      // elem 0 is the low nibble, 1 the high nibble
      auto val = in_data[i >> 1].GetElem(i & 0x1);
      out_data[i] = FromInt4Converter<SrcType, DstType>::Convert(val);
    }
  }
};

// string or numeric types -> (U)Int4x2
template <typename SrcType, typename DstType>
struct TensorCaster<SrcType, DstType,
                    std::enable_if_t<IsOrtInt4ConversionType<SrcType>::value && IsOrtInt4Type<DstType>::value>> {
  void Cast(const OpKernelContext&, const TensorShape& shape, const Tensor& in, Tensor& out) const {
    const ptrdiff_t shape_size = narrow<ptrdiff_t>(shape.Size());
    const auto* in_data = in.Data<SrcType>();
    auto* out_data = out.MutableData<DstType>();

    ptrdiff_t i = 0;
    for (; i < shape_size - 1; i += 2) {
      auto low_val = ToInt4Converter<SrcType, DstType>::Convert(in_data[i]);
      auto high_val = ToInt4Converter<SrcType, DstType>::Convert(in_data[i + 1]);
      out_data[i >> 1] = DstType(low_val, high_val);
    }

    if (i < shape_size) {
      auto low_val = ToInt4Converter<SrcType, DstType>::Convert(in_data[i]);
      out_data[i >> 1] = DstType(low_val, 0);
    }