text: #include "quant_decomp_zstd_op.h"
  #include <stdint.h>
  #include <zstd.h>
  #include "caffe2/core/tensor.h"
  #include "caffe2/proto/caffe2.pb.h"
  
  namespace caffe2 {
  
  namespace {
  
  #define REGISTER_TYPE(index, type)                                      \
    {                                                                     \
      index, [](TensorCPU* tensor_) -> uint8_t* {                         \
        return reinterpret_cast<uint8_t*>(tensor_->mutable_data<type>()); \
      }                                                                   \
    }
  
  // return a mutable pointer to the tensor in uint8_t format, the memory is
  //   allocated based on the type 'type_index'
  // supported type is defined in 'gTypeMapper'
  uint8_t* GetMutableData(int type_index, TensorCPU* tensor) {
    // see COMP_DATA_TYPE_MAPPER in mutils.py for the mapping
    static const std::map<int, std::function<uint8_t*(TensorCPU * tensor)>>
        gTypeMapper = {REGISTER_TYPE(TensorProto::UINT8, uint8_t),
                       REGISTER_TYPE(TensorProto::UINT16, uint16_t),
                       REGISTER_TYPE(TensorProto::INT32, int32_t),
                       REGISTER_TYPE(TensorProto::FLOAT, float)};
  
    CAFFE_ENFORCE_EQ(
        gTypeMapper.count(type_index),
        1,
        "Invalid type index " + caffe2::to_string(type_index) + ".");
    return gTypeMapper.at(type_index)(tensor);
  }
  
  const uint8_t* GetCompressedPtr(const TensorCPU& compressed, size_t* out_size) {
    CAFFE_ENFORCE(
        // array of uint8_t
        compressed.template IsType<uint8_t>() ||
        // array with one string
        compressed.template IsType<std::string>());
  
    if (compressed.template IsType<uint8_t>()) {
      *out_size = compressed.size();
      return compressed.data<uint8_t>();
    }
  
    // string type
    CAFFE_ENFORCE_EQ(compressed.size(), 1);
    auto& str = compressed.data<std::string>()[0];
    *out_size = str.size();
    return reinterpret_cast<const uint8_t*>(str.data());
  }
  
  // Deserialize the string to get TensorProtos, storing tensors in compressed
  // format
  TensorProtos GetTensorsProto(const TensorCPU& compressed) {
    size_t sz;
    auto* ptr = GetCompressedPtr(compressed, &sz);
    TensorProtos tensors;
    CAFFE_ENFORCE(tensors.ParseFromArray(ptr, sz));
    return tensors;
  }
  
  // Decompress tensor stored in compressed format
  // It is compressed using mutils.compress_data_list()
  void Decompress(const TensorProto& compressed, TensorCPU* outDecomp) {
    vector<TIndex> shape(compressed.dims().begin(), compressed.dims().end());
    // shape stores the dimensions of data before compression,
    //   see _compress_data_single() in mutils.py
    outDecomp->Resize(shape);
    auto* out_ptr = GetMutableData(compressed.data_type(), outDecomp);
  
    auto* src = reinterpret_cast<const uint8_t*>(compressed.byte_data().data());
    size_t comp_size = compressed.byte_data().size();
    size_t decomp_size = outDecomp->nbytes();
  
    // call zstd
    size_t dc_size = ZSTD_decompress(out_ptr, decomp_size, src, comp_size);
    CAFFE_ENFORCE(!ZSTD_isError(dc_size), ZSTD_getErrorName(dc_size));
    CAFFE_ENFORCE_EQ(decomp_size, dc_size);
  }
  
  } // namespace
  
  bool QuantDecompZstdOp::RunOnDevice() {
    const auto& op_compressed = Input(0);
  
    // Data could be an array of uint_t, or a string
    CAFFE_ENFORCE(
        // array of uint8_t
        op_compressed.template IsType<uint8_t>() ||
            // array with one string
            op_compressed.template IsType<std::string>(),
        op_compressed.meta().name());
  
    // op_compressed: compressed data, 1d
    if (op_compressed.template IsType<uint8_t>()) {
      CAFFE_ENFORCE_EQ(op_compressed.ndim(), 1, op_compressed.ndim());
    } else {
      // string type has 0 dimension
      CAFFE_ENFORCE_EQ(op_compressed.size(), 1, op_compressed.size());
    }
  
    auto tensors = GetTensorsProto(op_compressed);
    CAFFE_ENFORCE_EQ(tensors.protos_size(), OutputSize());
  
    for (int i = 0; i < OutputSize(); i++) {
      Decompress(tensors.protos(i), Output(i));
    }
  
    return true;
  }
  
  REGISTER_CPU_OPERATOR(QuantDecompZstd, QuantDecompZstdOp);
  
  OPERATOR_SCHEMA(QuantDecompZstd)
      .NumInputs(1)
      .NumOutputs(1, INT_MAX)
      .SetDoc(R"DOC(
   Decompress a set of tensors that are compressed using zstd.
   The data can be compressed using mutils.compress_data_list(), see
   quant_decomp_op_test.py for an example.
   The number of outputs depended on the input.
   )DOC")
      .Input(
          0,
          "compressed",
          "Compressed data in 1d tensor (uint8_t), "
          "or 0d tensor with one element in string type."
          "The data is compressed using mutils.compress_data_list().")
      .Output(0, "output0", "Decompressed data 0")
      .Output(1, "output1", "Decompressed data 1 if existed")
      .Output(2, "output2", "Decompressed data 2 if existed")
      .Output(3, "outputn", "Decompressed data n if existed");
  
  SHOULD_NOT_DO_GRADIENT(QuantDecompZstd);
  
  } // namespace caffe2
  