  static void Run(const Tensor& in_x, const Tensor& in_y, bool lower,
                  bool adjoint, const MatMulBCast& bcast, Tensor* out,
                  int start, int limit) {
    const bool should_bcast = bcast.IsBroadcastingRequired();
    const auto& x_batch_indices = bcast.x_batch_indices();
    const auto& y_batch_indices = bcast.y_batch_indices();
    int num_bands = in_x.dim_size(1);
    int matrix_size = in_x.dim_size(2);

    for (int64 i = start; i < limit; ++i) {
      const int64 x_batch_index = should_bcast ? x_batch_indices[i] : i;
      const int64 y_batch_index = should_bcast ? y_batch_indices[i] : i;
      auto matrix = ConstTensorSliceToEigenMatrix(in_x, x_batch_index);
      auto rhs = ConstTensorSliceToEigenMatrix(in_y, y_batch_index);
      auto output = TensorSliceToEigenMatrix(out, i);
      // Below, we use the standard algorithm for computing a triangular solve,
      // except we band limit it.
      // Given A x = b, where A is lower triangular,
      // x_i = (b_i - sum a_ij * x_j) / a_ii, where the sum is from
      // j = 0 to i - 1.
      //
      // Now, in a banded triangular matrix, when i exceeds the band size,
      // then the sum goes from j = i - band_size to i - 1, since the other
      // elements are zero.
      //
      // Finally, given the band storage format, we'll need to change the
      // indexing.
      if (lower) {
        if (!adjoint) {
          output.row(0) = rhs.row(0) / matrix(0, 0);
          for (int i = 1; i < matrix_size; ++i) {
            if (i < num_bands) {
              output.row(i).noalias() =
                  (rhs.row(i) - matrix.block(1, i, i, 1).reverse().transpose() *
                                    output.topRows(i)) /
                  matrix(0, i);
            } else {
              output.row(i).noalias() =
                  (rhs.row(i) -
                   matrix.block(1, i, num_bands - 1, 1).reverse().transpose() *
                       output.middleRows(i - (num_bands - 1), num_bands - 1)) /
                  matrix(0, i);
            }
          }
        } else {
          // In the adjoint case, here and below, we now have an upper (lower)
          // triangular matrix, and thus need to work through with the other
          // case. We can't simply conjugate `matrix` and use the upper (lower)
          // algorithm because the band storage format for upper and lower
          // triangular matrices are different (in the lower case, we pad
          // entries on the left, and in the upper case we pad entries on the
          // right.
          output.row(matrix_size - 1) =
              rhs.row(matrix_size - 1) / eigen_conj(matrix(0, matrix_size - 1));
          for (int i = matrix_size - 1; i >= 0; --i) {
            output.row(i).noalias() = rhs.row(i);
            for (int j = i + 1; j < std::min(matrix_size, i + num_bands); ++j) {
              output.row(i).noalias() -=
                  eigen_conj(matrix(j - i, j)) * output.row(j);
            }
            output.row(i) /= eigen_conj(matrix(0, i));
          }
        }
      } else {
        if (!adjoint) {
          output.row(matrix_size - 1) =
              rhs.row(matrix_size - 1) / matrix(num_bands - 1, matrix_size - 1);
          for (int i = 1; i < matrix_size; ++i) {
            int k = matrix_size - 1 - i;
            if (i < num_bands) {
              output.row(k).noalias() =
                  (rhs.row(k) - matrix.block(num_bands - 1 - i, k, i, 1)
                                        .reverse()
                                        .transpose() *
                                    output.bottomRows(i)) /
                  matrix(num_bands - 1, k);
            } else {
              output.row(k).noalias() =
                  (rhs.row(k) -
                   matrix.block(0, k, num_bands - 1, 1).reverse().transpose() *
                       output.middleRows(k + 1, num_bands - 1)) /
                  matrix(num_bands - 1, k);
            }
          }
        } else {
          output.row(0) = rhs.row(0) / eigen_conj(matrix(num_bands - 1, 0));
          for (int i = 1; i < matrix_size; ++i) {
            output.row(i).noalias() = rhs.row(i);
            for (int j = std::max(0, i - (num_bands - 1)); j < i; ++j) {
              output.row(i).noalias() -=
                  eigen_conj(matrix(num_bands - 1 - (i - j), j)) *
                  output.row(j);
            }
            output.row(i) /= eigen_conj(matrix(num_bands - 1, i));
          }
        }
      }
    }
  }


Scalar eigen_conj(const Scalar& scalar) {
  return Eigen::numext::conj<Scalar>(scalar);
}

Ensure validation sticks in banded_triangular_solve_op

PiperOrigin-RevId: 373275480
Change-Id: Id7717cf275b2d6fdb9441fbbe166d555182d2e79

/* No explicit global definitions found. */


Status TF_TensorToPyArray(Safe_TF_TensorPtr tensor, PyObject** out_ndarray) {
  // A fetched operation will correspond to a null tensor, and a None
  // in Python.
  if (tensor == nullptr) {
    Py_INCREF(Py_None);
    *out_ndarray = Py_None;
    return Status::OK();
  }
  int64 nelems = -1;
  gtl::InlinedVector<npy_intp, 4> dims;
  TF_RETURN_IF_ERROR(
      GetPyArrayDimensionsForTensor(tensor.get(), &dims, &nelems));

  // If the type is neither string nor resource we can reuse the Tensor memory.
  TF_Tensor* original = tensor.get();
  TF_Tensor* moved = TF_TensorMaybeMove(tensor.release());
  if (moved != nullptr) {
    if (ArrayFromMemory(
            dims.size(), dims.data(), TF_TensorData(moved),
            static_cast<DataType>(TF_TensorType(moved)),
            [moved] { TF_DeleteTensor(moved); }, out_ndarray)
            .ok()) {
      return Status::OK();
    }
  }
  tensor.reset(original);

  // Copy the TF_TensorData into a newly-created ndarray and return it.
  PyArray_Descr* descr = nullptr;
  TF_RETURN_IF_ERROR(GetPyArrayDescrForTensor(tensor.get(), &descr));
  Safe_PyObjectPtr safe_out_array =
      tensorflow::make_safe(PyArray_Empty(dims.size(), dims.data(), descr, 0));
  if (!safe_out_array) {
    return errors::Internal("Could not allocate ndarray");
  }
  PyArrayObject* py_array =
      reinterpret_cast<PyArrayObject*>(safe_out_array.get());
  if (TF_TensorType(tensor.get()) == TF_STRING) {
    Status s = CopyTF_TensorStringsToPyArray(tensor.get(), nelems, py_array);
    if (!s.ok()) {
      return s;
    }
  } else if (static_cast<size_t>(PyArray_NBYTES(py_array)) !=
             TF_TensorByteSize(tensor.get())) {
    return errors::Internal("ndarray was ", PyArray_NBYTES(py_array),
                            " bytes but TF_Tensor was ",
                            TF_TensorByteSize(tensor.get()), " bytes");
  } else {
    FastMemcpy(PyArray_DATA(py_array), TF_TensorData(tensor.get()),
               PyArray_NBYTES(py_array));
  }

  *out_ndarray = safe_out_array.release();
  return Status::OK();
}

Fix `tf.raw_ops.ResourceCountUpTo` null pointer dereference.

PiperOrigin-RevId: 368294347
Change-Id: I2c16fbfc9b4966c402c3d8e311f0d665a9c852d8


Status NdarrayToTensor(TFE_Context* ctx, PyObject* ndarray,
                       Safe_TF_TensorPtr* ret) {
  DCHECK(ret != nullptr);

  // Make sure we dereference this array object in case of error, etc.
  Safe_PyObjectPtr array_safe(make_safe(
      PyArray_FromAny(ndarray, nullptr, 0, 0, NPY_ARRAY_CARRAY_RO, nullptr)));
  if (!array_safe) return errors::InvalidArgument("Not a ndarray.");
  PyArrayObject* array = reinterpret_cast<PyArrayObject*>(array_safe.get());

  // Convert numpy dtype to TensorFlow dtype.
  TF_DataType dtype = TF_FLOAT;
  TF_RETURN_IF_ERROR(PyArray_TYPE_to_TF_DataType(array, &dtype));

  tensorflow::int64 nelems = 1;
  gtl::InlinedVector<int64_t, 4> dims;
  for (int i = 0; i < PyArray_NDIM(array); ++i) {
    dims.push_back(PyArray_SHAPE(array)[i]);
    nelems *= dims[i];
  }

  // Create a TF_Tensor based on the fed data. In the case of non-string data
  // type, this steals a reference to array, which will be relinquished when
  // the underlying buffer is deallocated. For string, a new temporary buffer
  // is allocated into which the strings are encoded.
  if (dtype == TF_RESOURCE) {
    size_t size = PyArray_NBYTES(array);
    array_safe.release();

    if (ctx) {
      *ret = make_safe(new TF_Tensor{tensorflow::unwrap(ctx)->CreateTensor(
          static_cast<tensorflow::DataType>(dtype), {}, 0, PyArray_DATA(array),
          size, &DelayedNumpyDecref, array)});
    } else {
      *ret = make_safe(TF_NewTensor(dtype, {}, 0, PyArray_DATA(array), size,
                                    &DelayedNumpyDecref, array));
    }

  } else if (dtype != TF_STRING) {
    size_t size = PyArray_NBYTES(array);
    array_safe.release();
    if (ctx) {
      *ret = make_safe(new TF_Tensor{tensorflow::unwrap(ctx)->CreateTensor(
          static_cast<tensorflow::DataType>(dtype), dims.data(), dims.size(),
          PyArray_DATA(array), size, &DelayedNumpyDecref, array)});
    } else {
      *ret = make_safe(TF_NewTensor(dtype, dims.data(), dims.size(),
                                    PyArray_DATA(array), size,
                                    &DelayedNumpyDecref, array));
    }

  } else {
    size_t size = 0;
    void* encoded = nullptr;
    TF_RETURN_IF_ERROR(EncodePyBytesArray(array, nelems, &size, &encoded));
    if (ctx) {
      *ret = make_safe(new TF_Tensor{tensorflow::unwrap(ctx)->CreateTensor(
          static_cast<tensorflow::DataType>(dtype), dims.data(), dims.size(),
          encoded, size,
          [](void* data, size_t len, void* arg) {
            delete[] reinterpret_cast<tensorflow::tstring*>(data);
          },
          nullptr)});
    } else {
      *ret = make_safe(TF_NewTensor(
          dtype, dims.data(), dims.size(), encoded, size,
          [](void* data, size_t len, void* arg) {
            delete[] reinterpret_cast<tensorflow::tstring*>(data);
          },
          nullptr));
    }
  }

  return Status::OK();
}


Status CopyTF_TensorStringsToPyArray(const TF_Tensor* src, uint64 nelems,
                                     PyArrayObject* dst) {
  const void* tensor_data = TF_TensorData(src);
  DCHECK(tensor_data != nullptr);
  DCHECK_EQ(TF_STRING, TF_TensorType(src));

  const tstring* tstr = static_cast<const tstring*>(tensor_data);

  std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
      TF_NewStatus(), TF_DeleteStatus);
  auto iter = make_safe(PyArray_IterNew(reinterpret_cast<PyObject*>(dst)));
  for (int64 i = 0; i < static_cast<int64>(nelems); ++i) {
    const tstring& tstr_i = tstr[i];
    auto py_string =
        make_safe(PyBytes_FromStringAndSize(tstr_i.data(), tstr_i.size()));
    if (py_string == nullptr) {
      return errors::Internal(
          "failed to create a python byte array when converting element #", i,
          " of a TF_STRING tensor to a numpy ndarray");
    }

    if (PyArray_SETITEM(dst, static_cast<char*>(PyArray_ITER_DATA(iter.get())),
                        py_string.get()) != 0) {
      return errors::Internal("Error settings element #", i,
                              " in the numpy ndarray");
    }
    PyArray_ITER_NEXT(iter.get());
  }
  return Status::OK();
}


Status GetPyArrayDimensionsForTensor(const TF_Tensor* tensor,
                                     gtl::InlinedVector<npy_intp, 4>* dims,
                                     tensorflow::int64* nelems) {
  dims->clear();
  const int ndims = TF_NumDims(tensor);
  if (TF_TensorType(tensor) == TF_RESOURCE) {
    if (ndims != 0) {
      return errors::InvalidArgument(
          "Fetching of non-scalar resource tensors is not supported.");
    }
    dims->push_back(TF_TensorByteSize(tensor));
    *nelems = dims->back();
  } else {
    *nelems = 1;
    for (int i = 0; i < ndims; ++i) {
      dims->push_back(TF_Dim(tensor, i));
      *nelems *= dims->back();
    }
  }
  return Status::OK();
}


Status TF_TensorToMaybeAliasedPyArray(Safe_TF_TensorPtr tensor,
                                      PyObject** out_ndarray) {
  auto dtype = TF_TensorType(tensor.get());
  if (dtype == TF_STRING || dtype == TF_RESOURCE) {
    return TF_TensorToPyArray(std::move(tensor), out_ndarray);
  }

  TF_Tensor* moved = tensor.release();
  int64 nelems = -1;
  gtl::InlinedVector<npy_intp, 4> dims;
  TF_RETURN_IF_ERROR(GetPyArrayDimensionsForTensor(moved, &dims, &nelems));
  return ArrayFromMemory(
      dims.size(), dims.data(), TF_TensorData(moved),
      static_cast<DataType>(dtype), [moved] { TF_DeleteTensor(moved); },
      out_ndarray);
}


Status NdarrayToTensor(PyObject* obj, Tensor* ret) {
  Safe_TF_TensorPtr tf_tensor = make_safe(static_cast<TF_Tensor*>(nullptr));
  Status s = NdarrayToTensor(nullptr /*ctx*/, obj, &tf_tensor);
  if (!s.ok()) {
    return s;
  }
  return TF_TensorToTensor(tf_tensor.get(), ret);
}


Status PyArrayDescr_to_TF_DataType(PyArray_Descr* descr,
                                   TF_DataType* out_tf_datatype) {
  PyObject* key;
  PyObject* value;
  Py_ssize_t pos = 0;

  // Return an error if the fields attribute is null.
  // Occurs with an improper conversion attempt to resource.
  if (descr->fields == nullptr) {
    return errors::Internal("Unexpected numpy data type");
  }

  if (PyDict_Next(descr->fields, &pos, &key, &value)) {
    // In Python 3, the keys of numpy custom struct types are unicode, unlike
    // Python 2, where the keys are bytes.
    const char* key_string =
        PyBytes_Check(key) ? PyBytes_AsString(key)
                           : PyBytes_AsString(PyUnicode_AsASCIIString(key));
    if (!key_string) {
      return errors::Internal("Corrupt numpy type descriptor");
    }
    tensorflow::string key = key_string;
    // The typenames here should match the field names in the custom struct
    // types constructed in test_util.py.
    // TODO(mrry,keveman): Investigate Numpy type registration to replace this
    // hard-coding of names.
    if (key == "quint8") {
      *out_tf_datatype = TF_QUINT8;
    } else if (key == "qint8") {
      *out_tf_datatype = TF_QINT8;
    } else if (key == "qint16") {
      *out_tf_datatype = TF_QINT16;
    } else if (key == "quint16") {
      *out_tf_datatype = TF_QUINT16;
    } else if (key == "qint32") {
      *out_tf_datatype = TF_QINT32;
    } else if (key == "resource") {
      *out_tf_datatype = TF_RESOURCE;
    } else {
      return errors::Internal("Unsupported numpy data type");
    }
    return Status::OK();
  }
  return errors::Internal("Unsupported numpy data type");
}


inline void FastMemcpy(void* dst, const void* src, size_t size) {
  // clang-format off
  switch (size) {
    // Most compilers will generate inline code for fixed sizes,
    // which is significantly faster for small copies.
    case  1: memcpy(dst, src, 1); break;
    case  2: memcpy(dst, src, 2); break;
    case  3: memcpy(dst, src, 3); break;
    case  4: memcpy(dst, src, 4); break;
    case  5: memcpy(dst, src, 5); break;
    case  6: memcpy(dst, src, 6); break;
    case  7: memcpy(dst, src, 7); break;
    case  8: memcpy(dst, src, 8); break;
    case  9: memcpy(dst, src, 9); break;
    case 10: memcpy(dst, src, 10); break;
    case 11: memcpy(dst, src, 11); break;
    case 12: memcpy(dst, src, 12); break;
    case 13: memcpy(dst, src, 13); break;
    case 14: memcpy(dst, src, 14); break;
    case 15: memcpy(dst, src, 15); break;
    case 16: memcpy(dst, src, 16); break;
#if defined(PLATFORM_GOOGLE) || defined(PLATFORM_POSIX) && \
    !defined(IS_MOBILE_PLATFORM)
    // On Linux, memmove appears to be faster than memcpy for
    // large sizes, strangely enough.
    default: memmove(dst, src, size); break;
#else
    default: memcpy(dst, src, size); break;
#endif
  }
  // clang-format on
}


Status GetPyArrayDescrForTensor(const TF_Tensor* tensor,
                                PyArray_Descr** descr) {
  if (TF_TensorType(tensor) == TF_RESOURCE) {
    PyObject* field = PyTuple_New(3);
#if PY_MAJOR_VERSION < 3
    PyTuple_SetItem(field, 0, PyBytes_FromString("resource"));
#else
    PyTuple_SetItem(field, 0, PyUnicode_FromString("resource"));
#endif
    PyTuple_SetItem(field, 1, PyArray_TypeObjectFromType(NPY_UBYTE));
    PyTuple_SetItem(field, 2, PyLong_FromLong(1));
    PyObject* fields = PyList_New(1);
    PyList_SetItem(fields, 0, field);
    int convert_result = PyArray_DescrConverter(fields, descr);
    Py_CLEAR(field);
    Py_CLEAR(fields);
    if (convert_result != 1) {
      return errors::Internal("Failed to create numpy array description for ",
                              "TF_RESOURCE-type tensor");
    }
  } else {
    int type_num = -1;
    TF_RETURN_IF_ERROR(
        TF_DataType_to_PyArray_TYPE(TF_TensorType(tensor), &type_num));
    *descr = PyArray_DescrFromType(type_num);
  }

  return Status::OK();
}


Status PyObjectToString(PyObject* obj, const char** ptr, Py_ssize_t* len,
                        PyObject** ptr_owner) {
  *ptr_owner = nullptr;
  if (PyBytes_Check(obj)) {
    char* buf;
    if (PyBytes_AsStringAndSize(obj, &buf, len) != 0) {
      return errors::Internal("Unable to get element as bytes.");
    }
    *ptr = buf;
    return Status::OK();
  } else if (PyUnicode_Check(obj)) {
#if (PY_MAJOR_VERSION > 3 || (PY_MAJOR_VERSION == 3 && PY_MINOR_VERSION >= 3))
    *ptr = PyUnicode_AsUTF8AndSize(obj, len);
    if (*ptr != nullptr) return Status::OK();
#else
    PyObject* utemp = PyUnicode_AsUTF8String(obj);
    char* buf;
    if (utemp != nullptr && PyBytes_AsStringAndSize(utemp, &buf, len) != -1) {
      *ptr = buf;
      *ptr_owner = utemp;
      return Status::OK();
    }
    Py_XDECREF(utemp);
#endif
    return errors::Internal("Unable to convert element to UTF-8");
  } else {
    return errors::Internal("Unsupported object type ", obj->ob_type->tp_name);
  }
}


Status PyArray_TYPE_to_TF_DataType(PyArrayObject* array,
                                   TF_DataType* out_tf_datatype) {
  int pyarray_type = PyArray_TYPE(array);
  PyArray_Descr* descr = PyArray_DESCR(array);
  switch (pyarray_type) {
    case NPY_FLOAT16:
      *out_tf_datatype = TF_HALF;
      break;
    case NPY_FLOAT32:
      *out_tf_datatype = TF_FLOAT;
      break;
    case NPY_FLOAT64:
      *out_tf_datatype = TF_DOUBLE;
      break;
    case NPY_INT32:
      *out_tf_datatype = TF_INT32;
      break;
    case NPY_UINT8:
      *out_tf_datatype = TF_UINT8;
      break;
    case NPY_UINT16:
      *out_tf_datatype = TF_UINT16;
      break;
    case NPY_UINT32:
      *out_tf_datatype = TF_UINT32;
      break;
    case NPY_UINT64:
      *out_tf_datatype = TF_UINT64;
      break;
    case NPY_INT8:
      *out_tf_datatype = TF_INT8;
      break;
    case NPY_INT16:
      *out_tf_datatype = TF_INT16;
      break;
    case NPY_INT64:
      *out_tf_datatype = TF_INT64;
      break;
    case NPY_BOOL:
      *out_tf_datatype = TF_BOOL;
      break;
    case NPY_COMPLEX64:
      *out_tf_datatype = TF_COMPLEX64;
      break;
    case NPY_COMPLEX128:
      *out_tf_datatype = TF_COMPLEX128;
      break;
    case NPY_OBJECT:
    case NPY_STRING:
    case NPY_UNICODE:
      *out_tf_datatype = TF_STRING;
      break;
    case NPY_VOID:
      // Quantized types are currently represented as custom struct types.
      // PyArray_TYPE returns NPY_VOID for structs, and we should look into
      // descr to derive the actual type.
      // Direct feeds of certain types of ResourceHandles are represented as a
      // custom struct type.
      return PyArrayDescr_to_TF_DataType(descr, out_tf_datatype);
    default:
      if (pyarray_type == Bfloat16NumpyType()) {
        *out_tf_datatype = TF_BFLOAT16;
        break;
      } else if (pyarray_type == NPY_ULONGLONG) {
        // NPY_ULONGLONG is equivalent to NPY_UINT64, while their enum values
        // might be different on certain platforms.
        *out_tf_datatype = TF_UINT64;
        break;
      } else if (pyarray_type == NPY_LONGLONG) {
        // NPY_LONGLONG is equivalent to NPY_INT64, while their enum values
        // might be different on certain platforms.
        *out_tf_datatype = TF_INT64;
        break;
      } else if (pyarray_type == NPY_INT) {
        // NPY_INT is equivalent to NPY_INT32, while their enum values might be
        // different on certain platforms.
        *out_tf_datatype = TF_INT32;
        break;
      } else if (pyarray_type == NPY_UINT) {
        // NPY_UINT is equivalent to NPY_UINT32, while their enum values might
        // be different on certain platforms.
        *out_tf_datatype = TF_UINT32;
        break;
      }
      return errors::Internal("Unsupported numpy type: ",
                              numpy_type_name(pyarray_type));
  }
  return Status::OK();
}


char const* numpy_type_name(int numpy_type) {
  switch (numpy_type) {
#define TYPE_CASE(s) \
  case s:            \
    return #s

    TYPE_CASE(NPY_BOOL);
    TYPE_CASE(NPY_BYTE);
    TYPE_CASE(NPY_UBYTE);
    TYPE_CASE(NPY_SHORT);
    TYPE_CASE(NPY_USHORT);
    TYPE_CASE(NPY_INT);
    TYPE_CASE(NPY_UINT);
    TYPE_CASE(NPY_LONG);
    TYPE_CASE(NPY_ULONG);
    TYPE_CASE(NPY_LONGLONG);
    TYPE_CASE(NPY_ULONGLONG);
    TYPE_CASE(NPY_FLOAT);
    TYPE_CASE(NPY_DOUBLE);
    TYPE_CASE(NPY_LONGDOUBLE);
    TYPE_CASE(NPY_CFLOAT);
    TYPE_CASE(NPY_CDOUBLE);
    TYPE_CASE(NPY_CLONGDOUBLE);
    TYPE_CASE(NPY_OBJECT);
    TYPE_CASE(NPY_STRING);
    TYPE_CASE(NPY_UNICODE);
    TYPE_CASE(NPY_VOID);
    TYPE_CASE(NPY_DATETIME);
    TYPE_CASE(NPY_TIMEDELTA);
    TYPE_CASE(NPY_HALF);
    TYPE_CASE(NPY_NTYPES);
    TYPE_CASE(NPY_NOTYPE);
    TYPE_CASE(NPY_CHAR);
    TYPE_CASE(NPY_USERDEF);
    default:
      return "not a numpy type";
  }
}


Status PyBytesArrayMap(PyArrayObject* array, F f) {
  Safe_PyObjectPtr iter = tensorflow::make_safe(
      PyArray_IterNew(reinterpret_cast<PyObject*>(array)));
  while (PyArray_ITER_NOTDONE(iter.get())) {
    auto item = tensorflow::make_safe(PyArray_GETITEM(
        array, static_cast<char*>(PyArray_ITER_DATA(iter.get()))));
    if (!item) {
      return errors::Internal("Unable to get element from the feed - no item.");
    }
    Py_ssize_t len;
    const char* ptr;
    PyObject* ptr_owner = nullptr;
    TF_RETURN_IF_ERROR(PyObjectToString(item.get(), &ptr, &len, &ptr_owner));
    f(ptr, len);
    Py_XDECREF(ptr_owner);
    PyArray_ITER_NEXT(iter.get());
  }
  return Status::OK();
}


Status TensorToNdarray(const Tensor& t, PyObject** ret) {
  Status status;
  Safe_TF_TensorPtr tf_tensor = make_safe(TF_TensorFromTensor(t, &status));
  if (!status.ok()) {
    return status;
  }
  return TF_TensorToPyArray(std::move(tf_tensor), ret);
}


Status EncodePyBytesArray(PyArrayObject* array, tensorflow::int64 nelems,
                          size_t* size, void** buffer) {
  // Encode all strings.
  *size = nelems * sizeof(tensorflow::tstring);
  std::unique_ptr<tensorflow::tstring[]> base_ptr(
      new tensorflow::tstring[nelems]);
  tensorflow::tstring* dst = base_ptr.get();

  TF_RETURN_IF_ERROR(
      PyBytesArrayMap(array, [&dst](const char* ptr, Py_ssize_t len) {
        dst->assign(ptr, len);
        dst++;
      }));
  *buffer = base_ptr.release();
  return Status::OK();
}


std::vector<std::vector<float>> DefaultColorTable(int depth) {
  std::vector<std::vector<float>> color_table;
  color_table.emplace_back(std::vector<float>({1, 1, 0, 1}));      // 0: yellow
  color_table.emplace_back(std::vector<float>({0, 0, 1, 1}));      // 1: blue
  color_table.emplace_back(std::vector<float>({1, 0, 0, 1}));      // 2: red
  color_table.emplace_back(std::vector<float>({0, 1, 0, 1}));      // 3: lime
  color_table.emplace_back(std::vector<float>({0.5, 0, 0.5, 1}));  // 4: purple
  color_table.emplace_back(std::vector<float>({0.5, 0.5, 0, 1}));  // 5: olive
  color_table.emplace_back(std::vector<float>({0.5, 0, 0, 1}));    // 6: maroon
  color_table.emplace_back(std::vector<float>({0, 0, 0.5, 1}));  // 7: navy blue
  color_table.emplace_back(std::vector<float>({0, 1, 1, 1}));    // 8: aqua
  color_table.emplace_back(std::vector<float>({1, 0, 1, 1}));    // 9: fuchsia

  if (depth == 1) {
    for (int64 i = 0; i < color_table.size(); i++) {
      color_table[i][0] = 1;
    }
  }
  return color_table;
}

Fix overflow CHECK issue with `tf.raw_ops.DrawBoundingBoxes`.

PiperOrigin-RevId: 369753591
Change-Id: I3b45fc98ee0d28a3c20b7e9c995aa647c976ec40


  explicit DrawBoundingBoxesOp(OpKernelConstruction* context)
      : OpKernel(context) {}


virtio_pci_modern_pio_read(struct vmctx *ctx, int vcpu, struct pci_vdev *dev,
			   int baridx, uint64_t offset, int size)
{
	/* guest driver should not read notify pio */
	return 0;
}

dm: validate inputs in vq_endchains

 inputs shall be validated to avoid NULL pointer access.

Tracked-On: #6129
Signed-off-by: Yonghua Huang <yonghua.huang@intel.com>

struct config_reg {
	uint16_t	offset;	/* register offset */
	uint8_t		size;	/* size (bytes) */
	uint8_t		ro;	/* true => reg is read only */
	const char	*name;	/* name of reg */
}

static struct config_reg legacy_config_regs[] = {
	{ VIRTIO_PCI_HOST_FEATURES,	4, 1, "HOSTCAP" },
	{ VIRTIO_PCI_GUEST_FEATURES,	4, 0, "GUESTCAP" },
	{ VIRTIO_PCI_QUEUE_PFN,		4, 0, "PFN" },
	{ VIRTIO_PCI_QUEUE_NUM,		2, 1, "QNUM" },
	{ VIRTIO_PCI_QUEUE_SEL,		2, 0, "QSEL" },
	{ VIRTIO_PCI_QUEUE_NOTIFY,	2, 0, "QNOTIFY" },
	{ VIRTIO_PCI_STATUS,		1, 0, "STATUS" },
	{ VIRTIO_PCI_ISR,		1, 0, "ISR" },
	{ VIRTIO_MSI_CONFIG_VECTOR,	2, 0, "CFGVEC" },
	{ VIRTIO_MSI_QUEUE_VECTOR,	2, 0, "QVEC" },
};

static struct config_reg modern_config_regs[] = {
	{ VIRTIO_PCI_COMMON_DFSELECT,		4, 0, "DFSELECT" },
	{ VIRTIO_PCI_COMMON_DF,			4, 1, "DF" },
	{ VIRTIO_PCI_COMMON_GFSELECT,		4, 0, "GFSELECT" },
	{ VIRTIO_PCI_COMMON_GF,			4, 0, "GF" },
	{ VIRTIO_PCI_COMMON_MSIX,		2, 0, "MSIX" },
	{ VIRTIO_PCI_COMMON_NUMQ,		2, 1, "NUMQ" },
	{ VIRTIO_PCI_COMMON_STATUS,		1, 0, "STATUS" },
	{ VIRTIO_PCI_COMMON_CFGGENERATION,	1, 1, "CFGGENERATION" },
	{ VIRTIO_PCI_COMMON_Q_SELECT,		2, 0, "Q_SELECT" },
	{ VIRTIO_PCI_COMMON_Q_SIZE,		2, 0, "Q_SIZE" },
	{ VIRTIO_PCI_COMMON_Q_MSIX,		2, 0, "Q_MSIX" },
	{ VIRTIO_PCI_COMMON_Q_ENABLE,		2, 0, "Q_ENABLE" },
	{ VIRTIO_PCI_COMMON_Q_NOFF,		2, 1, "Q_NOFF" },
	{ VIRTIO_PCI_COMMON_Q_DESCLO,		4, 0, "Q_DESCLO" },
	{ VIRTIO_PCI_COMMON_Q_DESCHI,		4, 0, "Q_DESCHI" },
	{ VIRTIO_PCI_COMMON_Q_AVAILLO,		4, 0, "Q_AVAILLO" },
	{ VIRTIO_PCI_COMMON_Q_AVAILHI,		4, 0, "Q_AVAILHI" },
	{ VIRTIO_PCI_COMMON_Q_USEDLO,		4, 0, "Q_USEDLO" },
	{ VIRTIO_PCI_COMMON_Q_USEDHI,		4, 0, "Q_USEDHI" },
};


virtio_pci_read(struct vmctx *ctx, int vcpu, struct pci_vdev *dev,
		int baridx, uint64_t offset, int size)
{
	struct virtio_base *base = dev->arg;

	if (base->flags & VIRTIO_USE_MSIX) {
		if (baridx == pci_msix_table_bar(dev) ||
		    baridx == pci_msix_pba_bar(dev)) {
			return pci_emul_msix_tread(dev, offset, size);
		}
	}

	if (baridx == base->legacy_pio_bar_idx)
		return virtio_pci_legacy_read(ctx, vcpu, dev, baridx,
			offset, size);

	if (baridx == base->modern_mmio_bar_idx)
		return virtio_pci_modern_mmio_read(ctx, vcpu, dev, baridx,
			offset, size);

	if (baridx == base->modern_pio_bar_idx)
		return virtio_pci_modern_pio_read(ctx, vcpu, dev, baridx,
			offset, size);

	pr_err("%s: read unexpected baridx %d\r\n",
		base->vops->name, baridx);
	return size == 1 ? 0xff : size == 2 ? 0xffff : 0xffffffff;
}


virtio_find_cr(const struct config_reg *p_cr_array, u_int array_size,
	       int offset) {
	u_int hi, lo, mid;
	const struct config_reg *cr;

	lo = 0;
	hi = array_size - 1;
	while (hi >= lo) {
		mid = (hi + lo) >> 1;
		cr = p_cr_array + mid;
		if (cr->offset == offset)
			return cr;
		if (cr->offset < offset)
			lo = mid + 1;
		else
			hi = mid - 1;
	}
	return NULL;
}


virtio_poll_timer(void *arg, uint64_t nexp)
{
	struct virtio_base *base;
	struct virtio_ops *vops;
	struct virtio_vq_info *vq;
	const char *name;
	int i;

	base = arg;
	vops = base->vops;
	name = vops->name;

	if (base->mtx)
		pthread_mutex_lock(base->mtx);

	base->polling_in_progress = 1;

	for (i = 0; i < base->vops->nvq; i++) {
		vq = &base->queues[i];
		if(!vq_ring_ready(vq))
			continue;
		vq->used->flags |= VRING_USED_F_NO_NOTIFY;
		/* TODO: call notify when necessary */
		if (vq->notify)
			(*vq->notify)(DEV_STRUCT(base), vq);
		else if (vops->qnotify)
			(*vops->qnotify)(DEV_STRUCT(base), vq);
		else
			pr_err("%s: qnotify queue %d: missing vq/vops notify\r\n",
				name, i);
	}

	if (base->mtx)
		pthread_mutex_unlock(base->mtx);

	virtio_start_timer(&base->polling_timer, 0, virtio_poll_interval);
}

#define DEV_STRUCT(vs) ((void *)(vs))

static size_t virtio_poll_interval;

struct config_reg {
	uint16_t	offset;	/* register offset */
	uint8_t		size;	/* size (bytes) */
	uint8_t		ro;	/* true => reg is read only */
	const char	*name;	/* name of reg */
}

static struct config_reg legacy_config_regs[] = {
	{ VIRTIO_PCI_HOST_FEATURES,	4, 1, "HOSTCAP" },
	{ VIRTIO_PCI_GUEST_FEATURES,	4, 0, "GUESTCAP" },
	{ VIRTIO_PCI_QUEUE_PFN,		4, 0, "PFN" },
	{ VIRTIO_PCI_QUEUE_NUM,		2, 1, "QNUM" },
	{ VIRTIO_PCI_QUEUE_SEL,		2, 0, "QSEL" },
	{ VIRTIO_PCI_QUEUE_NOTIFY,	2, 0, "QNOTIFY" },
	{ VIRTIO_PCI_STATUS,		1, 0, "STATUS" },
	{ VIRTIO_PCI_ISR,		1, 0, "ISR" },
	{ VIRTIO_MSI_CONFIG_VECTOR,	2, 0, "CFGVEC" },
	{ VIRTIO_MSI_QUEUE_VECTOR,	2, 0, "QVEC" },
};

static struct config_reg modern_config_regs[] = {
	{ VIRTIO_PCI_COMMON_DFSELECT,		4, 0, "DFSELECT" },
	{ VIRTIO_PCI_COMMON_DF,			4, 1, "DF" },
	{ VIRTIO_PCI_COMMON_GFSELECT,		4, 0, "GFSELECT" },
	{ VIRTIO_PCI_COMMON_GF,			4, 0, "GF" },
	{ VIRTIO_PCI_COMMON_MSIX,		2, 0, "MSIX" },
	{ VIRTIO_PCI_COMMON_NUMQ,		2, 1, "NUMQ" },
	{ VIRTIO_PCI_COMMON_STATUS,		1, 0, "STATUS" },
	{ VIRTIO_PCI_COMMON_CFGGENERATION,	1, 1, "CFGGENERATION" },
	{ VIRTIO_PCI_COMMON_Q_SELECT,		2, 0, "Q_SELECT" },
	{ VIRTIO_PCI_COMMON_Q_SIZE,		2, 0, "Q_SIZE" },
	{ VIRTIO_PCI_COMMON_Q_MSIX,		2, 0, "Q_MSIX" },
	{ VIRTIO_PCI_COMMON_Q_ENABLE,		2, 0, "Q_ENABLE" },
	{ VIRTIO_PCI_COMMON_Q_NOFF,		2, 1, "Q_NOFF" },
	{ VIRTIO_PCI_COMMON_Q_DESCLO,		4, 0, "Q_DESCLO" },
	{ VIRTIO_PCI_COMMON_Q_DESCHI,		4, 0, "Q_DESCHI" },
	{ VIRTIO_PCI_COMMON_Q_AVAILLO,		4, 0, "Q_AVAILLO" },
	{ VIRTIO_PCI_COMMON_Q_AVAILHI,		4, 0, "Q_AVAILHI" },
	{ VIRTIO_PCI_COMMON_Q_USEDLO,		4, 0, "Q_USEDLO" },
	{ VIRTIO_PCI_COMMON_Q_USEDHI,		4, 0, "Q_USEDHI" },
};


_vq_record(int i, volatile struct vring_desc *vd, struct vmctx *ctx,
	   struct iovec *iov, int n_iov, uint16_t *flags) {

	void *host_addr;

	if (i >= n_iov)
		return -1;
	host_addr = paddr_guest2host(ctx, vd->addr, vd->len);
	if (!host_addr)
		return -1;
	iov[i].iov_base = host_addr;
	iov[i].iov_len = vd->len;
	if (flags != NULL)
		flags[i] = vd->flags;
	return 0;
}

#define DEV_STRUCT(vs) ((void *)(vs))

struct config_reg {
	uint16_t	offset;	/* register offset */
	uint8_t		size;	/* size (bytes) */
	uint8_t		ro;	/* true => reg is read only */
	const char	*name;	/* name of reg */
}

static struct config_reg legacy_config_regs[] = {
	{ VIRTIO_PCI_HOST_FEATURES,	4, 1, "HOSTCAP" },
	{ VIRTIO_PCI_GUEST_FEATURES,	4, 0, "GUESTCAP" },
	{ VIRTIO_PCI_QUEUE_PFN,		4, 0, "PFN" },
	{ VIRTIO_PCI_QUEUE_NUM,		2, 1, "QNUM" },
	{ VIRTIO_PCI_QUEUE_SEL,		2, 0, "QSEL" },
	{ VIRTIO_PCI_QUEUE_NOTIFY,	2, 0, "QNOTIFY" },
	{ VIRTIO_PCI_STATUS,		1, 0, "STATUS" },
	{ VIRTIO_PCI_ISR,		1, 0, "ISR" },
	{ VIRTIO_MSI_CONFIG_VECTOR,	2, 0, "CFGVEC" },
	{ VIRTIO_MSI_QUEUE_VECTOR,	2, 0, "QVEC" },
};

static struct config_reg modern_config_regs[] = {
	{ VIRTIO_PCI_COMMON_DFSELECT,		4, 0, "DFSELECT" },
	{ VIRTIO_PCI_COMMON_DF,			4, 1, "DF" },
	{ VIRTIO_PCI_COMMON_GFSELECT,		4, 0, "GFSELECT" },
	{ VIRTIO_PCI_COMMON_GF,			4, 0, "GF" },
	{ VIRTIO_PCI_COMMON_MSIX,		2, 0, "MSIX" },
	{ VIRTIO_PCI_COMMON_NUMQ,		2, 1, "NUMQ" },
	{ VIRTIO_PCI_COMMON_STATUS,		1, 0, "STATUS" },
	{ VIRTIO_PCI_COMMON_CFGGENERATION,	1, 1, "CFGGENERATION" },
	{ VIRTIO_PCI_COMMON_Q_SELECT,		2, 0, "Q_SELECT" },
	{ VIRTIO_PCI_COMMON_Q_SIZE,		2, 0, "Q_SIZE" },
	{ VIRTIO_PCI_COMMON_Q_MSIX,		2, 0, "Q_MSIX" },
	{ VIRTIO_PCI_COMMON_Q_ENABLE,		2, 0, "Q_ENABLE" },
	{ VIRTIO_PCI_COMMON_Q_NOFF,		2, 1, "Q_NOFF" },
	{ VIRTIO_PCI_COMMON_Q_DESCLO,		4, 0, "Q_DESCLO" },
	{ VIRTIO_PCI_COMMON_Q_DESCHI,		4, 0, "Q_DESCHI" },
	{ VIRTIO_PCI_COMMON_Q_AVAILLO,		4, 0, "Q_AVAILLO" },
	{ VIRTIO_PCI_COMMON_Q_AVAILHI,		4, 0, "Q_AVAILHI" },
	{ VIRTIO_PCI_COMMON_Q_USEDLO,		4, 0, "Q_USEDLO" },
	{ VIRTIO_PCI_COMMON_Q_USEDHI,		4, 0, "Q_USEDHI" },
};


virtio_get_cap_id(uint64_t offset, int size)
{
	int i, rc = -1;

	for (i = 0; i < ARRAY_SIZE(cap_regions); i++) {
		if (offset >= cap_regions[i].cap_offset &&
			offset + size <= cap_regions[i].cap_offset +
			cap_regions[i].cap_size)
			return cap_regions[i].cap_id;
	}

	return rc;
}

struct config_reg {
	uint16_t	offset;	/* register offset */
	uint8_t		size;	/* size (bytes) */
	uint8_t		ro;	/* true => reg is read only */
	const char	*name;	/* name of reg */
}

static struct cap_region {
	uint64_t	cap_offset;	/* offset of capability region */
	int		cap_size;	/* size of capability region */
	int		cap_id;		/* capability id */
} cap_regions[] = {
	{VIRTIO_CAP_COMMON_OFFSET, VIRTIO_CAP_COMMON_SIZE,
		VIRTIO_PCI_CAP_COMMON_CFG},
	{VIRTIO_CAP_ISR_OFFSET, VIRTIO_CAP_ISR_SIZE,
		VIRTIO_PCI_CAP_ISR_CFG},
	{VIRTIO_CAP_DEVICE_OFFSET, VIRTIO_CAP_DEVICE_SIZE,
		VIRTIO_PCI_CAP_DEVICE_CFG},
	{VIRTIO_CAP_NOTIFY_OFFSET, VIRTIO_CAP_NOTIFY_SIZE,
		VIRTIO_PCI_CAP_NOTIFY_CFG},
};


virtio_pci_modern_mmio_write(struct vmctx *ctx, int vcpu, struct pci_vdev *dev,
			     int baridx, uint64_t offset, int size,
			     uint64_t value)
{
	struct virtio_base *base = dev->arg;
	struct virtio_ops *vops;
	const char *name;
	int capid;

	vops = base->vops;
	name = vops->name;

	if (size != 1 && size != 2 && size != 4) {
		pr_err("%s: write to [%d:0x%lx] bad size %d\r\n",
			name, baridx, offset, size);
		return;
	}

	capid = virtio_get_cap_id(offset, size);
	if (capid < 0) {
		pr_err("%s: write to [%d:0x%lx] bad range %d\r\n",
			name, baridx, offset, size);
		return;
	}

	if (base->mtx)
		pthread_mutex_lock(base->mtx);

	switch (capid) {
	case VIRTIO_PCI_CAP_COMMON_CFG:
		offset -= VIRTIO_CAP_COMMON_OFFSET;
		virtio_common_cfg_write(dev, offset, size, value);
		break;
	case VIRTIO_PCI_CAP_DEVICE_CFG:
		offset -= VIRTIO_CAP_DEVICE_OFFSET;
		virtio_device_cfg_write(dev, offset, size, value);
		break;
	case VIRTIO_PCI_CAP_NOTIFY_CFG:
		offset -= VIRTIO_CAP_NOTIFY_OFFSET;
		virtio_notify_cfg_write(dev, offset, size, value);
		break;
	default: /* guest driver should not write to ISR region */
		pr_err("%s: write to [%d:0x%lx] size %d not supported\r\n",
			name, baridx, offset, size);
	}

	if (base->mtx)
		pthread_mutex_unlock(base->mtx);
}


virtio_set_modern_bar(struct virtio_base *base, bool use_notify_pio)
{
	struct virtio_ops *vops;
	int rc = 0;

	vops = base->vops;

	if (!vops || (base->device_caps & (1UL << VIRTIO_F_VERSION_1)) == 0)
		return -1;

	if (use_notify_pio)
		rc = virtio_set_modern_pio_bar(base,
			VIRTIO_MODERN_PIO_BAR_IDX);
	if (!rc)
		rc = virtio_set_modern_mmio_bar(base,
			VIRTIO_MODERN_MMIO_BAR_IDX);

	return rc;
}


virtio_pci_write(struct vmctx *ctx, int vcpu, struct pci_vdev *dev,
		 int baridx, uint64_t offset, int size, uint64_t value)
{
	struct virtio_base *base = dev->arg;

	if (base->flags & VIRTIO_USE_MSIX) {
		if (baridx == pci_msix_table_bar(dev) ||
		    baridx == pci_msix_pba_bar(dev)) {
			pci_emul_msix_twrite(dev, offset, size, value);
			return;
		}
	}

	if (baridx == base->legacy_pio_bar_idx) {
		virtio_pci_legacy_write(ctx, vcpu, dev, baridx,
			offset, size, value);
		return;
	}

	if (baridx == base->modern_mmio_bar_idx) {
		virtio_pci_modern_mmio_write(ctx, vcpu, dev, baridx,
			offset, size, value);
		return;
	}

	if (baridx == base->modern_pio_bar_idx) {
		virtio_pci_modern_pio_write(ctx, vcpu, dev, baridx,
			offset, size, value);
		return;
	}

	pr_err("%s: write unexpected baridx %d\r\n",
		base->vops->name, baridx);
}


vq_getchain(struct virtio_vq_info *vq, uint16_t *pidx,
	    struct iovec *iov, int n_iov, uint16_t *flags)
{
	int i;
	u_int ndesc, n_indir;
	u_int idx, next;

	volatile struct vring_desc *vdir, *vindir, *vp;
	struct vmctx *ctx;
	struct virtio_base *base;
	const char *name;

	base = vq->base;
	name = base->vops->name;

	/*
	 * Note: it's the responsibility of the guest not to
	 * update vq->avail->idx until all of the descriptors
	 * the guest has written are valid (including all their
	 * next fields and vd_flags).
	 *
	 * Compute (last_avail - idx) in integers mod 2**16.  This is
	 * the number of descriptors the device has made available
	 * since the last time we updated vq->last_avail.
	 *
	 * We just need to do the subtraction as an unsigned int,
	 * then trim off excess bits.
	 */
	idx = vq->last_avail;
	ndesc = (uint16_t)((u_int)vq->avail->idx - idx);
	if (ndesc == 0)
		return 0;
	if (ndesc > vq->qsize) {
		/* XXX need better way to diagnose issues */
		pr_err("%s: ndesc (%u) out of range, driver confused?\r\n",
		    name, (u_int)ndesc);
		return -1;
	}

	/*
	 * Now count/parse "involved" descriptors starting from
	 * the head of the chain.
	 *
	 * To prevent loops, we could be more complicated and
	 * check whether we're re-visiting a previously visited
	 * index, but we just abort if the count gets excessive.
	 */
	ctx = base->dev->vmctx;
	*pidx = next = vq->avail->ring[idx & (vq->qsize - 1)];
	vq->last_avail++;
	for (i = 0; i < VQ_MAX_DESCRIPTORS; next = vdir->next) {
		if (next >= vq->qsize) {
			pr_err("%s: descriptor index %u out of range, "
			    "driver confused?\r\n",
			    name, next);
			return -1;
		}
		vdir = &vq->desc[next];
		if ((vdir->flags & VRING_DESC_F_INDIRECT) == 0) {
			if (_vq_record(i, vdir, ctx, iov, n_iov, flags)) {
				pr_err("%s: mapping to host failed\r\n", name);
				return -1;
			}
			i++;
		} else if ((base->device_caps &
		    (1 << VIRTIO_RING_F_INDIRECT_DESC)) == 0) {
			pr_err("%s: descriptor has forbidden INDIRECT flag, "
			    "driver confused?\r\n",
			    name);
			return -1;
		} else {
			n_indir = vdir->len / 16;
			if ((vdir->len & 0xf) || n_indir == 0) {
				pr_err("%s: invalid indir len 0x%x, "
				    "driver confused?\r\n",
				    name, (u_int)vdir->len);
				return -1;
			}
			vindir = paddr_guest2host(ctx,
			    vdir->addr, vdir->len);

			if (!vindir) {
				pr_err("%s cannot get host memory\r\n", name);
				return -1;
			}
			/*
			 * Indirects start at the 0th, then follow
			 * their own embedded "next"s until those run
			 * out.  Each one's indirect flag must be off
			 * (we don't really have to check, could just
			 * ignore errors...).
			 */
			next = 0;
			for (;;) {
				vp = &vindir[next];
				if (vp->flags & VRING_DESC_F_INDIRECT) {
					pr_err("%s: indirect desc has INDIR flag,"
					    " driver confused?\r\n",
					    name);
					return -1;
				}
				if (_vq_record(i, vp, ctx, iov, n_iov, flags)) {
					pr_err("%s: mapping to host failed\r\n", name);
					return -1;
				}
				if (++i > VQ_MAX_DESCRIPTORS)
					goto loopy;
				if ((vp->flags & VRING_DESC_F_NEXT) == 0)
					break;
				next = vp->next;
				if (next >= n_indir) {
					pr_err("%s: invalid next %u > %u, "
					    "driver confused?\r\n",
					    name, (u_int)next, n_indir);
					return -1;
				}
			}
		}
		if ((vdir->flags & VRING_DESC_F_NEXT) == 0)
			return i;
	}
loopy:
	pr_err("%s: descriptor loop? count > %d - driver confused?\r\n",
	    name, i);
	return -1;
}

#define	VQ_MAX_DESCRIPTORS	512	/* see below */

struct config_reg {
	uint16_t	offset;	/* register offset */
	uint8_t		size;	/* size (bytes) */
	uint8_t		ro;	/* true => reg is read only */
	const char	*name;	/* name of reg */
}

static struct config_reg legacy_config_regs[] = {
	{ VIRTIO_PCI_HOST_FEATURES,	4, 1, "HOSTCAP" },
	{ VIRTIO_PCI_GUEST_FEATURES,	4, 0, "GUESTCAP" },
	{ VIRTIO_PCI_QUEUE_PFN,		4, 0, "PFN" },
	{ VIRTIO_PCI_QUEUE_NUM,		2, 1, "QNUM" },
	{ VIRTIO_PCI_QUEUE_SEL,		2, 0, "QSEL" },
	{ VIRTIO_PCI_QUEUE_NOTIFY,	2, 0, "QNOTIFY" },
	{ VIRTIO_PCI_STATUS,		1, 0, "STATUS" },
	{ VIRTIO_PCI_ISR,		1, 0, "ISR" },
	{ VIRTIO_MSI_CONFIG_VECTOR,	2, 0, "CFGVEC" },
	{ VIRTIO_MSI_QUEUE_VECTOR,	2, 0, "QVEC" },
};

static struct config_reg modern_config_regs[] = {
	{ VIRTIO_PCI_COMMON_DFSELECT,		4, 0, "DFSELECT" },
	{ VIRTIO_PCI_COMMON_DF,			4, 1, "DF" },
	{ VIRTIO_PCI_COMMON_GFSELECT,		4, 0, "GFSELECT" },
	{ VIRTIO_PCI_COMMON_GF,			4, 0, "GF" },
	{ VIRTIO_PCI_COMMON_MSIX,		2, 0, "MSIX" },
	{ VIRTIO_PCI_COMMON_NUMQ,		2, 1, "NUMQ" },
	{ VIRTIO_PCI_COMMON_STATUS,		1, 0, "STATUS" },
	{ VIRTIO_PCI_COMMON_CFGGENERATION,	1, 1, "CFGGENERATION" },
	{ VIRTIO_PCI_COMMON_Q_SELECT,		2, 0, "Q_SELECT" },
	{ VIRTIO_PCI_COMMON_Q_SIZE,		2, 0, "Q_SIZE" },
	{ VIRTIO_PCI_COMMON_Q_MSIX,		2, 0, "Q_MSIX" },
	{ VIRTIO_PCI_COMMON_Q_ENABLE,		2, 0, "Q_ENABLE" },
	{ VIRTIO_PCI_COMMON_Q_NOFF,		2, 1, "Q_NOFF" },
	{ VIRTIO_PCI_COMMON_Q_DESCLO,		4, 0, "Q_DESCLO" },
	{ VIRTIO_PCI_COMMON_Q_DESCHI,		4, 0, "Q_DESCHI" },
	{ VIRTIO_PCI_COMMON_Q_AVAILLO,		4, 0, "Q_AVAILLO" },
	{ VIRTIO_PCI_COMMON_Q_AVAILHI,		4, 0, "Q_AVAILHI" },
	{ VIRTIO_PCI_COMMON_Q_USEDLO,		4, 0, "Q_USEDLO" },
	{ VIRTIO_PCI_COMMON_Q_USEDHI,		4, 0, "Q_USEDHI" },
};


virtio_intr_init(struct virtio_base *base, int barnum, int use_msix)
{
	int nvec;

	if (use_msix) {
		base->flags |= VIRTIO_USE_MSIX;
		VIRTIO_BASE_LOCK(base);
		virtio_reset_dev(base); /* set all vectors to NO_VECTOR */
		VIRTIO_BASE_UNLOCK(base);
		nvec = base->vops->nvq + 1;
		if (pci_emul_add_msixcap(base->dev, nvec, barnum))
			return -1;
	} else
		base->flags &= ~VIRTIO_USE_MSIX;

	/* Only 1 MSI vector for acrn-dm */
	pci_emul_add_msicap(base->dev, 1);

	/* Legacy interrupts are mandatory for virtio devices */
	pci_lintr_request(base->dev);

	return 0;
}


virtio_device_cfg_read(struct pci_vdev *dev, uint64_t offset, int size)
{
	struct virtio_base *base = dev->arg;
	struct virtio_ops *vops;
	const char *name;
	uint32_t value;
	uint64_t max;
	int error = -1;

	vops = base->vops;
	name = vops->name;
	value = size == 1 ? 0xff : size == 2 ? 0xffff : 0xffffffff;
	max = vops->cfgsize ? vops->cfgsize : 0x100000000;

	if (offset + size > max) {
		pr_err("%s: reading from 0x%lx size %d exceeds limit\r\n",
			name, offset, size);
		return value;
	}

	if (vops->cfgread) {
		error = (*vops->cfgread)(DEV_STRUCT(base), offset, size, &value);
	}
	if (error) {
		pr_err("%s: reading from 0x%lx size %d failed %d\r\n",
			name, offset, size, error);
		value = size == 1 ? 0xff : size == 2 ? 0xffff : 0xffffffff;
	}

	return value;
}


acrn_parse_virtio_poll_interval(const char *optarg)
{
	char *ptr;

	virtio_poll_interval = strtoul(optarg, &ptr, 0);

	/* poll interval is limited from 1us to 10ms */
	if (virtio_poll_interval < 1 || virtio_poll_interval > 10000000)
		return -1;

	virtio_poll_enabled = 1;

	return 0;
}

static uint8_t virtio_poll_enabled;

static size_t virtio_poll_interval;

struct config_reg {
	uint16_t	offset;	/* register offset */
	uint8_t		size;	/* size (bytes) */
	uint8_t		ro;	/* true => reg is read only */
	const char	*name;	/* name of reg */
}


virtio_pci_legacy_read(struct vmctx *ctx, int vcpu, struct pci_vdev *dev,
		       int baridx, uint64_t offset, int size)
{
	struct virtio_base *base = dev->arg;
	struct virtio_ops *vops;
	const struct config_reg *cr;
	uint64_t virtio_config_size, max;
	const char *name;
	uint32_t newoff;
	uint32_t value;
	int error = -1;

	if (base->mtx)
		pthread_mutex_lock(base->mtx);

	vops = base->vops;
	name = vops->name;
	value = size == 1 ? 0xff : size == 2 ? 0xffff : 0xffffffff;

	if (size != 1 && size != 2 && size != 4)
		goto bad;

	if (pci_msix_enabled(dev))
		virtio_config_size = VIRTIO_PCI_CONFIG_OFF(1);
	else
		virtio_config_size = VIRTIO_PCI_CONFIG_OFF(0);

	if (offset >= virtio_config_size) {
		/*
		 * Subtract off the standard size (including MSI-X
		 * registers if enabled) and dispatch to underlying driver.
		 * If that fails, fall into general code.
		 */
		newoff = offset - virtio_config_size;
		max = vops->cfgsize ? vops->cfgsize : 0x100000000;
		if (newoff + size > max)
			goto bad;

		if (vops->cfgread) {
			error = (*vops->cfgread)(DEV_STRUCT(base), newoff,
						 size, &value);
		}
		if (!error)
			goto done;
	}

bad:
	cr = virtio_find_legacy_cr(offset);
	if (cr == NULL || cr->size != size) {
		if (cr != NULL) {
			/* offset must be OK, so size must be bad */
			pr_err("%s: read from %s: bad size %d\r\n",
			    name, cr->name, size);
		} else {
			pr_err("%s: read from bad offset/size %jd/%d\r\n",
			    name, (uintmax_t)offset, size);
		}
		goto done;
	}

	switch (offset) {
	case VIRTIO_PCI_HOST_FEATURES:
		value = base->device_caps;
		break;
	case VIRTIO_PCI_GUEST_FEATURES:
		value = base->negotiated_caps;
		break;
	case VIRTIO_PCI_QUEUE_PFN:
		if (base->curq < vops->nvq)
			value = base->queues[base->curq].pfn;
		break;
	case VIRTIO_PCI_QUEUE_NUM:
		value = base->curq < vops->nvq ?
			base->queues[base->curq].qsize : 0;
		break;
	case VIRTIO_PCI_QUEUE_SEL:
		value = base->curq;
		break;
	case VIRTIO_PCI_QUEUE_NOTIFY:
		value = 0;	/* XXX */
		break;
	case VIRTIO_PCI_STATUS:
		value = base->status;
		break;
	case VIRTIO_PCI_ISR:
		value = base->isr;
		base->isr = 0;		/* a read clears this flag */
		if (value)
			pci_lintr_deassert(dev);
		break;
	case VIRTIO_MSI_CONFIG_VECTOR:
		value = base->msix_cfg_idx;
		break;
	case VIRTIO_MSI_QUEUE_VECTOR:
		value = base->curq < vops->nvq ?
		    base->queues[base->curq].msix_idx :
		    VIRTIO_MSI_NO_VECTOR;
		break;
	}
done:
	if (base->mtx)
		pthread_mutex_unlock(base->mtx);
	return value;
}


virtio_pci_modern_pio_write(struct vmctx *ctx, int vcpu, struct pci_vdev *dev,
			    int baridx, uint64_t offset, int size,
			    uint64_t value)
{
	struct virtio_base *base = dev->arg;
	struct virtio_vq_info *vq;
	struct virtio_ops *vops;
	const char *name;
	uint64_t idx;

	vops = base->vops;
	name = vops->name;
	idx = value;

	if (size != 1 && size != 2 && size != 4) {
		pr_err("%s: write to [%d:0x%lx] bad size %d\r\n",
			name, baridx, offset, size);
		return;
	}

	if (idx >= vops->nvq) {
		pr_err("%s: queue %lu notify out of range\r\n", name, idx);
		return;
	}

	if (base->mtx)
		pthread_mutex_lock(base->mtx);

	vq = &base->queues[idx];
	if (vq->notify)
		(*vq->notify)(DEV_STRUCT(base), vq);
	else if (vops->qnotify)
		(*vops->qnotify)(DEV_STRUCT(base), vq);
	else
		pr_err("%s: qnotify queue %lu: missing vq/vops notify\r\n",
			name, idx);

	if (base->mtx)
		pthread_mutex_unlock(base->mtx);
}


vq_endchains(struct virtio_vq_info *vq, int used_all_avail)
{
	struct virtio_base *base;
	uint16_t event_idx, new_idx, old_idx;
	int intr;

	if (!vq || !vq->used)
		return;

	/*
	 * Interrupt generation: if we're using EVENT_IDX,
	 * interrupt if we've crossed the event threshold.
	 * Otherwise interrupt is generated if we added "used" entries,
	 * but suppressed by VRING_AVAIL_F_NO_INTERRUPT.
	 *
	 * In any case, though, if NOTIFY_ON_EMPTY is set and the
	 * entire avail was processed, we need to interrupt always.
	 */

	atomic_thread_fence();

	base = vq->base;
	old_idx = vq->save_used;
	vq->save_used = new_idx = vq->used->idx;
	if (used_all_avail &&
	    (base->negotiated_caps & (1 << VIRTIO_F_NOTIFY_ON_EMPTY)))
		intr = 1;
	else if (base->negotiated_caps & (1 << VIRTIO_RING_F_EVENT_IDX)) {
		event_idx = VQ_USED_EVENT_IDX(vq);
		/*
		 * This calculation is per docs and the kernel
		 * (see src/sys/dev/virtio/virtio_ring.h).
		 */
		intr = (uint16_t)(new_idx - event_idx - 1) <
			(uint16_t)(new_idx - old_idx);
	} else {
		intr = new_idx != old_idx &&
		    !(vq->avail->flags & VRING_AVAIL_F_NO_INTERRUPT);
	}
	if (intr)
		vq_interrupt(base, vq);
}


virtio_reset_dev(struct virtio_base *base)
{
	struct virtio_vq_info *vq;
	int i, nvq;

/* if (base->mtx) */
/* assert(pthread_mutex_isowned_np(base->mtx)); */

	acrn_timer_deinit(&base->polling_timer);
	base->polling_in_progress = 0;

	nvq = base->vops->nvq;
	for (vq = base->queues, i = 0; i < nvq; vq++, i++) {
		vq->flags = 0;
		vq->last_avail = 0;
		vq->save_used = 0;
		vq->pfn = 0;
		vq->msix_idx = VIRTIO_MSI_NO_VECTOR;
		vq->gpa_desc[0] = 0;
		vq->gpa_desc[1] = 0;
		vq->gpa_avail[0] = 0;
		vq->gpa_avail[1] = 0;
		vq->gpa_used[0] = 0;
		vq->gpa_used[1] = 0;
		vq->enabled = 0;
	}
	base->negotiated_caps = 0;
	base->curq = 0;
	/* base->status = 0; -- redundant */
	if (base->isr)
		pci_lintr_deassert(base->dev);
	base->isr = 0;
	base->msix_cfg_idx = VIRTIO_MSI_NO_VECTOR;
	base->device_feature_select = 0;
	base->driver_feature_select = 0;
	base->config_generation = 0;
}


virtio_find_legacy_cr(int offset) {
	return virtio_find_cr(legacy_config_regs,
		sizeof(legacy_config_regs) / sizeof(*legacy_config_regs),
		offset);
}

struct config_reg {
	uint16_t	offset;	/* register offset */
	uint8_t		size;	/* size (bytes) */
	uint8_t		ro;	/* true => reg is read only */
	const char	*name;	/* name of reg */
}

static struct config_reg legacy_config_regs[] = {
	{ VIRTIO_PCI_HOST_FEATURES,	4, 1, "HOSTCAP" },
	{ VIRTIO_PCI_GUEST_FEATURES,	4, 0, "GUESTCAP" },
	{ VIRTIO_PCI_QUEUE_PFN,		4, 0, "PFN" },
	{ VIRTIO_PCI_QUEUE_NUM,		2, 1, "QNUM" },
	{ VIRTIO_PCI_QUEUE_SEL,		2, 0, "QSEL" },
	{ VIRTIO_PCI_QUEUE_NOTIFY,	2, 0, "QNOTIFY" },
	{ VIRTIO_PCI_STATUS,		1, 0, "STATUS" },
	{ VIRTIO_PCI_ISR,		1, 0, "ISR" },
	{ VIRTIO_MSI_CONFIG_VECTOR,	2, 0, "CFGVEC" },
	{ VIRTIO_MSI_QUEUE_VECTOR,	2, 0, "QVEC" },
};

static struct cap_region {
	uint64_t	cap_offset;	/* offset of capability region */
	int		cap_size;	/* size of capability region */
	int		cap_id;		/* capability id */
} cap_regions[] = {
	{VIRTIO_CAP_COMMON_OFFSET, VIRTIO_CAP_COMMON_SIZE,
		VIRTIO_PCI_CAP_COMMON_CFG},
	{VIRTIO_CAP_ISR_OFFSET, VIRTIO_CAP_ISR_SIZE,
		VIRTIO_PCI_CAP_ISR_CFG},
	{VIRTIO_CAP_DEVICE_OFFSET, VIRTIO_CAP_DEVICE_SIZE,
		VIRTIO_PCI_CAP_DEVICE_CFG},
	{VIRTIO_CAP_NOTIFY_OFFSET, VIRTIO_CAP_NOTIFY_SIZE,
		VIRTIO_PCI_CAP_NOTIFY_CFG},
};


virtio_notify_cfg_write(struct pci_vdev *dev, uint64_t offset, int size,
			uint64_t value)
{
	struct virtio_base *base = dev->arg;
	struct virtio_vq_info *vq;
	struct virtio_ops *vops;
	const char *name;
	uint64_t idx;

	idx = offset / VIRTIO_MODERN_NOTIFY_OFF_MULT;
	vops = base->vops;
	name = vops->name;

	if (idx >= vops->nvq) {
		pr_err("%s: queue %lu notify out of range\r\n", name, idx);
		return;
	}

	vq = &base->queues[idx];
	if (vq->notify)
		(*vq->notify)(DEV_STRUCT(base), vq);
	else if (vops->qnotify)
		(*vops->qnotify)(DEV_STRUCT(base), vq);
	else
		pr_err("%s: qnotify queue %lu: missing vq/vops notify\r\n",
			name, idx);
}


virtio_device_cfg_write(struct pci_vdev *dev, uint64_t offset, int size,
			uint64_t value)
{
	struct virtio_base *base = dev->arg;
	struct virtio_ops *vops;
	const char *name;
	uint64_t max;
	int error = -1;

	vops = base->vops;
	name = vops->name;
	max = vops->cfgsize ? vops->cfgsize : 0x100000000;

	if (offset + size > max) {
		pr_err("%s: writing to 0x%lx size %d exceeds limit\r\n",
			name, offset, size);
		return;
	}

	if (vops->cfgwrite) {
		error = (*vops->cfgwrite)(DEV_STRUCT(base), offset, size, value);
	}
	if (error)
		pr_err("%s: writing ot 0x%lx size %d failed %d\r\n",
			name, offset, size, error);
}


void vq_clear_used_ring_flags(struct virtio_base *base, struct virtio_vq_info *vq)
{
	int backend_type = base->backend_type;
	int polling_in_progress = base->polling_in_progress;

	/* we should never unmask notification in polling mode */
	if (virtio_poll_enabled && backend_type == BACKEND_VBSU && polling_in_progress == 1)
		return;

	vq->used->flags &= ~VRING_USED_F_NO_NOTIFY;
}

#define DEV_STRUCT(vs) ((void *)(vs))

static uint8_t virtio_poll_enabled;

struct config_reg {
	uint16_t	offset;	/* register offset */
	uint8_t		size;	/* size (bytes) */
	uint8_t		ro;	/* true => reg is read only */
	const char	*name;	/* name of reg */
}

static struct config_reg legacy_config_regs[] = {
	{ VIRTIO_PCI_HOST_FEATURES,	4, 1, "HOSTCAP" },
	{ VIRTIO_PCI_GUEST_FEATURES,	4, 0, "GUESTCAP" },
	{ VIRTIO_PCI_QUEUE_PFN,		4, 0, "PFN" },
	{ VIRTIO_PCI_QUEUE_NUM,		2, 1, "QNUM" },
	{ VIRTIO_PCI_QUEUE_SEL,		2, 0, "QSEL" },
	{ VIRTIO_PCI_QUEUE_NOTIFY,	2, 0, "QNOTIFY" },
	{ VIRTIO_PCI_STATUS,		1, 0, "STATUS" },
	{ VIRTIO_PCI_ISR,		1, 0, "ISR" },
	{ VIRTIO_MSI_CONFIG_VECTOR,	2, 0, "CFGVEC" },
	{ VIRTIO_MSI_QUEUE_VECTOR,	2, 0, "QVEC" },
};

static struct config_reg modern_config_regs[] = {
	{ VIRTIO_PCI_COMMON_DFSELECT,		4, 0, "DFSELECT" },
	{ VIRTIO_PCI_COMMON_DF,			4, 1, "DF" },
	{ VIRTIO_PCI_COMMON_GFSELECT,		4, 0, "GFSELECT" },
	{ VIRTIO_PCI_COMMON_GF,			4, 0, "GF" },
	{ VIRTIO_PCI_COMMON_MSIX,		2, 0, "MSIX" },
	{ VIRTIO_PCI_COMMON_NUMQ,		2, 1, "NUMQ" },
	{ VIRTIO_PCI_COMMON_STATUS,		1, 0, "STATUS" },
	{ VIRTIO_PCI_COMMON_CFGGENERATION,	1, 1, "CFGGENERATION" },
	{ VIRTIO_PCI_COMMON_Q_SELECT,		2, 0, "Q_SELECT" },
	{ VIRTIO_PCI_COMMON_Q_SIZE,		2, 0, "Q_SIZE" },
	{ VIRTIO_PCI_COMMON_Q_MSIX,		2, 0, "Q_MSIX" },
	{ VIRTIO_PCI_COMMON_Q_ENABLE,		2, 0, "Q_ENABLE" },
	{ VIRTIO_PCI_COMMON_Q_NOFF,		2, 1, "Q_NOFF" },
	{ VIRTIO_PCI_COMMON_Q_DESCLO,		4, 0, "Q_DESCLO" },
	{ VIRTIO_PCI_COMMON_Q_DESCHI,		4, 0, "Q_DESCHI" },
	{ VIRTIO_PCI_COMMON_Q_AVAILLO,		4, 0, "Q_AVAILLO" },
	{ VIRTIO_PCI_COMMON_Q_AVAILHI,		4, 0, "Q_AVAILHI" },
	{ VIRTIO_PCI_COMMON_Q_USEDLO,		4, 0, "Q_USEDLO" },
	{ VIRTIO_PCI_COMMON_Q_USEDHI,		4, 0, "Q_USEDHI" },
};


virtio_isr_cfg_read(struct pci_vdev *dev, uint64_t offset, int size)
{
	struct virtio_base *base = dev->arg;
	uint32_t value = 0;

	value = base->isr;
	base->isr = 0;		/* a read clears this flag */
	if (value)
		pci_lintr_deassert(dev);

	return value;
}


virtio_set_modern_mmio_bar(struct virtio_base *base, int barnum)
{
	struct virtio_ops *vops;
	int rc;
	struct virtio_pci_cap cap = {
		.cap_vndr = PCIY_VENDOR,
		.cap_next = 0,
		.cap_len = sizeof(cap),
		.bar = barnum,
	};
	struct virtio_pci_notify_cap notify = {
		.cap.cap_vndr = PCIY_VENDOR,
		.cap.cap_next = 0,
		.cap.cap_len = sizeof(notify),
		.cap.cfg_type = VIRTIO_PCI_CAP_NOTIFY_CFG,
		.cap.bar = barnum,
		.cap.offset = VIRTIO_CAP_NOTIFY_OFFSET,
		.cap.length = VIRTIO_CAP_NOTIFY_SIZE,
		.notify_off_multiplier = VIRTIO_MODERN_NOTIFY_OFF_MULT,
	};
	struct virtio_pci_cfg_cap cfg = {
		.cap.cap_vndr = PCIY_VENDOR,
		.cap.cap_next = 0,
		.cap.cap_len = sizeof(cfg),
		.cap.cfg_type = VIRTIO_PCI_CAP_PCI_CFG,
	};

	vops = base->vops;

	if (vops->cfgsize > VIRTIO_CAP_DEVICE_SIZE) {
		pr_err("%s: cfgsize %lu > max %d\r\n",
			vops->name, vops->cfgsize, VIRTIO_CAP_DEVICE_SIZE);
		return -1;
	}

	/* common configuration capability */
	cap.cfg_type = VIRTIO_PCI_CAP_COMMON_CFG;
	cap.offset = VIRTIO_CAP_COMMON_OFFSET;
	cap.length = VIRTIO_CAP_COMMON_SIZE;
	rc = pci_emul_add_capability(base->dev, (u_char *)&cap, sizeof(cap));
	if (rc != 0) {
		pr_err("pci emulation add common configuration capability failed\n");
		return -1;
	}

	/* isr status capability */
	cap.cfg_type = VIRTIO_PCI_CAP_ISR_CFG;
	cap.offset = VIRTIO_CAP_ISR_OFFSET;
	cap.length = VIRTIO_CAP_ISR_SIZE;
	rc = pci_emul_add_capability(base->dev, (u_char *)&cap, sizeof(cap));
	if (rc != 0) {
		pr_err("pci emulation add isr status capability failed\n");
		return -1;
	}

	/* device specific configuration capability */
	cap.cfg_type = VIRTIO_PCI_CAP_DEVICE_CFG;
	cap.offset = VIRTIO_CAP_DEVICE_OFFSET;
	cap.length = VIRTIO_CAP_DEVICE_SIZE;
	rc = pci_emul_add_capability(base->dev, (u_char *)&cap, sizeof(cap));
	if (rc != 0) {
		pr_err("pci emulation add device specific configuration capability failed\n");
		return -1;
	}

	/* notification capability */
	rc = pci_emul_add_capability(base->dev, (u_char *)&notify,
		sizeof(notify));
	if (rc != 0) {
		pr_err("pci emulation add notification capability failed\n");
		return -1;
	}

	/* pci alternative configuration access capability */
	rc = pci_emul_add_capability(base->dev, (u_char *)&cfg, sizeof(cfg));
	if (rc != 0) {
		pr_err("pci emulation add alternative configuration access capability failed\n");
		return -1;
	}

	/* allocate and register modern memory bar */
	rc = pci_emul_alloc_bar(base->dev, barnum, PCIBAR_MEM64,
				VIRTIO_MODERN_MEM_BAR_SIZE);
	if (rc != 0) {
		pr_err("allocate and register modern memory bar failed\n");
		return -1;
	}

	base->cfg_coff = virtio_find_capability(base, VIRTIO_PCI_CAP_PCI_CFG);
	if (base->cfg_coff < 0) {
		pr_err("%s: VIRTIO_PCI_CAP_PCI_CFG not found\r\n",
			vops->name);
		return -1;
	}

	base->modern_mmio_bar_idx = barnum;
	return 0;
}


virtio_pci_legacy_write(struct vmctx *ctx, int vcpu, struct pci_vdev *dev,
			int baridx, uint64_t offset, int size, uint64_t value)
{
	struct virtio_base *base = dev->arg;
	struct virtio_vq_info *vq;
	struct virtio_ops *vops;
	const struct config_reg *cr;
	uint64_t virtio_config_size, max;
	const char *name;
	uint32_t newoff;
	int error = -1;

	if (base->mtx)
		pthread_mutex_lock(base->mtx);

	vops = base->vops;
	name = vops->name;

	if (size != 1 && size != 2 && size != 4)
		goto bad;

	if (pci_msix_enabled(dev))
		virtio_config_size = VIRTIO_PCI_CONFIG_OFF(1);
	else
		virtio_config_size = VIRTIO_PCI_CONFIG_OFF(0);

	if (offset >= virtio_config_size) {
		/*
		 * Subtract off the standard size (including MSI-X
		 * registers if enabled) and dispatch to underlying driver.
		 */
		newoff = offset - virtio_config_size;
		max = vops->cfgsize ? vops->cfgsize : 0x100000000;
		if (newoff + size > max)
			goto bad;
		if (vops->cfgwrite) {
			error = (*vops->cfgwrite)(DEV_STRUCT(base), newoff,
						  size, value);
		}
		if (!error)
			goto done;
	}

bad:
	cr = virtio_find_legacy_cr(offset);
	if (cr == NULL || cr->size != size || cr->ro) {
		if (cr != NULL) {
			/* offset must be OK, wrong size and/or reg is R/O */
			if (cr->size != size)
				pr_err("%s: write to %s: bad size %d\r\n",
				    name, cr->name, size);
			if (cr->ro)
				pr_err("%s: write to read-only reg %s\r\n",
				    name, cr->name);
		} else {
			pr_err("%s: write to bad offset/size %jd/%d\r\n",
			    name, (uintmax_t)offset, size);
		}
		goto done;
	}

	switch (offset) {
	case VIRTIO_PCI_GUEST_FEATURES:
		base->negotiated_caps = value & base->device_caps;
		if (vops->apply_features)
			(*vops->apply_features)(DEV_STRUCT(base),
			    base->negotiated_caps);
		break;
	case VIRTIO_PCI_QUEUE_PFN:
		if (base->curq >= vops->nvq)
			goto bad_qindex;
		virtio_vq_init(base, value);
		break;
	case VIRTIO_PCI_QUEUE_SEL:
		/*
		 * Note that the guest is allowed to select an
		 * invalid queue; we just need to return a QNUM
		 * of 0 while the bad queue is selected.
		 */
		base->curq = value;
		break;
	case VIRTIO_PCI_QUEUE_NOTIFY:
		if (value >= vops->nvq) {
			pr_err("%s: queue %d notify out of range\r\n",
				name, (int)value);
			goto done;
		}
		vq = &base->queues[value];
		if (vq->notify)
			(*vq->notify)(DEV_STRUCT(base), vq);
		else if (vops->qnotify)
			(*vops->qnotify)(DEV_STRUCT(base), vq);
		else
			pr_err("%s: qnotify queue %d: missing vq/vops notify\r\n",
				name, (int)value);
		break;
	case VIRTIO_PCI_STATUS:
		base->status = value;
		if (vops->set_status)
			(*vops->set_status)(DEV_STRUCT(base), value);
		if ((value == 0) && (vops->reset))
			(*vops->reset)(DEV_STRUCT(base));
		if ((value & VIRTIO_CONFIG_S_DRIVER_OK) &&
		     base->backend_type == BACKEND_VBSU &&
		     virtio_poll_enabled) {
			base->polling_timer.clockid = CLOCK_MONOTONIC;
			acrn_timer_init(&base->polling_timer, virtio_poll_timer, base);
			/* wait 5s to start virtio poll mode
			 * skip vsbl and make sure device initialization completed
			 * FIXME: Need optimization in the future
			 */
			virtio_start_timer(&base->polling_timer, 5, 0);
		}
		break;
	case VIRTIO_MSI_CONFIG_VECTOR:
		base->msix_cfg_idx = value;
		break;
	case VIRTIO_MSI_QUEUE_VECTOR:
		if (base->curq >= vops->nvq)
			goto bad_qindex;
		vq = &base->queues[base->curq];
		vq->msix_idx = value;
		break;
	}
	goto done;

bad_qindex:
	pr_err("%s: write config reg %s: curq %d >= max %d\r\n",
	    name, cr->name, base->curq, vops->nvq);
done:
	if (base->mtx)
		pthread_mutex_unlock(base->mtx);
}


virtio_vq_enable(struct virtio_base *base)
{
	struct virtio_vq_info *vq;
	uint16_t qsz;
	uint64_t phys;
	size_t size;
	char *vb;

	vq = &base->queues[base->curq];
	qsz = vq->qsize;

	/* descriptors */
	phys = (((uint64_t)vq->gpa_desc[1]) << 32) | vq->gpa_desc[0];
	size = qsz * sizeof(struct vring_desc);
	vb = paddr_guest2host(base->dev->vmctx, phys, size);
	if (!vb)
		goto error;
	vq->desc = (struct vring_desc *)vb;

	/* available ring */
	phys = (((uint64_t)vq->gpa_avail[1]) << 32) | vq->gpa_avail[0];
	size = (2 + qsz + 1) * sizeof(uint16_t);
	vb = paddr_guest2host(base->dev->vmctx, phys, size);
	if (!vb)
		goto error;

	vq->avail = (struct vring_avail *)vb;

	/* used ring */
	phys = (((uint64_t)vq->gpa_used[1]) << 32) | vq->gpa_used[0];
	size = sizeof(uint16_t) * 3 + sizeof(struct vring_used_elem) * qsz;
	vb = paddr_guest2host(base->dev->vmctx, phys, size);
	if (!vb)
		goto error;
	vq->used = (struct vring_used *)vb;

	/* Start at 0 when we use it. */
	vq->last_avail = 0;
	vq->save_used = 0;

	/* Mark queue as enabled. */
	vq->enabled = true;

	/* Mark queue as allocated after initialization is complete. */
	mb();
	vq->flags = VQ_ALLOC;
	return;
 error:
	vq->flags = 0;
	pr_err("%s: vq enable failed\n", __func__);
}

static size_t virtio_poll_interval;

struct config_reg {
	uint16_t	offset;	/* register offset */
	uint8_t		size;	/* size (bytes) */
	uint8_t		ro;	/* true => reg is read only */
	const char	*name;	/* name of reg */
}

static struct config_reg legacy_config_regs[] = {
	{ VIRTIO_PCI_HOST_FEATURES,	4, 1, "HOSTCAP" },
	{ VIRTIO_PCI_GUEST_FEATURES,	4, 0, "GUESTCAP" },
	{ VIRTIO_PCI_QUEUE_PFN,		4, 0, "PFN" },
	{ VIRTIO_PCI_QUEUE_NUM,		2, 1, "QNUM" },
	{ VIRTIO_PCI_QUEUE_SEL,		2, 0, "QSEL" },
	{ VIRTIO_PCI_QUEUE_NOTIFY,	2, 0, "QNOTIFY" },
	{ VIRTIO_PCI_STATUS,		1, 0, "STATUS" },
	{ VIRTIO_PCI_ISR,		1, 0, "ISR" },
	{ VIRTIO_MSI_CONFIG_VECTOR,	2, 0, "CFGVEC" },
	{ VIRTIO_MSI_QUEUE_VECTOR,	2, 0, "QVEC" },
};

static struct config_reg modern_config_regs[] = {
	{ VIRTIO_PCI_COMMON_DFSELECT,		4, 0, "DFSELECT" },
	{ VIRTIO_PCI_COMMON_DF,			4, 1, "DF" },
	{ VIRTIO_PCI_COMMON_GFSELECT,		4, 0, "GFSELECT" },
	{ VIRTIO_PCI_COMMON_GF,			4, 0, "GF" },
	{ VIRTIO_PCI_COMMON_MSIX,		2, 0, "MSIX" },
	{ VIRTIO_PCI_COMMON_NUMQ,		2, 1, "NUMQ" },
	{ VIRTIO_PCI_COMMON_STATUS,		1, 0, "STATUS" },
	{ VIRTIO_PCI_COMMON_CFGGENERATION,	1, 1, "CFGGENERATION" },
	{ VIRTIO_PCI_COMMON_Q_SELECT,		2, 0, "Q_SELECT" },
	{ VIRTIO_PCI_COMMON_Q_SIZE,		2, 0, "Q_SIZE" },
	{ VIRTIO_PCI_COMMON_Q_MSIX,		2, 0, "Q_MSIX" },
	{ VIRTIO_PCI_COMMON_Q_ENABLE,		2, 0, "Q_ENABLE" },
	{ VIRTIO_PCI_COMMON_Q_NOFF,		2, 1, "Q_NOFF" },
	{ VIRTIO_PCI_COMMON_Q_DESCLO,		4, 0, "Q_DESCLO" },
	{ VIRTIO_PCI_COMMON_Q_DESCHI,		4, 0, "Q_DESCHI" },
	{ VIRTIO_PCI_COMMON_Q_AVAILLO,		4, 0, "Q_AVAILLO" },
	{ VIRTIO_PCI_COMMON_Q_AVAILHI,		4, 0, "Q_AVAILHI" },
	{ VIRTIO_PCI_COMMON_Q_USEDLO,		4, 0, "Q_USEDLO" },
	{ VIRTIO_PCI_COMMON_Q_USEDHI,		4, 0, "Q_USEDHI" },
};


virtio_set_modern_pio_bar(struct virtio_base *base, int barnum)
{
	int rc;
	struct virtio_pci_notify_cap notify_pio = {
		.cap.cap_vndr = PCIY_VENDOR,
		.cap.cap_next = 0,
		.cap.cap_len = sizeof(notify_pio),
		.cap.cfg_type = VIRTIO_PCI_CAP_NOTIFY_CFG,
		.cap.bar = barnum,
		.cap.offset = 0,
		.cap.length = 4,
		.notify_off_multiplier = 0,
	};

	/* notification capability */
	rc = pci_emul_add_capability(base->dev, (u_char *)&notify_pio,
		sizeof(notify_pio));
	if (rc != 0) {
		pr_err("pci emulation add notification capability for virtio modern PIO BAR failed\n");
		return -1;
	}

	/* allocate and register modern pio bar */
	rc = pci_emul_alloc_bar(base->dev, barnum, PCIBAR_IO, 4);
	if (rc != 0) {
		pr_err("allocate and register modern pio bar failed\n");
		return -1;
	}

	base->modern_pio_bar_idx = barnum;
	return 0;
}


virtio_common_cfg_read(struct pci_vdev *dev, uint64_t offset, int size)
{
	struct virtio_base *base = dev->arg;
	struct virtio_ops *vops;
	const struct config_reg *cr;
	const char *name;
	uint32_t value;

	vops = base->vops;
	name = vops->name;
	value = size == 1 ? 0xff : size == 2 ? 0xffff : 0xffffffff;

	cr = virtio_find_modern_cr(offset);
	if (cr == NULL || cr->size != size) {
		if (cr != NULL) {
			/* offset must be OK, so size must be bad */
			pr_err("%s: read from %s: bad size %d\r\n",
				name, cr->name, size);
		} else {
			pr_err("%s: read from bad offset/size %jd/%d\r\n",
				name, (uintmax_t)offset, size);
		}

		return value;
	}

	switch (offset) {
	case VIRTIO_PCI_COMMON_DFSELECT:
		value = base->device_feature_select;
		break;
	case VIRTIO_PCI_COMMON_DF:
		if (base->device_feature_select == 0)
			value = base->device_caps & 0xffffffff;
		else if (base->device_feature_select == 1)
			value = (base->device_caps >> 32) & 0xffffffff;
		else /* present 0, see 4.1.4.3.1 */
			value = 0;
		break;
	case VIRTIO_PCI_COMMON_GFSELECT:
		value = base->driver_feature_select;
		break;
	case VIRTIO_PCI_COMMON_GF:
		/* see 4.1.4.3.1. Present any valid feature bits the driver
		 * has written in driver_feature. Valid feature bits are those
		 * which are subset of the corresponding device_feature bits
		 */
		if (base->driver_feature_select == 0)
			value = base->negotiated_caps & 0xffffffff;
		else if (base->driver_feature_select == 1)
			value = (base->negotiated_caps >> 32) & 0xffffffff;
		else
			value = 0;
		break;
	case VIRTIO_PCI_COMMON_MSIX:
		value = base->msix_cfg_idx;
		break;
	case VIRTIO_PCI_COMMON_NUMQ:
		value = vops->nvq;
		break;
	case VIRTIO_PCI_COMMON_STATUS:
		value = base->status;
		break;
	case VIRTIO_PCI_COMMON_CFGGENERATION:
		value = base->config_generation;
		break;
	case VIRTIO_PCI_COMMON_Q_SELECT:
		value = base->curq;
		break;
	case VIRTIO_PCI_COMMON_Q_SIZE:
		value = base->curq < vops->nvq ?
			base->queues[base->curq].qsize : 0;
		break;
	case VIRTIO_PCI_COMMON_Q_MSIX:
		value = base->curq < vops->nvq ?
			base->queues[base->curq].msix_idx :
			VIRTIO_MSI_NO_VECTOR;
		break;
	case VIRTIO_PCI_COMMON_Q_ENABLE:
		value = base->curq < vops->nvq ?
			base->queues[base->curq].enabled : 0;
		break;
	case VIRTIO_PCI_COMMON_Q_NOFF:
		value = base->curq;
		break;
	case VIRTIO_PCI_COMMON_Q_DESCLO:
		value = base->curq < vops->nvq ?
			base->queues[base->curq].gpa_desc[0] : 0;
		break;
	case VIRTIO_PCI_COMMON_Q_DESCHI:
		value = base->curq < vops->nvq ?
			base->queues[base->curq].gpa_desc[1] : 0;
		break;
	case VIRTIO_PCI_COMMON_Q_AVAILLO:
		value = base->curq < vops->nvq ?
			base->queues[base->curq].gpa_avail[0] : 0;
		break;
	case VIRTIO_PCI_COMMON_Q_AVAILHI:
		value = base->curq < vops->nvq ?
			base->queues[base->curq].gpa_avail[1] : 0;
		break;
	case VIRTIO_PCI_COMMON_Q_USEDLO:
		value = base->curq < vops->nvq ?
			base->queues[base->curq].gpa_used[0] : 0;
		break;
	case VIRTIO_PCI_COMMON_Q_USEDHI:
		value = base->curq < vops->nvq ?
			base->queues[base->curq].gpa_used[1] : 0;
		break;
	}

	return value;
}


virtio_vq_init(struct virtio_base *base, uint32_t pfn)
{
	struct virtio_vq_info *vq;
	uint64_t phys;
	size_t size;
	char *vb;

	vq = &base->queues[base->curq];
	vq->pfn = pfn;
	phys = (uint64_t)pfn << VRING_PAGE_BITS;
	size = vring_size(vq->qsize, VIRTIO_PCI_VRING_ALIGN);
	vb = paddr_guest2host(base->dev->vmctx, phys, size);
	if (!vb)
		goto error;

	/* First page(s) are descriptors... */
	vq->desc = (struct vring_desc *)vb;
	vb += vq->qsize * sizeof(struct vring_desc);

	/* ... immediately followed by "avail" ring (entirely uint16_t's) */
	vq->avail = (struct vring_avail *)vb;
	vb += (2 + vq->qsize + 1) * sizeof(uint16_t);

	/* Then it's rounded up to the next page... */
	vb = (char *)roundup2((uintptr_t)vb, VIRTIO_PCI_VRING_ALIGN);

	/* ... and the last page(s) are the used ring. */
	vq->used = (struct vring_used *)vb;

	/* Start at 0 when we use it. */
	vq->last_avail = 0;
	vq->save_used = 0;

	/* Mark queue as allocated after initialization is complete. */
	mb();
	vq->flags = VQ_ALLOC;

	return;

error:
	vq->flags = 0;
	pr_err("%s: vq enable failed\n", __func__);
}


virtio_find_capability(struct virtio_base *base, uint8_t cfg_type)
{
	struct pci_vdev *dev = base->dev;
	uint8_t type;
	int rc, coff = 0;

	rc = pci_emul_find_capability(dev, PCIY_VENDOR, &coff);
	while (!rc) {
		type = pci_get_cfgdata8(dev,
			coff + offsetof(struct virtio_pci_cap, cfg_type));
		if (type == cfg_type)
			return coff;
		rc = pci_emul_find_capability(dev, PCIY_VENDOR, &coff);
	}

	return -1;
}

static uint8_t virtio_poll_enabled;

struct config_reg {
	uint16_t	offset;	/* register offset */
	uint8_t		size;	/* size (bytes) */
	uint8_t		ro;	/* true => reg is read only */
	const char	*name;	/* name of reg */
}

static struct config_reg legacy_config_regs[] = {
	{ VIRTIO_PCI_HOST_FEATURES,	4, 1, "HOSTCAP" },
	{ VIRTIO_PCI_GUEST_FEATURES,	4, 0, "GUESTCAP" },
	{ VIRTIO_PCI_QUEUE_PFN,		4, 0, "PFN" },
	{ VIRTIO_PCI_QUEUE_NUM,		2, 1, "QNUM" },
	{ VIRTIO_PCI_QUEUE_SEL,		2, 0, "QSEL" },
	{ VIRTIO_PCI_QUEUE_NOTIFY,	2, 0, "QNOTIFY" },
	{ VIRTIO_PCI_STATUS,		1, 0, "STATUS" },
	{ VIRTIO_PCI_ISR,		1, 0, "ISR" },
	{ VIRTIO_MSI_CONFIG_VECTOR,	2, 0, "CFGVEC" },
	{ VIRTIO_MSI_QUEUE_VECTOR,	2, 0, "QVEC" },
};

static struct config_reg modern_config_regs[] = {
	{ VIRTIO_PCI_COMMON_DFSELECT,		4, 0, "DFSELECT" },
	{ VIRTIO_PCI_COMMON_DF,			4, 1, "DF" },
	{ VIRTIO_PCI_COMMON_GFSELECT,		4, 0, "GFSELECT" },
	{ VIRTIO_PCI_COMMON_GF,			4, 0, "GF" },
	{ VIRTIO_PCI_COMMON_MSIX,		2, 0, "MSIX" },
	{ VIRTIO_PCI_COMMON_NUMQ,		2, 1, "NUMQ" },
	{ VIRTIO_PCI_COMMON_STATUS,		1, 0, "STATUS" },
	{ VIRTIO_PCI_COMMON_CFGGENERATION,	1, 1, "CFGGENERATION" },
	{ VIRTIO_PCI_COMMON_Q_SELECT,		2, 0, "Q_SELECT" },
	{ VIRTIO_PCI_COMMON_Q_SIZE,		2, 0, "Q_SIZE" },
	{ VIRTIO_PCI_COMMON_Q_MSIX,		2, 0, "Q_MSIX" },
	{ VIRTIO_PCI_COMMON_Q_ENABLE,		2, 0, "Q_ENABLE" },
	{ VIRTIO_PCI_COMMON_Q_NOFF,		2, 1, "Q_NOFF" },
	{ VIRTIO_PCI_COMMON_Q_DESCLO,		4, 0, "Q_DESCLO" },
	{ VIRTIO_PCI_COMMON_Q_DESCHI,		4, 0, "Q_DESCHI" },
	{ VIRTIO_PCI_COMMON_Q_AVAILLO,		4, 0, "Q_AVAILLO" },
	{ VIRTIO_PCI_COMMON_Q_AVAILHI,		4, 0, "Q_AVAILHI" },
	{ VIRTIO_PCI_COMMON_Q_USEDLO,		4, 0, "Q_USEDLO" },
	{ VIRTIO_PCI_COMMON_Q_USEDHI,		4, 0, "Q_USEDHI" },
};


virtio_linkup(struct virtio_base *base, struct virtio_ops *vops,
	      void *pci_virtio_dev, struct pci_vdev *dev,
	      struct virtio_vq_info *queues,
	      int backend_type)
{
	int i;

	/* base and pci_virtio_dev addresses must match */
	if ((void *)base != pci_virtio_dev) {
		pr_err("virtio_base and pci_virtio_dev addresses don't match!\n");
		return;
	}
	base->vops = vops;
	base->dev = dev;
	dev->arg = base;
	base->backend_type = backend_type;

	base->queues = queues;
	for (i = 0; i < vops->nvq; i++) {
		queues[i].base = base;
		queues[i].num = i;
	}
}


vq_retchain(struct virtio_vq_info *vq)
{
	vq->last_avail--;
}


virtio_set_io_bar(struct virtio_base *base, int barnum)
{
	size_t size;

	/*
	 * ??? should we use VIRTIO_PCI_CONFIG_OFF(0) if MSI-X
	 * is disabled? Existing code did not...
	 */
	size = VIRTIO_PCI_CONFIG_OFF(1) + base->vops->cfgsize;
	pci_emul_alloc_bar(base->dev, barnum, PCIBAR_IO, size);
	base->legacy_pio_bar_idx = barnum;
}


virtio_start_timer(struct acrn_timer *timer, time_t sec, time_t nsec)
{
	struct itimerspec ts;

	/* setting the interval time */
	ts.it_interval.tv_sec = 0;
	ts.it_interval.tv_nsec = 0;
	/* set the delay time it will be started when timer_setting */
	ts.it_value.tv_sec = sec;
	ts.it_value.tv_nsec = nsec;
	if (acrn_timer_settime(timer, &ts) != 0) {
		pr_err("acrn timer set time failed\n");
		return;
	}
}


virtio_pci_modern_mmio_read(struct vmctx *ctx, int vcpu, struct pci_vdev *dev,
			    int baridx, uint64_t offset, int size)
{
	struct virtio_base *base = dev->arg;
	struct virtio_ops *vops;
	const char *name;
	uint32_t value;
	int capid;

	vops = base->vops;
	name = vops->name;
	value = size == 1 ? 0xff : size == 2 ? 0xffff : 0xffffffff;

	if (size != 1 && size != 2 && size != 4) {
		pr_err("%s: read from [%d:0x%lx] bad size %d\r\n",
			name, baridx, offset, size);
		return value;
	}

	capid = virtio_get_cap_id(offset, size);
	if (capid < 0) {
		pr_err("%s: read from [%d:0x%lx] bad range %d\r\n",
			name, baridx, offset, size);
		return value;
	}

	if (base->mtx)
		pthread_mutex_lock(base->mtx);

	switch (capid) {
	case VIRTIO_PCI_CAP_COMMON_CFG:
		offset -= VIRTIO_CAP_COMMON_OFFSET;
		value = virtio_common_cfg_read(dev, offset, size);
		break;
	case VIRTIO_PCI_CAP_ISR_CFG:
		offset -= VIRTIO_CAP_ISR_OFFSET;
		value = virtio_isr_cfg_read(dev, offset, size);
		break;
	case VIRTIO_PCI_CAP_DEVICE_CFG:
		offset -= VIRTIO_CAP_DEVICE_OFFSET;
		value = virtio_device_cfg_read(dev, offset, size);
		break;
	default: /* guest driver should not read from notify region */
		pr_err("%s: read from [%d:0x%lx] size %d not supported\r\n",
			name, baridx, offset, size);
	}

	if (base->mtx)
		pthread_mutex_unlock(base->mtx);
	return value;
}


virtio_interrupt_init(struct virtio_base *base, int use_msix)
{
	return virtio_intr_init(base, 1, use_msix);
}


virtio_common_cfg_write(struct pci_vdev *dev, uint64_t offset, int size,
			uint64_t value)
{
	struct virtio_base *base = dev->arg;
	struct virtio_vq_info *vq;
	struct virtio_ops *vops;
	const struct config_reg *cr;
	const char *name;

	vops = base->vops;
	name = vops->name;

	cr = virtio_find_modern_cr(offset);
	if (cr == NULL || cr->size != size || cr->ro) {
		if (cr != NULL) {
			/* offset must be OK, wrong size and/or reg is R/O */
			if (cr->size != size)
				pr_err("%s: write to %s: bad size %d\r\n",
					name, cr->name, size);
			if (cr->ro)
				pr_err("%s: write to read-only reg %s\r\n",
					name, cr->name);
		} else {
			pr_err("%s: write to bad offset/size %jd/%d\r\n",
				name, (uintmax_t)offset, size);
		}

		return;
	}

	switch (offset) {
	case VIRTIO_PCI_COMMON_DFSELECT:
		base->device_feature_select = value;
		break;
	case VIRTIO_PCI_COMMON_GFSELECT:
		base->driver_feature_select = value;
		break;
	case VIRTIO_PCI_COMMON_GF:
		if (base->status & VIRTIO_CONFIG_S_DRIVER_OK)
			break;
		if (base->driver_feature_select < 2) {
			value &= 0xffffffff;
			base->negotiated_caps =
				(value << (base->driver_feature_select * 32))
				& base->device_caps;
			if (vops->apply_features)
				(*vops->apply_features)(DEV_STRUCT(base),
					base->negotiated_caps);
		}
		break;
	case VIRTIO_PCI_COMMON_MSIX:
		base->msix_cfg_idx = value;
		break;
	case VIRTIO_PCI_COMMON_STATUS:
		base->status = value & 0xff;
		if (vops->set_status)
			(*vops->set_status)(DEV_STRUCT(base), value);
		if ((base->status == 0) && (vops->reset))
			(*vops->reset)(DEV_STRUCT(base));
		/* TODO: virtio poll mode for modern devices */
		break;
	case VIRTIO_PCI_COMMON_Q_SELECT:
		/*
		 * Note that the guest is allowed to select an
		 * invalid queue; we just need to return a QNUM
		 * of 0 while the bad queue is selected.
		 */
		base->curq = value;
		break;
	case VIRTIO_PCI_COMMON_Q_SIZE:
		if (base->curq >= vops->nvq)
			goto bad_qindex;
		vq = &base->queues[base->curq];
		vq->qsize = value;
		break;
	case VIRTIO_PCI_COMMON_Q_MSIX:
		if (base->curq >= vops->nvq)
			goto bad_qindex;
		vq = &base->queues[base->curq];
		vq->msix_idx = value;
		break;
	case VIRTIO_PCI_COMMON_Q_ENABLE:
		if (base->curq >= vops->nvq)
			goto bad_qindex;
		virtio_vq_enable(base);
		break;
	case VIRTIO_PCI_COMMON_Q_DESCLO:
		if (base->curq >= vops->nvq)
			goto bad_qindex;
		vq = &base->queues[base->curq];
		vq->gpa_desc[0] = value;
		break;
	case VIRTIO_PCI_COMMON_Q_DESCHI:
		if (base->curq >= vops->nvq)
			goto bad_qindex;
		vq = &base->queues[base->curq];
		vq->gpa_desc[1] = value;
		break;
	case VIRTIO_PCI_COMMON_Q_AVAILLO:
		if (base->curq >= vops->nvq)
			goto bad_qindex;
		vq = &base->queues[base->curq];
		vq->gpa_avail[0] = value;
		break;
	case VIRTIO_PCI_COMMON_Q_AVAILHI:
		if (base->curq >= vops->nvq)
			goto bad_qindex;
		vq = &base->queues[base->curq];
		vq->gpa_avail[1] = value;
		break;
	case VIRTIO_PCI_COMMON_Q_USEDLO:
		if (base->curq >= vops->nvq)
			goto bad_qindex;
		vq = &base->queues[base->curq];
		vq->gpa_used[0] = value;
		break;
	case VIRTIO_PCI_COMMON_Q_USEDHI:
		if (base->curq >= vops->nvq)
			goto bad_qindex;
		vq = &base->queues[base->curq];
		vq->gpa_used[1] = value;
		break;
	}

	return;

bad_qindex:
	pr_err("%s: write config reg %s: curq %d >= max %d\r\n",
		name, cr->name, base->curq, vops->nvq);
}


vq_relchain(struct virtio_vq_info *vq, uint16_t idx, uint32_t iolen)
{
	uint16_t uidx, mask;
	volatile struct vring_used *vuh;
	volatile struct vring_used_elem *vue;

	/*
	 * Notes:
	 *  - mask is N-1 where N is a power of 2 so computes x % N
	 *  - vuh points to the "used" data shared with guest
	 *  - vue points to the "used" ring entry we want to update
	 *  - head is the same value we compute in vq_iovecs().
	 *
	 * (I apologize for the two fields named idx; the
	 * virtio spec calls the one that vue points to, "id"...)
	 */
	mask = vq->qsize - 1;
	vuh = vq->used;

	uidx = vuh->idx;
	vue = &vuh->ring[uidx++ & mask];
	vue->id = idx;
	vue->len = iolen;
	vuh->idx = uidx;
}


virtio_find_modern_cr(int offset) {
	return virtio_find_cr(modern_config_regs,
		sizeof(modern_config_regs) / sizeof(*modern_config_regs),
		offset);
}

struct config_reg {
	uint16_t	offset;	/* register offset */
	uint8_t		size;	/* size (bytes) */
	uint8_t		ro;	/* true => reg is read only */
	const char	*name;	/* name of reg */
}

static struct config_reg modern_config_regs[] = {
	{ VIRTIO_PCI_COMMON_DFSELECT,		4, 0, "DFSELECT" },
	{ VIRTIO_PCI_COMMON_DF,			4, 1, "DF" },
	{ VIRTIO_PCI_COMMON_GFSELECT,		4, 0, "GFSELECT" },
	{ VIRTIO_PCI_COMMON_GF,			4, 0, "GF" },
	{ VIRTIO_PCI_COMMON_MSIX,		2, 0, "MSIX" },
	{ VIRTIO_PCI_COMMON_NUMQ,		2, 1, "NUMQ" },
	{ VIRTIO_PCI_COMMON_STATUS,		1, 0, "STATUS" },
	{ VIRTIO_PCI_COMMON_CFGGENERATION,	1, 1, "CFGGENERATION" },
	{ VIRTIO_PCI_COMMON_Q_SELECT,		2, 0, "Q_SELECT" },
	{ VIRTIO_PCI_COMMON_Q_SIZE,		2, 0, "Q_SIZE" },
	{ VIRTIO_PCI_COMMON_Q_MSIX,		2, 0, "Q_MSIX" },
	{ VIRTIO_PCI_COMMON_Q_ENABLE,		2, 0, "Q_ENABLE" },
	{ VIRTIO_PCI_COMMON_Q_NOFF,		2, 1, "Q_NOFF" },
	{ VIRTIO_PCI_COMMON_Q_DESCLO,		4, 0, "Q_DESCLO" },
	{ VIRTIO_PCI_COMMON_Q_DESCHI,		4, 0, "Q_DESCHI" },
	{ VIRTIO_PCI_COMMON_Q_AVAILLO,		4, 0, "Q_AVAILLO" },
	{ VIRTIO_PCI_COMMON_Q_AVAILHI,		4, 0, "Q_AVAILHI" },
	{ VIRTIO_PCI_COMMON_Q_USEDLO,		4, 0, "Q_USEDLO" },
	{ VIRTIO_PCI_COMMON_Q_USEDHI,		4, 0, "Q_USEDHI" },
};

static struct cap_region {
	uint64_t	cap_offset;	/* offset of capability region */
	int		cap_size;	/* size of capability region */
	int		cap_id;		/* capability id */
} cap_regions[] = {
	{VIRTIO_CAP_COMMON_OFFSET, VIRTIO_CAP_COMMON_SIZE,
		VIRTIO_PCI_CAP_COMMON_CFG},
	{VIRTIO_CAP_ISR_OFFSET, VIRTIO_CAP_ISR_SIZE,
		VIRTIO_PCI_CAP_ISR_CFG},
	{VIRTIO_CAP_DEVICE_OFFSET, VIRTIO_CAP_DEVICE_SIZE,
		VIRTIO_PCI_CAP_DEVICE_CFG},
	{VIRTIO_CAP_NOTIFY_OFFSET, VIRTIO_CAP_NOTIFY_SIZE,
		VIRTIO_PCI_CAP_NOTIFY_CFG},
};


explicit SplitVOpCPU(OpKernelConstruction* c) : Base(c) {}

A negative size in one of the split sizes allowed the computed size of another
to exceed the total dimension, leading to a segfault and security vulnerability.
Adding a check for negative sizes prevents this.

PiperOrigin-RevId: 401035665
Change-Id: I79bbe329787dac82aa4bf60397a9129b716aedab

#define EIGEN_USE_THREADS

#define EIGEN_USE_GPU


  std::tuple<IndexType, IndexType, IndexType> SetDims(
      const TensorShape& input_shape, const int32_t split_dim) const {
    static_assert(std::is_integral<IndexType>::value,
                  "IndexType must be an integer type");
    int32_t prefix_dim_size = 1;
    for (int i = 0; i < split_dim; ++i) {
      prefix_dim_size *= input_shape.dim_size(i);
    }

    // Caller must ensure that dim_size and suffix_dim_size are <
    // std::numeric_limits<IndexType>::max()
    IndexType split_dim_size =
        static_cast<IndexType>(input_shape.dim_size(split_dim));

    IndexType suffix_dim_size = 1;
    for (int i = split_dim + 1; i < input_shape.dims(); ++i) {
      suffix_dim_size *= static_cast<IndexType>(input_shape.dim_size(i));
    }
    return std::make_tuple(prefix_dim_size, split_dim_size, suffix_dim_size);
  }


explicit SplitVOpGPU(OpKernelConstruction* c) : Base(c) {}


  void ParallelSplitByInputData(OpKernelContext* context,
                                const InputReshapedType& input_reshaped,
                                const TensorShape& input_shape,
                                const std::vector<Tlen>& split_sizes_vec,
                                const int32_t split_dim) const {
    const T* p_data = input_reshaped.data();
    const uint32 elem_pkg = input_reshaped.dimensions().rank() == 3
                                ? input_reshaped.dimension(2)
                                : 1;
    const uint32 line_elem_num =
        (input_reshaped.dimensions().rank() >= 2 ? input_reshaped.dimension(1)
                                                 : 1) *
        elem_pkg;
    const uint32 line_num = input_reshaped.dimension(0);

    // Prepare the output matrix.
    std::vector<T*> outputs(split_sizes_vec.size());
    for (uint64 i = 0; i < split_sizes_vec.size(); ++i) {
      TensorShape output_shape(input_shape);
      output_shape.set_dim(split_dim, split_sizes_vec[i]);
      Tensor* result = nullptr;
      OP_REQUIRES_OK(context,
                     context->allocate_output(i, output_shape, &result));
      outputs[i] = static_cast<T*>(&result->flat<T>()(0));
    }

    auto sub_split_func = [&split_sizes_vec, &p_data, elem_pkg, &outputs,
                           line_elem_num](int32_t start_part,
                                          int32_t end_part) {
      int start = start_part * line_elem_num;
      int end = end_part * line_elem_num;
      uint32 times = 0;
      for (int32_t i = start; i < end;) {
        for (uint32 j = 0; j < split_sizes_vec.size(); ++j) {
          const auto copy_elem_num = split_sizes_vec[j] * elem_pkg;
          std::copy_n(p_data + i, copy_elem_num,
                      &(outputs[j][(start_part + times) * copy_elem_num]));
          i += copy_elem_num;
        }
        ++times;
      }
    };

    uint32 part_size =
        context->device()->tensorflow_cpu_worker_threads()->num_threads;
    Shard(part_size,
          context->device()->tensorflow_cpu_worker_threads()->workers, line_num,
          line_num, sub_split_func);
  }


  void ComputeEasyCases(OpKernelContext* context, bool* done,
                        std::vector<Tlen>* split_sizes_vec) {
    const int32_t num_split = context->num_outputs();
    const Tensor& input = context->input(0);
    const TensorShape& input_shape = input.shape();
    const Tensor& split_tensor = context->input(1);
    const Tensor& split_dim_tensor = context->input(2);

    OP_REQUIRES(context, split_dim_tensor.NumElements() == 1,
                errors::InvalidArgument("split_dim_tensor must have "
                                        "exactly one element."));

    const int32_t split_dim_orig = split_dim_tensor.flat<int32>()(0);
    const int32_t split_dim =
        split_dim_orig < 0 ? split_dim_orig + input.dims() : split_dim_orig;

    OP_REQUIRES(
        context,
        split_tensor.dims() == 1 && split_tensor.NumElements() == num_split,
        errors::InvalidArgument("size of the split_tensor must be 1-D and have "
                                "the same elements as outputs got ",
                                split_tensor.dims(), " -D and ",
                                split_tensor.NumElements(), " elements"));

    auto split_sizes_d = split_tensor.vec<Tlen>();

    split_sizes_vec->resize(split_sizes_d.size());

    std::copy(split_sizes_d.data(), split_sizes_d.data() + split_sizes_d.size(),
              split_sizes_vec->begin());

    OP_REQUIRES(
        context, num_split > 0,
        errors::InvalidArgument(
            "Number of ways to split should be > 0, but got ", num_split));

    OP_REQUIRES(
        context, 0 <= split_dim && split_dim < input.dims(),
        errors::InvalidArgument("-input rank(-", input.dims(),
                                ") <= split_dim < input rank (", input.dims(),
                                "), but got ", split_dim_orig));

    Tlen input_size_split_dim = input_shape.dim_size(split_dim);

    // Special case 1: num_split == 1. Nothing to do.
    if (num_split == 1) {
      context->set_output(0, context->input(0));
      OP_REQUIRES(
          context, (*split_sizes_vec)[0] == input_size_split_dim,
          errors::InvalidArgument("If there is only one output, it must have "
                                  "the same size as the input. Input size: ",
                                  input_size_split_dim,
                                  " output size: ", (*split_sizes_vec)[0]));
      *done = true;
      return;
    }

    // Determine sizes of output, in case of a -1 input value
    int neg_one_dim = -1;
    Tlen determined_size = 0;
    for (int d = 0; d < split_sizes_vec->size(); ++d) {
      Tlen size = (*split_sizes_vec)[d];

      if (size == -1) {
        OP_REQUIRES(context, neg_one_dim == -1,
                    errors::InvalidArgument("There can only be one -1 in the "
                                            "input."));
        neg_one_dim = d;
      } else {
        determined_size += size;
      }
    }

    OP_REQUIRES(
        context,
        (neg_one_dim == -1 && determined_size == input_size_split_dim) ||
            (neg_one_dim >= 0 && determined_size <= input_size_split_dim),
        errors::InvalidArgument("Determined shape must either match "
                                "input shape along split_dim exactly if "
                                "fully specified, or be less than the size of "
                                "the input along split_dim if not fully "
                                "specified.  Got: ",
                                determined_size));

    if (neg_one_dim >= 0) {
      (*split_sizes_vec)[neg_one_dim] = input_size_split_dim - determined_size;
    }

    for (int i = 0; i < split_sizes_vec->size(); ++i) {
      const Tlen& split_size = (*split_sizes_vec)[i];
      OP_REQUIRES(context, split_size >= Tlen(0),
                  errors::InvalidArgument("Split size at index ", i,
                                          " must be >= 0. Got: ", split_size));
    }

    // Special case 2: split along the 1st dimension. The requirements are that
    // either we are splitting the outer dimension of two or more such that
    // every outer subpart is aligned or that the split sizes mean that they are
    // always aligned. In these cases, we can share the underlying buffer.
    //
    // Apply this optimization conservatively: if input is aligned,
    // the resulting tensors must be aligned. It's conservative
    // because if the immediate consumer of the resulting tensors are
    // not using eigen for computation, its perfectly fine to avoid
    // the copying.
    if (SplitHasAlignedOutputsInFirstDimension(
            input_shape, split_dim, absl::MakeConstSpan(*split_sizes_vec))) {
      Tlen start = 0;
      for (int i = 0; i < num_split; ++i) {
        context->set_output(i,
                            input.Slice(start, start + (*split_sizes_vec)[i]));
        start += (*split_sizes_vec)[i];
      }
      *done = true;
      return;
    }
  }


  void operator()(OpKernelContext* context,
                  const InputReshapedType& input_reshaped,
                  const std::vector<int64_t>& split_start_points,
                  const TensorShape& input_shape, int32_t split_dim,
                  Eigen::DenseIndex prefix_dim_size,
                  Eigen::DenseIndex split_dim_size,
                  Eigen::DenseIndex suffix_dim_size,
                  std::vector<Tlen>& split_sizes_vec,
                  const MakeSizesType& make_sizes,
                  const ReshapeResultType& reshape_result) const {
    Eigen::DSizes<Eigen::DenseIndex, NDims> indices;
    for (int i = 0; i < NDims; ++i) {
      indices[i] = 0;
    }
    const auto num_threads =
        context->device()->tensorflow_cpu_worker_threads()->num_threads;
    // TODO(jewillco): Tune heuristic further.
    const auto input_element_count = input_shape.num_elements();
    const int num_split = split_start_points.size();
    const bool use_parallelism_between_outputs =
        (num_split >= kMinimumSplitNum &&
         input_element_count >= std::min(num_threads, num_split) * 4096 &&
         input_element_count < num_split * 180 * 1024);

    auto range_output_func = [&indices, context, &input_shape, split_dim,
                              &split_sizes_vec, &split_start_points,
                              use_parallelism_between_outputs, &input_reshaped,
                              &make_sizes,
                              &reshape_result](int64_t start, int64_t limit) {
      for (int64_t i = start; i < limit; ++i) {
        TensorShape output_shape(input_shape);
        output_shape.set_dim(split_dim, split_sizes_vec[i]);
        Tensor* result = nullptr;
        OP_REQUIRES_OK(context,
                       context->allocate_output(i, output_shape, &result));

        const auto sizes = make_sizes(split_sizes_vec[i]);

        if (sizes.TotalSize() > 0) {
          auto result_shaped = reshape_result(result, split_sizes_vec[i]);

          auto current_indices = indices;
          current_indices[NDims - 2] = split_start_points[i];
          if (use_parallelism_between_outputs) {
            // Use sequential implementation for single output.
            result_shaped = input_reshaped.slice(current_indices, sizes);
          } else {
            // This implementation may be parallel internally.
            functor::Split<CPUDevice, T, NDims>()(
                context->eigen_device<CPUDevice>(), result_shaped,
                input_reshaped, current_indices, sizes);
          }
        }
      }
    };

    // 1. Parallel performance is not as good as serial when the amount of data
    // is too small (<kMinimumInputSize);
    // 2. There is sufficient data on the 0th dimension to ensure parallelism;
    // 3. This method only supports non-zero split.
    if ((input_element_count >= kMinimumInputSize) &&
        input_reshaped.dimension(0) > kMinimumDim0Size && split_dim) {
      // Each thread processes the same amount of data, and then copies data
      // to all output tensors .
      ParallelSplitByInputData(context, input_reshaped, input_shape,
                               split_sizes_vec, split_dim);
    } else if (use_parallelism_between_outputs) {
      // A thread maps a output tensor, this thread will traverse all the data,
      // and then put specified data to mapped output tensor. Run in parallel,
      // disabling parallelism in functor.
      Shard(num_split,
            context->device()->tensorflow_cpu_worker_threads()->workers,
            num_split, input_element_count / num_split, range_output_func);
    } else {
      // Run sequentially, but allow internal parallelism in functor.
      range_output_func(0, num_split);
    }
  }


explicit SplitVOpBase(OpKernelConstruction* c) : OpKernel(c) {}


  void Compute(OpKernelContext* context) override {
    bool done = false;
    std::vector<Tlen> split_sizes_vec;
    Base::ComputeEasyCases(context, &done, &split_sizes_vec);
    if (!context->status().ok() || done) {
      return;
    }
    const int32_t num_split = Base::num_outputs();
    const Tensor& input = context->input(0);
    const TensorShape& input_shape = input.shape();
    const int32_t split_dim_orig = context->input(2).flat<int32>()(0);
    const int32_t split_dim =
        split_dim_orig < 0 ? split_dim_orig + input.dims() : split_dim_orig;
    OP_REQUIRES(
        context,
        FastBoundsCheck(input.NumElements(), std::numeric_limits<int32>::max()),
        errors::InvalidArgument("Split on GPU requires input size "
                                "< max int32"));

    int32_t prefix_dim_size;
    int32_t split_dim_size;
    int32_t suffix_dim_size;
    std::tie(prefix_dim_size, split_dim_size, suffix_dim_size) =
        Base::template SetDims<int32>(input_shape, split_dim);

    // use the same approach as concat (see documentation there)
    // reshape to 2D

    if (num_split > 16) {
      GpuDeviceArrayOnHost<T*> ptrs(context, num_split);
      OP_REQUIRES_OK(context, ptrs.Init());

      GpuDeviceArrayOnHost<Tlen> offsets(context, num_split + 1);
      OP_REQUIRES_OK(context, offsets.Init());

      Tlen offset = 0;
      int entry = split_sizes_vec[0];
      bool fixed_size =
          std::all_of(split_sizes_vec.begin(), split_sizes_vec.end(),
                      [&entry](int n) { return n == entry; });

      for (int i = 0; i < num_split; ++i) {
        TensorShape output_shape(input_shape);
        output_shape.set_dim(split_dim, split_sizes_vec[i]);
        Tensor* result = nullptr;
        OP_REQUIRES_OK(context,
                       context->allocate_output(i, output_shape, &result));
        ptrs.Set(i, result->flat<T>().data());
        offsets.Set(i, offset);
        offset += split_sizes_vec[i] * suffix_dim_size;
      }
      offsets.Set(num_split, offset);
      OP_REQUIRES_OK(context, ptrs.Finalize());
      OP_REQUIRES_OK(context, offsets.Finalize());

      if (input.NumElements() > 0) {
        SplitVOpGPULaunch<T, Tlen>().Run(
            context->eigen_device<GPUDevice>(), fixed_size,
            input.flat<T>().data(), prefix_dim_size,
            input.NumElements() / prefix_dim_size, offsets.data(), ptrs.data());
        OP_REQUIRES(
            context, context->op_device_context()->stream()->ok(),
            errors::Internal("Launch of gpu kernel for SplitVOp failed"));
      }
    } else {
      Eigen::DenseIndex prefix_dim_size;
      Eigen::DenseIndex split_dim_size;
      Eigen::DenseIndex suffix_dim_size;

      std::tie(prefix_dim_size, split_dim_size, suffix_dim_size) =
          Base::template SetDims<Eigen::DenseIndex>(input_shape, split_dim);
      auto input_reshaped = input.shaped<T, 2>(
          {prefix_dim_size, split_dim_size * suffix_dim_size});

      Eigen::DSizes<Eigen::DenseIndex, 2> indices{0, 0};

      for (int i = 0; i < num_split; ++i) {
        TensorShape output_shape(input_shape);
        output_shape.set_dim(split_dim, split_sizes_vec[i]);
        Tensor* result = nullptr;
        OP_REQUIRES_OK(context,
                       context->allocate_output(i, output_shape, &result));

        Eigen::DSizes<Eigen::DenseIndex, 2> sizes{
            prefix_dim_size, split_sizes_vec[i] * suffix_dim_size};

        if (sizes.TotalSize() > 0) {
          auto result_shaped = result->shaped<T, 2>(
              {prefix_dim_size, split_sizes_vec[i] * suffix_dim_size});

          functor::SplitCustom<GPUDevice, T>()(
              context->eigen_device<GPUDevice>(), result_shaped, input_reshaped,
              indices, sizes);
        }
        indices[1] += split_sizes_vec[i] * suffix_dim_size;
      }
    }
  }


  void Compute(OpKernelContext* context) override {
    bool done = false;
    std::vector<Tlen> split_sizes_vec;
    Base::ComputeEasyCases(context, &done, &split_sizes_vec);
    if (!context->status().ok() || done) {
      return;
    }
    const int32_t num_split = Base::num_outputs();
    const Tensor& input = context->input(0);
    const TensorShape& input_shape = input.shape();
    const int32_t split_dim_orig = context->input(2).flat<int32>()(0);
    const int32_t split_dim =
        split_dim_orig < 0 ? split_dim_orig + input.dims() : split_dim_orig;

    // Android also uses int32 indexing, so check here also.
    OP_REQUIRES(
        context,
        FastBoundsCheck(input.NumElements(),
                        std::numeric_limits<Eigen::DenseIndex>::max()),
        errors::InvalidArgument("Split requires input size < ",
                                std::numeric_limits<Eigen::DenseIndex>::max()));

    Eigen::DenseIndex prefix_dim_size;
    Eigen::DenseIndex split_dim_size;
    Eigen::DenseIndex suffix_dim_size;

    std::tie(prefix_dim_size, split_dim_size, suffix_dim_size) =
        Base::template SetDims<Eigen::DenseIndex>(input_shape, split_dim);
    std::vector<int64_t> split_start_points(num_split);
    for (int i = 0; i < num_split; ++i) {
      if (i == 0) {
        split_start_points[i] = 0;
      } else {
        split_start_points[i] =
            split_start_points[i - 1] + split_sizes_vec[i - 1];
      }
    }

    if (prefix_dim_size == 1) {
      auto input_reshaped =
          input.shaped<T, 2>({split_dim_size, suffix_dim_size});
      auto make_sizes = [&](Eigen::DenseIndex split_size) {
        return Eigen::DSizes<Eigen::DenseIndex, 2>{split_size, suffix_dim_size};
      };
      auto reshape_result = [&](Tensor* result, Tlen split_size) {
        return result->shaped<T, 2>({split_size, suffix_dim_size});
      };
      SplitVOpCPUImpl<T, Tlen, decltype(input_reshaped), 2>{}(
          context, input_reshaped, split_start_points, input_shape, split_dim,
          prefix_dim_size, split_dim_size, suffix_dim_size, split_sizes_vec,
          make_sizes, reshape_result);
    } else {
      auto input_reshaped = input.shaped<T, 3>(
          {prefix_dim_size, split_dim_size, suffix_dim_size});
      auto make_sizes = [&](Eigen::DenseIndex split_size) {
        return Eigen::DSizes<Eigen::DenseIndex, 3>{prefix_dim_size, split_size,
                                                   suffix_dim_size};
      };
      auto reshape_result = [&](Tensor* result, Tlen split_size) {
        return result->shaped<T, 3>(
            {prefix_dim_size, split_size, suffix_dim_size});
      };
      SplitVOpCPUImpl<T, Tlen, decltype(input_reshaped), 3>{}(
          context, input_reshaped, split_start_points, input_shape, split_dim,
          prefix_dim_size, split_dim_size, suffix_dim_size, split_sizes_vec,
          make_sizes, reshape_result);
    }
  }


  static bool SplitHasAlignedOutputsInFirstDimension(
      const TensorShape& input_shape, int32_t split_dim,
      absl::Span<const Tlen> split_sizes) {
    if (split_dim != 0) {
      return false;
    }
    Tlen start = 0;
    for (const Tlen split_size : split_sizes) {
      if (!IsDim0SliceAligned<T>(input_shape, start, start + split_size)) {
        return false;
      }
      start += split_size;
    }
    return true;
  }


int ndpi_netbios_name_interpret(char *in, size_t in_len, char *out, u_int out_len) {
  u_int ret = 0, len, idx = in_len, out_idx = 0;

  len = (*in++)/2, in_len--;
  out_len--;
  out[out_idx] = 0;

  if((len > out_len) || (len < 1) || ((2*len) > in_len))
    return(-1);

  while((len--) && (out_idx < out_len)) {
    if((idx < 2) || (in[0] < 'A') || (in[0] > 'P') || (in[1] < 'A') || (in[1] > 'P')) {
      out[out_idx] = 0;
      break;
    }

    out[out_idx] = ((in[0] - 'A') << 4) + (in[1] - 'A');
    in += 2, idx -= 2;

    if(isprint(out[out_idx]))
      out_idx++, ret++;
  }

  /* Trim trailing whitespace from the returned string */
  if(out_idx > 0) {
    out[out_idx] = 0;
    out_idx--;

    while((out_idx > 0) && (out[out_idx] == ' ')) {
      out[out_idx] = 0;
      out_idx--;
    }

  }

  return(ret);
}

Added further checks

#define NDPI_CURRENT_PROTO NDPI_PROTOCOL_NETBIOS

struct netbios_header {
  u_int16_t transaction_id, flags, questions, answer_rrs, authority_rrs, additional_rrs;
}