<commit_before>/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#if GOOGLE_CUDA

#define EIGEN_USE_GPU

#include "tensorflow/core/kernels/parameterized_truncated_normal_op.h"

#include <assert.h>
#include <stdio.h>
#include <cmath>

#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "tensorflow/core/framework/tensor_types.h"
#include "tensorflow/core/lib/random/philox_random.h"
#include "tensorflow/core/lib/random/random_distributions.h"
#include "tensorflow/core/util/cuda_kernel_helper.h"

#define UNROLL _Pragma("unroll")

namespace tensorflow {

class OpKernelContext;

namespace functor {

typedef Eigen::GpuDevice GPUDevice;

template <typename T>
__global__ void __launch_bounds__(1024)
    TruncatedNormalKernel(random::PhiloxRandom gen, T* data, int64 num_batches,
                          int64 samples_per_batch, int64 num_elements,
                          const T* means, bool single_mean, const T* stddevs,
                          bool single_stddev, const T* minvals,
                          bool single_minval, const T* maxvals,
                          bool single_maxval, int64 kMaxIterations) {
  const int32 max_samples_per_item = 2 * kMaxIterations;
  // Initial offset as given by CUDA_1D_KERNEL_LOOP.
  const int32 initial_offset = blockIdx.x * blockDim.x + threadIdx.x;
  gen.Skip(max_samples_per_item * initial_offset);
  typedef random::UniformDistribution<random::PhiloxRandom, T> Uniform;
  Uniform dist;
  const int kDistSize = Uniform::kResultElementCount;
  const T quietNaN = Eigen::NumTraits<T>::quiet_NaN();

  // We skip the total number of threads to get to the next element. To produce
  // deterministic results between devices, each element in the output array
  // skips max_samples_per_item in the generator. Then after generating this
  // item, we need to skip the samples for one element for every thread to get
  // to the next element that we actually process.
  const int32 samples_between_processed_elements =
      max_samples_per_item * (gridDim.x * blockDim.x);

  CUDA_1D_KERNEL_LOOP(offset, num_elements) {
    // Track how many more samples we need to skip before we process the next
    // element.
    int32 remaining_samples = samples_between_processed_elements;

    const int64 batch_id = offset / samples_per_batch;
    T mean = means[single_mean ? 0 : batch_id];
    const T input_stddev = stddevs[single_stddev ? 0 : batch_id];
    T minval = minvals[single_minval ? 0 : batch_id];
    T maxval = maxvals[single_maxval ? 0 : batch_id];

    // Flip the distribution if we can make the lower bound positive.
    T stddev;
    if (Eigen::numext::isinf(minval) || maxval < mean) {
      // Reverse all calculations. normMin and normMax will be flipped.
      // std::swap is a host function (not available in CUDA).
      T temp = minval;
      minval = maxval;
      maxval = temp;
      stddev = -input_stddev;
    } else {
      stddev = input_stddev;
    }

    // Calculate normalized samples, then scale them.
    const T normMin = (minval - mean) / stddev;
    const T normMax = (maxval - mean) / stddev;

    // Determine the method to use.
    const T sqrtFactor = Eigen::numext::sqrt((normMin * normMin) + T(4));
    const T cutoff =
        T(2) *
        Eigen::numext::exp(T(0.5) + (normMin * (normMin - sqrtFactor)) / T(4)) /
        (normMin + sqrtFactor);
    const T diff = normMax - normMin;

    // Validate the normalized min and max, because the originals may have been
    // flipped already.
    if (!(input_stddev > T(0) && normMin < normMax &&
          (Eigen::numext::isfinite(normMin) ||
           Eigen::numext::isfinite(normMax)))) {
      data[offset] = quietNaN;
    } else if (diff < cutoff) {
      // Sample from a uniform distribution on [normMin, normMax].

      // Vectorized intermediate calculations for uniform rejection sampling.
      // We always generate at most 4 samples.
      Eigen::array<T, 4> z;
      Eigen::array<T, 4> g;

      const T plusFactor = (normMin < T(0)) ? T(0) : normMin * normMin;

      int numIterations = 0;
      while (numIterations < kMaxIterations) {
        const auto rand = dist(&gen);
        remaining_samples -= gen.kResultElementCount;
        UNROLL for (int i = 0; i < kDistSize; i++) {
          z[i] = rand[i] * diff + normMin;
        }
        UNROLL for (int i = 0; i < kDistSize; i++) {
          g[i] = (plusFactor - z[i] * z[i]) / 2.0;
        }

        const auto u = dist(&gen);
        remaining_samples -= gen.kResultElementCount;
        UNROLL for (int i = 0; i < kDistSize; i++) {
          if (u[i] <= Eigen::numext::exp(g[i]) ||
              numIterations + 1 >= kMaxIterations) {
            // Accept the sample z.
            // If we run out of iterations, just use the current uniform
            // sample. Emperically, the probability of accepting each sample
            // is at least 50% for typical inputs, so we will always accept
            // by 100 iterations.
            // This introduces a slight inaccuracy when at least one bound
            // is large, minval is negative and maxval is positive.
            data[offset] = z[i] * stddev + mean;
            // Break out of the nested loop by updating numIterations.
            numIterations = kMaxIterations;
            break;
          } else {
            numIterations++;
          }
        }
      }
    } else {
      // Sample from an exponential distribution with alpha maximizing
      // acceptance probability, offset by normMin from the origin.
      // Accept only if less than normMax.
      const T alpha =
          (normMin + Eigen::numext::sqrt((normMin * normMin) + T(4))) / T(2);
      int numIterations = 0;
      while (numIterations < kMaxIterations) {
        auto rand = dist(&gen);
        remaining_samples -= gen.kResultElementCount;
        UNROLL for (int i = 0; i < kDistSize; i += 2) {
          const T z = -Eigen::numext::log(rand[i]) / alpha + normMin;
          const T x = normMin < alpha ? alpha - z : normMin - alpha;
          const T g = Eigen::numext::exp(-x * x / 2.0);
          const T u = rand[i + 1];
          if ((u <= g && z < normMax) || numIterations + 1 >= kMaxIterations) {
            data[offset] = z * stddev + mean;
            // Break out of the nested loop by updating numIterations.
            numIterations = kMaxIterations;
            break;
          } else {
            numIterations++;
          }
        }
      }
    }

    gen.Skip(remaining_samples);
  }
}

// Partial specialization for GPU
template <typename T>
struct TruncatedNormalFunctor<GPUDevice, T> {
  static const int kMaxIterations = 100;

  void operator()(OpKernelContext* ctx, const GPUDevice& d, int64 num_batches,
                  int64 samples_per_batch, int64 num_elements,
                  typename TTypes<T>::ConstFlat means,
                  typename TTypes<T>::ConstFlat stddevs,
                  typename TTypes<T>::ConstFlat minvals,
                  typename TTypes<T>::ConstFlat maxvals,
                  const random::PhiloxRandom& gen,
                  typename TTypes<T>::Flat output) {
    const auto config = GetCudaLaunchConfig(num_elements, d);

    TruncatedNormalKernel<
        T><<<config.block_count, config.thread_per_block, 0, d.stream()>>>(
        gen, output.data(), num_batches, samples_per_batch, num_elements,
        means.data(), means.dimension(0) == 1, stddevs.data(),
        stddevs.dimension(0) == 1, minvals.data(), minvals.dimension(0) == 1,
        maxvals.data(), maxvals.dimension(0) == 1, kMaxIterations);
  };
};

// Explicit instantiation of the GPU distributions functors
template struct TruncatedNormalFunctor<GPUDevice, Eigen::half>;
template struct TruncatedNormalFunctor<GPUDevice, float>;
template struct TruncatedNormalFunctor<GPUDevice, double>;

}  // namespace functor
}  // namespace tensorflow

#endif  // GOOGLE_CUDA
<commit_msg>fix windows gpu build (#5421)<commit_after>/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#if GOOGLE_CUDA

#define EIGEN_USE_GPU

#include "tensorflow/core/kernels/parameterized_truncated_normal_op.h"

#include <assert.h>
#include <stdio.h>
#include <cmath>

#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "tensorflow/core/framework/tensor_types.h"
#include "tensorflow/core/lib/random/philox_random.h"
#include "tensorflow/core/lib/random/random_distributions.h"
#include "tensorflow/core/util/cuda_kernel_helper.h"

#ifdef COMPILER_MSVC
// msvc does not support unroll. One could try the loop pragma but we need to
// take a closer look if this generates better code in this case. For now let
// the compiler take care of of it.
#define UNROLL
#else
#define UNROLL _Pragma("unroll")
#endif

namespace tensorflow {

class OpKernelContext;

namespace functor {

typedef Eigen::GpuDevice GPUDevice;

template <typename T>
__global__ void __launch_bounds__(1024)
    TruncatedNormalKernel(random::PhiloxRandom gen, T* data, int64 num_batches,
                          int64 samples_per_batch, int64 num_elements,
                          const T* means, bool single_mean, const T* stddevs,
                          bool single_stddev, const T* minvals,
                          bool single_minval, const T* maxvals,
                          bool single_maxval, int64 kMaxIterations) {
  const int32 max_samples_per_item = 2 * kMaxIterations;
  // Initial offset as given by CUDA_1D_KERNEL_LOOP.
  const int32 initial_offset = blockIdx.x * blockDim.x + threadIdx.x;
  gen.Skip(max_samples_per_item * initial_offset);
  typedef random::UniformDistribution<random::PhiloxRandom, T> Uniform;
  Uniform dist;
  const int kDistSize = Uniform::kResultElementCount;
  const T quietNaN = Eigen::NumTraits<T>::quiet_NaN();

  // We skip the total number of threads to get to the next element. To produce
  // deterministic results between devices, each element in the output array
  // skips max_samples_per_item in the generator. Then after generating this
  // item, we need to skip the samples for one element for every thread to get
  // to the next element that we actually process.
  const int32 samples_between_processed_elements =
      max_samples_per_item * (gridDim.x * blockDim.x);

  CUDA_1D_KERNEL_LOOP(offset, num_elements) {
    // Track how many more samples we need to skip before we process the next
    // element.
    int32 remaining_samples = samples_between_processed_elements;

    const int64 batch_id = offset / samples_per_batch;
    T mean = means[single_mean ? 0 : batch_id];
    const T input_stddev = stddevs[single_stddev ? 0 : batch_id];
    T minval = minvals[single_minval ? 0 : batch_id];
    T maxval = maxvals[single_maxval ? 0 : batch_id];

    // Flip the distribution if we can make the lower bound positive.
    T stddev;
    if (Eigen::numext::isinf(minval) || maxval < mean) {
      // Reverse all calculations. normMin and normMax will be flipped.
      // std::swap is a host function (not available in CUDA).
      T temp = minval;
      minval = maxval;
      maxval = temp;
      stddev = -input_stddev;
    } else {
      stddev = input_stddev;
    }

    // Calculate normalized samples, then scale them.
    const T normMin = (minval - mean) / stddev;
    const T normMax = (maxval - mean) / stddev;

    // Determine the method to use.
    const T sqrtFactor = Eigen::numext::sqrt((normMin * normMin) + T(4));
    const T cutoff =
        T(2) *
        Eigen::numext::exp(T(0.5) + (normMin * (normMin - sqrtFactor)) / T(4)) /
        (normMin + sqrtFactor);
    const T diff = normMax - normMin;
    const T two = T(2.0);

    // Validate the normalized min and max, because the originals may have been
    // flipped already.
    if (!(input_stddev > T(0) && normMin < normMax &&
          (Eigen::numext::isfinite(normMin) ||
           Eigen::numext::isfinite(normMax)))) {
      data[offset] = quietNaN;
    } else if (diff < cutoff) {
      // Sample from a uniform distribution on [normMin, normMax].

      // Vectorized intermediate calculations for uniform rejection sampling.
      // We always generate at most 4 samples.
      Eigen::array<T, 4> z;
      Eigen::array<T, 4> g;

      const T plusFactor = (normMin < T(0)) ? T(0) : normMin * normMin;

      int numIterations = 0;
      while (numIterations < kMaxIterations) {
        const auto rand = dist(&gen);
        remaining_samples -= gen.kResultElementCount;
        UNROLL for (int i = 0; i < kDistSize; i++) {
          z[i] = rand[i] * diff + normMin;
        }
        UNROLL for (int i = 0; i < kDistSize; i++) {
          g[i] = (plusFactor - z[i] * z[i]) / two;
        }

        const auto u = dist(&gen);
        remaining_samples -= gen.kResultElementCount;
        UNROLL for (int i = 0; i < kDistSize; i++) {
          if (u[i] <= Eigen::numext::exp(g[i]) ||
              numIterations + 1 >= kMaxIterations) {
            // Accept the sample z.
            // If we run out of iterations, just use the current uniform
            // sample. Emperically, the probability of accepting each sample
            // is at least 50% for typical inputs, so we will always accept
            // by 100 iterations.
            // This introduces a slight inaccuracy when at least one bound
            // is large, minval is negative and maxval is positive.
            data[offset] = z[i] * stddev + mean;
            // Break out of the nested loop by updating numIterations.
            numIterations = kMaxIterations;
            break;
          } else {
            numIterations++;
          }
        }
      }
    } else {
      // Sample from an exponential distribution with alpha maximizing
      // acceptance probability, offset by normMin from the origin.
      // Accept only if less than normMax.
      const T alpha =
          (normMin + Eigen::numext::sqrt((normMin * normMin) + T(4))) / T(2);
      int numIterations = 0;
      while (numIterations < kMaxIterations) {
        auto rand = dist(&gen);
        remaining_samples -= gen.kResultElementCount;
        UNROLL for (int i = 0; i < kDistSize; i += 2) {
          const T z = -Eigen::numext::log(rand[i]) / alpha + normMin;
          const T x = normMin < alpha ? alpha - z : normMin - alpha;
          const T g = Eigen::numext::exp(-x * x / two);
          const T u = rand[i + 1];
          if ((u <= g && z < normMax) || numIterations + 1 >= kMaxIterations) {
            data[offset] = z * stddev + mean;
            // Break out of the nested loop by updating numIterations.
            numIterations = kMaxIterations;
            break;
          } else {
            numIterations++;
          }
        }
      }
    }

    gen.Skip(remaining_samples);
  }
}

// Partial specialization for GPU
template <typename T>
struct TruncatedNormalFunctor<GPUDevice, T> {
  static const int kMaxIterations = 100;

  void operator()(OpKernelContext* ctx, const GPUDevice& d, int64 num_batches,
                  int64 samples_per_batch, int64 num_elements,
                  typename TTypes<T>::ConstFlat means,
                  typename TTypes<T>::ConstFlat stddevs,
                  typename TTypes<T>::ConstFlat minvals,
                  typename TTypes<T>::ConstFlat maxvals,
                  const random::PhiloxRandom& gen,
                  typename TTypes<T>::Flat output) {
    const auto config = GetCudaLaunchConfig(num_elements, d);

    TruncatedNormalKernel<
        T><<<config.block_count, config.thread_per_block, 0, d.stream()>>>(
        gen, output.data(), num_batches, samples_per_batch, num_elements,
        means.data(), means.dimension(0) == 1, stddevs.data(),
        stddevs.dimension(0) == 1, minvals.data(), minvals.dimension(0) == 1,
        maxvals.data(), maxvals.dimension(0) == 1, kMaxIterations);
  };
};

// Explicit instantiation of the GPU distributions functors
template struct TruncatedNormalFunctor<GPUDevice, Eigen::half>;
template struct TruncatedNormalFunctor<GPUDevice, float>;
template struct TruncatedNormalFunctor<GPUDevice, double>;

}  // namespace functor
}  // namespace tensorflow

#endif  // GOOGLE_CUDA
<|endoftext|>


<commit_before>/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include <limits>

#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/experimental/kernels/gru_cell.h"
#include "tensorflow/lite/kernels/cpu_backend_context.h"
#include "tensorflow/lite/kernels/internal/tensor.h"
#include "tensorflow/lite/kernels/kernel_util.h"

namespace tflite {
namespace ops {
namespace experimental {
namespace unidirectional_sequence_gru {
namespace {

void GruImpl(const TfLiteTensor* input, const TfLiteTensor* input_state,
             const TfLiteTensor* gate_weight, const TfLiteTensor* gate_bias,
             const TfLiteTensor* candidate_weight,
             const TfLiteTensor* candidate_bias, TfLiteTensor* output,
             TfLiteTensor* output_state, TfLiteTensor* activation,
             TfLiteTensor* concat,
             tflite::CpuBackendContext* cpu_backend_context) {
  const int n_time = input->dims->data[0];
  const int n_batch = input->dims->data[1];
  const int n_input = input->dims->data[2];
  const int n_output = output->dims->data[2];
  const int n_batch_input = n_batch * n_input;
  const int n_batch_output = n_batch * n_output;
  const RuntimeShape input_shape({n_batch, n_input});
  const float* input_data = GetTensorData<float>(input);
  const RuntimeShape state_shape = GetTensorShape(input_state);
  const float* input_state_data = GetTensorData<float>(input_state);
  const RuntimeShape gate_weight_shape = GetTensorShape(gate_weight);
  const float* gate_weight_data = GetTensorData<float>(gate_weight);
  const RuntimeShape gate_bias_shape = GetTensorShape(gate_bias);
  const float* gate_bias_data = GetTensorData<float>(gate_bias);
  const RuntimeShape candidate_weight_shape = GetTensorShape(candidate_weight);
  const float* candidate_weight_data = GetTensorData<float>(candidate_weight);
  const RuntimeShape candidate_bias_shape = GetTensorShape(candidate_bias);
  const float* candidate_bias_data = GetTensorData<float>(candidate_bias);
  const RuntimeShape activation_shape = GetTensorShape(activation);
  const RuntimeShape output_shape = RuntimeShape({n_batch, n_output});
  float* output_data = GetTensorData<float>(output);
  float* output_state_data = GetTensorData<float>(output_state);
  float* activation_data = GetTensorData<float>(activation);
  const RuntimeShape concat_shape = GetTensorShape(concat);
  float* concat_data = GetTensorData<float>(concat);
  tflite::FullyConnectedParams fc_params;
  fc_params.float_activation_min = std::numeric_limits<float>::lowest();
  fc_params.float_activation_max = std::numeric_limits<float>::max();
  for (int i = 0; i < n_time; ++i) {
    gru_cell::GruCell(
        input_shape, input_data, state_shape, input_state_data,
        gate_weight_shape, gate_weight_data, gate_bias_shape, gate_bias_data,
        candidate_weight_shape, candidate_weight_data, candidate_bias_shape,
        candidate_bias_data, output_shape, output_data, output_state_data,
        activation_shape, activation_data, concat_shape, concat_data, fc_params,
        cpu_backend_context);
    input_data += n_batch_input;
    output_data += n_batch_output;
    input_state_data = output_state_data;
  }
}

}  // namespace

enum InputTensor {
  // Input tensor of size [n_time, n_batch, n_input]
  kInput = 0,
  // Input state tensor of size [n_batch, n_output]
  kInputState = 1,
  // Gate weight tensor of size [2*n_output, n_input+n_output]
  kGateWeight = 2,
  // Gate bias tensor of size [2*n_output]
  kGateBias = 3,
  // Candidate weight tensor of size [n_output, n_input+n_output]
  kCandidateWeight = 4,
  // Candidate bias tensor of size [n_output]
  kCandidateBias = 5,
  kInputNum = 6
};

enum OutputTensor {
  // Input tensor of size [n_time, n_batch, n_output]
  kOutput = 0,
  // Output state tensor of size [n_batch, n_output]
  kOutputState = 1,
  kOutputNum = 2
};

enum TemporaryTensor {
  // Scratch buffer for activation of size [n_batch, 2*n_output]
  kActivation = 0,
  // Scratch buffer for activation of size [n_batch, n_input+n_output]
  kConcat = 1,
  kTemporaryNum = 2
};

void* Init(TfLiteContext* context, const char* buffer, size_t length) {
  auto* scratch_tensor_index = new int;
  context->AddTensors(context, kTemporaryNum, scratch_tensor_index);
  return scratch_tensor_index;
}

void Free(TfLiteContext* context, void* buffer) {
  delete reinterpret_cast<int*>(buffer);
}

TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
  int* scratch_tensor_index = reinterpret_cast<int*>(node->user_data);

  TF_LITE_ENSURE_EQ(context, node->inputs->size, kInputNum);
  TF_LITE_ENSURE_EQ(context, node->outputs->size, kOutputNum);

  // input's dim = [n_time, n_batch, n_input]
  const TfLiteTensor* input;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInput, &input));
  TF_LITE_ENSURE_EQ(context, input->dims->size, 3);
  const int n_time = input->dims->data[0];
  const int n_batch = input->dims->data[1];
  const int n_input = input->dims->data[2];

  // input_state's dim = [n_batch, n_output]
  const TfLiteTensor* input_state;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kInputState, &input_state));
  TF_LITE_ENSURE_EQ(context, input_state->dims->size, 2);
  TF_LITE_ENSURE_EQ(context, input_state->dims->data[0], n_batch);
  const int n_output = input_state->dims->data[1];

  // gate_weight' dim = [2 * n_output, n_input + n_output]
  const TfLiteTensor* gate_weight;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kGateWeight, &gate_weight));
  TF_LITE_ENSURE_EQ(context, gate_weight->dims->size, 2);
  TF_LITE_ENSURE_EQ(context, gate_weight->dims->data[0], 2 * n_output);
  TF_LITE_ENSURE_EQ(context, gate_weight->dims->data[1], n_input + n_output);

  // gate_bias' dim = [2 * n_output]
  const TfLiteTensor* gate_bias;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kGateBias, &gate_bias));
  TF_LITE_ENSURE_EQ(context, gate_bias->dims->size, 1);
  TF_LITE_ENSURE_EQ(context, gate_bias->dims->data[0], 2 * n_output);

  // candidate_weight' dim = [n_output, n_input + n_output]
  const TfLiteTensor* candidate_weight;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kCandidateWeight,
                                          &candidate_weight));
  TF_LITE_ENSURE_EQ(context, candidate_weight->dims->size, 2);
  TF_LITE_ENSURE_EQ(context, candidate_weight->dims->data[0], n_output);
  TF_LITE_ENSURE_EQ(context, candidate_weight->dims->data[1],
                    n_input + n_output);

  // candidate_bias' dim = [n_output]
  const TfLiteTensor* candidate_bias;
  TF_LITE_ENSURE_OK(
      context, GetInputSafe(context, node, kCandidateBias, &candidate_bias));
  TF_LITE_ENSURE_EQ(context, candidate_bias->dims->size, 1);
  TF_LITE_ENSURE_EQ(context, candidate_bias->dims->data[0], n_output);

  // output's dim = [n_time, n_batch, n_output]
  TfLiteTensor* output;
  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, kOutput, &output));
  TfLiteIntArray* output_size = TfLiteIntArrayCreate(3);
  output_size->data[0] = n_time;
  output_size->data[1] = n_batch;
  output_size->data[2] = n_output;
  TF_LITE_ENSURE_OK(context,
                    context->ResizeTensor(context, output, output_size));

  // output_state's dim = [n_batch, n_output]
  TfLiteTensor* output_state;
  TF_LITE_ENSURE_OK(context,
                    GetOutputSafe(context, node, kOutputState, &output_state));
  TF_LITE_ENSURE_OK(
      context, context->ResizeTensor(context, output_state,
                                     TfLiteIntArrayCopy(input_state->dims)));

  TfLiteIntArrayFree(node->temporaries);
  node->temporaries = TfLiteIntArrayCreate(kTemporaryNum);

  // activation's dim = [n_batch, 2 * n_output]
  node->temporaries->data[kActivation] = *scratch_tensor_index;
  TfLiteTensor* activation;
  TF_LITE_ENSURE_OK(context,
                    GetTemporarySafe(context, node, kActivation, &activation));
  activation->type = input->type;
  activation->allocation_type = kTfLiteArenaRw;
  TfLiteIntArray* activation_size = TfLiteIntArrayCreate(2);
  activation_size->data[0] = n_batch;
  activation_size->data[1] = 2 * n_output;
  TF_LITE_ENSURE_OK(
      context, context->ResizeTensor(context, activation, activation_size));

  // concat's dim  = [n_batch, n_input + n_output]
  node->temporaries->data[kConcat] = (*scratch_tensor_index) + kConcat;
  TfLiteTensor* concat;
  TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, kConcat, &concat));
  concat->type = input->type;
  concat->allocation_type = kTfLiteArenaRw;
  TfLiteIntArray* concat_size = TfLiteIntArrayCreate(2);
  concat_size->data[0] = n_batch;
  concat_size->data[1] = n_input + n_output;
  TF_LITE_ENSURE_OK(context,
                    context->ResizeTensor(context, concat, concat_size));

  return kTfLiteOk;
}

TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
  const TfLiteTensor* input;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInput, &input));
  const TfLiteTensor* input_state;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kInputState, &input_state));
  const TfLiteTensor* gate_weight;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kGateWeight, &gate_weight));
  const TfLiteTensor* gate_bias;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kGateBias, &gate_bias));
  const TfLiteTensor* candidate_weight;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kCandidateWeight,
                                          &candidate_weight));
  const TfLiteTensor* candidate_bias;
  TF_LITE_ENSURE_OK(
      context, GetInputSafe(context, node, kCandidateBias, &candidate_bias));
  TfLiteTensor* output;
  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, kOutput, &output));
  TfLiteTensor* output_state;
  TF_LITE_ENSURE_OK(context,
                    GetOutputSafe(context, node, kOutputState, &output_state));
  TfLiteTensor* activation;
  TF_LITE_ENSURE_OK(context,
                    GetTemporarySafe(context, node, kActivation, &activation));
  TfLiteTensor* concat;
  TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, kConcat, &concat));
  auto cpu_backend_context = CpuBackendContext::GetFromContext(context);

  if (gate_weight->type == kTfLiteFloat32) {
    GruImpl(input, input_state, gate_weight, gate_bias, candidate_weight,
            candidate_bias, output, output_state, activation, concat,
            cpu_backend_context);
  } else {
    context->ReportError(context,
                         "Unsupported combination of data types for GruCell");
    return kTfLiteError;
  }

  return kTfLiteOk;
}

}  // namespace unidirectional_sequence_gru

TfLiteRegistration* Register_UNIDIRECTIONAL_SEQUENCE_GRU() {
  static TfLiteRegistration r = {
      unidirectional_sequence_gru::Init, unidirectional_sequence_gru::Free,
      unidirectional_sequence_gru::Prepare, unidirectional_sequence_gru::Eval};
  return &r;
}

}  // namespace experimental
}  // namespace ops
}  // namespace tflite
<commit_msg>Fix asan test for unidirectional_sequence_gru by setting fully_connected lhs/rhs cacheable.<commit_after>/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include <limits>

#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/experimental/kernels/gru_cell.h"
#include "tensorflow/lite/kernels/cpu_backend_context.h"
#include "tensorflow/lite/kernels/internal/tensor.h"
#include "tensorflow/lite/kernels/kernel_util.h"

namespace tflite {
namespace ops {
namespace experimental {
namespace unidirectional_sequence_gru {
namespace {

void GruImpl(const TfLiteTensor* input, const TfLiteTensor* input_state,
             const TfLiteTensor* gate_weight, const TfLiteTensor* gate_bias,
             const TfLiteTensor* candidate_weight,
             const TfLiteTensor* candidate_bias, TfLiteTensor* output,
             TfLiteTensor* output_state, TfLiteTensor* activation,
             TfLiteTensor* concat,
             tflite::CpuBackendContext* cpu_backend_context) {
  const int n_time = input->dims->data[0];
  const int n_batch = input->dims->data[1];
  const int n_input = input->dims->data[2];
  const int n_output = output->dims->data[2];
  const int n_batch_input = n_batch * n_input;
  const int n_batch_output = n_batch * n_output;
  const RuntimeShape input_shape({n_batch, n_input});
  const float* input_data = GetTensorData<float>(input);
  const RuntimeShape state_shape = GetTensorShape(input_state);
  const float* input_state_data = GetTensorData<float>(input_state);
  const RuntimeShape gate_weight_shape = GetTensorShape(gate_weight);
  const float* gate_weight_data = GetTensorData<float>(gate_weight);
  const RuntimeShape gate_bias_shape = GetTensorShape(gate_bias);
  const float* gate_bias_data = GetTensorData<float>(gate_bias);
  const RuntimeShape candidate_weight_shape = GetTensorShape(candidate_weight);
  const float* candidate_weight_data = GetTensorData<float>(candidate_weight);
  const RuntimeShape candidate_bias_shape = GetTensorShape(candidate_bias);
  const float* candidate_bias_data = GetTensorData<float>(candidate_bias);
  const RuntimeShape activation_shape = GetTensorShape(activation);
  const RuntimeShape output_shape = RuntimeShape({n_batch, n_output});
  float* output_data = GetTensorData<float>(output);
  float* output_state_data = GetTensorData<float>(output_state);
  float* activation_data = GetTensorData<float>(activation);
  const RuntimeShape concat_shape = GetTensorShape(concat);
  float* concat_data = GetTensorData<float>(concat);
  tflite::FullyConnectedParams fc_params;
  fc_params.float_activation_min = std::numeric_limits<float>::lowest();
  fc_params.float_activation_max = std::numeric_limits<float>::max();

  // The lhs is cacheable only when both gate weight & candidate weight are both
  // constants.
  fc_params.lhs_cacheable =
      IsConstantTensor(gate_weight) && IsConstantTensor(candidate_weight);
  fc_params.rhs_cacheable = false;
  for (int i = 0; i < n_time; ++i) {
    gru_cell::GruCell(
        input_shape, input_data, state_shape, input_state_data,
        gate_weight_shape, gate_weight_data, gate_bias_shape, gate_bias_data,
        candidate_weight_shape, candidate_weight_data, candidate_bias_shape,
        candidate_bias_data, output_shape, output_data, output_state_data,
        activation_shape, activation_data, concat_shape, concat_data, fc_params,
        cpu_backend_context);
    input_data += n_batch_input;
    output_data += n_batch_output;
    input_state_data = output_state_data;
  }
}

}  // namespace

enum InputTensor {
  // Input tensor of size [n_time, n_batch, n_input]
  kInput = 0,
  // Input state tensor of size [n_batch, n_output]
  kInputState = 1,
  // Gate weight tensor of size [2*n_output, n_input+n_output]
  kGateWeight = 2,
  // Gate bias tensor of size [2*n_output]
  kGateBias = 3,
  // Candidate weight tensor of size [n_output, n_input+n_output]
  kCandidateWeight = 4,
  // Candidate bias tensor of size [n_output]
  kCandidateBias = 5,
  kInputNum = 6
};

enum OutputTensor {
  // Input tensor of size [n_time, n_batch, n_output]
  kOutput = 0,
  // Output state tensor of size [n_batch, n_output]
  kOutputState = 1,
  kOutputNum = 2
};

enum TemporaryTensor {
  // Scratch buffer for activation of size [n_batch, 2*n_output]
  kActivation = 0,
  // Scratch buffer for activation of size [n_batch, n_input+n_output]
  kConcat = 1,
  kTemporaryNum = 2
};

void* Init(TfLiteContext* context, const char* buffer, size_t length) {
  auto* scratch_tensor_index = new int;
  context->AddTensors(context, kTemporaryNum, scratch_tensor_index);
  return scratch_tensor_index;
}

void Free(TfLiteContext* context, void* buffer) {
  delete reinterpret_cast<int*>(buffer);
}

TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
  int* scratch_tensor_index = reinterpret_cast<int*>(node->user_data);

  TF_LITE_ENSURE_EQ(context, node->inputs->size, kInputNum);
  TF_LITE_ENSURE_EQ(context, node->outputs->size, kOutputNum);

  // input's dim = [n_time, n_batch, n_input]
  const TfLiteTensor* input;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInput, &input));
  TF_LITE_ENSURE_EQ(context, input->dims->size, 3);
  const int n_time = input->dims->data[0];
  const int n_batch = input->dims->data[1];
  const int n_input = input->dims->data[2];

  // input_state's dim = [n_batch, n_output]
  const TfLiteTensor* input_state;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kInputState, &input_state));
  TF_LITE_ENSURE_EQ(context, input_state->dims->size, 2);
  TF_LITE_ENSURE_EQ(context, input_state->dims->data[0], n_batch);
  const int n_output = input_state->dims->data[1];

  // gate_weight' dim = [2 * n_output, n_input + n_output]
  const TfLiteTensor* gate_weight;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kGateWeight, &gate_weight));
  TF_LITE_ENSURE_EQ(context, gate_weight->dims->size, 2);
  TF_LITE_ENSURE_EQ(context, gate_weight->dims->data[0], 2 * n_output);
  TF_LITE_ENSURE_EQ(context, gate_weight->dims->data[1], n_input + n_output);

  // gate_bias' dim = [2 * n_output]
  const TfLiteTensor* gate_bias;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kGateBias, &gate_bias));
  TF_LITE_ENSURE_EQ(context, gate_bias->dims->size, 1);
  TF_LITE_ENSURE_EQ(context, gate_bias->dims->data[0], 2 * n_output);

  // candidate_weight' dim = [n_output, n_input + n_output]
  const TfLiteTensor* candidate_weight;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kCandidateWeight,
                                          &candidate_weight));
  TF_LITE_ENSURE_EQ(context, candidate_weight->dims->size, 2);
  TF_LITE_ENSURE_EQ(context, candidate_weight->dims->data[0], n_output);
  TF_LITE_ENSURE_EQ(context, candidate_weight->dims->data[1],
                    n_input + n_output);

  // candidate_bias' dim = [n_output]
  const TfLiteTensor* candidate_bias;
  TF_LITE_ENSURE_OK(
      context, GetInputSafe(context, node, kCandidateBias, &candidate_bias));
  TF_LITE_ENSURE_EQ(context, candidate_bias->dims->size, 1);
  TF_LITE_ENSURE_EQ(context, candidate_bias->dims->data[0], n_output);

  // output's dim = [n_time, n_batch, n_output]
  TfLiteTensor* output;
  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, kOutput, &output));
  TfLiteIntArray* output_size = TfLiteIntArrayCreate(3);
  output_size->data[0] = n_time;
  output_size->data[1] = n_batch;
  output_size->data[2] = n_output;
  TF_LITE_ENSURE_OK(context,
                    context->ResizeTensor(context, output, output_size));

  // output_state's dim = [n_batch, n_output]
  TfLiteTensor* output_state;
  TF_LITE_ENSURE_OK(context,
                    GetOutputSafe(context, node, kOutputState, &output_state));
  TF_LITE_ENSURE_OK(
      context, context->ResizeTensor(context, output_state,
                                     TfLiteIntArrayCopy(input_state->dims)));

  TfLiteIntArrayFree(node->temporaries);
  node->temporaries = TfLiteIntArrayCreate(kTemporaryNum);

  // activation's dim = [n_batch, 2 * n_output]
  node->temporaries->data[kActivation] = *scratch_tensor_index;
  TfLiteTensor* activation;
  TF_LITE_ENSURE_OK(context,
                    GetTemporarySafe(context, node, kActivation, &activation));
  activation->type = input->type;
  activation->allocation_type = kTfLiteArenaRw;
  TfLiteIntArray* activation_size = TfLiteIntArrayCreate(2);
  activation_size->data[0] = n_batch;
  activation_size->data[1] = 2 * n_output;
  TF_LITE_ENSURE_OK(
      context, context->ResizeTensor(context, activation, activation_size));

  // concat's dim  = [n_batch, n_input + n_output]
  node->temporaries->data[kConcat] = (*scratch_tensor_index) + kConcat;
  TfLiteTensor* concat;
  TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, kConcat, &concat));
  concat->type = input->type;
  concat->allocation_type = kTfLiteArenaRw;
  TfLiteIntArray* concat_size = TfLiteIntArrayCreate(2);
  concat_size->data[0] = n_batch;
  concat_size->data[1] = n_input + n_output;
  TF_LITE_ENSURE_OK(context,
                    context->ResizeTensor(context, concat, concat_size));

  return kTfLiteOk;
}

TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
  const TfLiteTensor* input;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInput, &input));
  const TfLiteTensor* input_state;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kInputState, &input_state));
  const TfLiteTensor* gate_weight;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kGateWeight, &gate_weight));
  const TfLiteTensor* gate_bias;
  TF_LITE_ENSURE_OK(context,
                    GetInputSafe(context, node, kGateBias, &gate_bias));
  const TfLiteTensor* candidate_weight;
  TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kCandidateWeight,
                                          &candidate_weight));
  const TfLiteTensor* candidate_bias;
  TF_LITE_ENSURE_OK(
      context, GetInputSafe(context, node, kCandidateBias, &candidate_bias));
  TfLiteTensor* output;
  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, kOutput, &output));
  TfLiteTensor* output_state;
  TF_LITE_ENSURE_OK(context,
                    GetOutputSafe(context, node, kOutputState, &output_state));
  TfLiteTensor* activation;
  TF_LITE_ENSURE_OK(context,
                    GetTemporarySafe(context, node, kActivation, &activation));
  TfLiteTensor* concat;
  TF_LITE_ENSURE_OK(context, GetTemporarySafe(context, node, kConcat, &concat));
  auto cpu_backend_context = CpuBackendContext::GetFromContext(context);

  if (gate_weight->type == kTfLiteFloat32) {
    GruImpl(input, input_state, gate_weight, gate_bias, candidate_weight,
            candidate_bias, output, output_state, activation, concat,
            cpu_backend_context);
  } else {
    context->ReportError(context,
                         "Unsupported combination of data types for GruCell");
    return kTfLiteError;
  }

  return kTfLiteOk;
}

}  // namespace unidirectional_sequence_gru

TfLiteRegistration* Register_UNIDIRECTIONAL_SEQUENCE_GRU() {
  static TfLiteRegistration r = {
      unidirectional_sequence_gru::Init, unidirectional_sequence_gru::Free,
      unidirectional_sequence_gru::Prepare, unidirectional_sequence_gru::Eval};
  return &r;
}

}  // namespace experimental
}  // namespace ops
}  // namespace tflite
<|endoftext|>


<commit_before>/* microflo_component yaml
name: ReadCapacitivePin
description: "Emits true if measured capacitance (in iterations) on @pin exceeds @threshold"
inports:
  trigger:
    type: all
    description: ""
  pin:
    type: all
    description: ""
  threshold:
    type: all
    description: ""
outports:
  out:
    type: all
    description: ""
microflo_component */
#ifdef ARDUINO0

// readCapacitivePin
//  Input: Arduino pin number
//  Output: A number, from 0 to 17 expressing
//          how much capacitance is on the pin
//  When you touch the pin, or whatever you have
//  attached to it, the number will get higher
//  In order for this to work now,
// The pin should have a 1+Megaohm resistor pulling
//  it up to +5v.
uint8_t readCapacitivePin(int pinToMeasure) {
  // This is how you declare a variable which
  //  will hold the PORT, PIN, and DDR registers
  //  on an AVR
  volatile uint8_t* port;
  volatile uint8_t* ddr;
  volatile uint8_t* pin;
  // Here we translate the input pin number from
  //  Arduino pin number to the AVR PORT, PIN, DDR,
  //  and which bit of those registers we care about.
  byte bitmask;
  if ((pinToMeasure >= 0) && (pinToMeasure <= 7)){
    port = &PORTD;
    ddr = &DDRD;
    bitmask = 1 << pinToMeasure;
    pin = &PIND;
  }
  if ((pinToMeasure > 7) && (pinToMeasure <= 13)){
    port = &PORTB;
    ddr = &DDRB;
    bitmask = 1 << (pinToMeasure - 8);
    pin = &PINB;
  }
  if ((pinToMeasure > 13) && (pinToMeasure <= 19)){
    port = &PORTC;
    ddr = &DDRC;
    bitmask = 1 << (pinToMeasure - 13);
    pin = &PINC;
  }
  // Discharge the pin first by setting it low and output
  *port &= ~(bitmask);
  *ddr  |= bitmask;
  delay(1);
  // Make the pin an input WITHOUT the internal pull-up on
  *ddr &= ~(bitmask);
  // Now see how long the pin to get pulled up
  int cycles = 16000;
  for(int i = 0; i < cycles; i++){
    if (*pin & bitmask){
      cycles = i;
      break;
    }
  }
  // Discharge the pin again by setting it low and output
  //  It's important to leave the pins low if you want to
  //  be able to touch more than 1 sensor at a time - if
  //  the sensor is left pulled high, when you touch
  //  two sensors, your body will transfer the charge between
  //  sensors.
  *port &= ~(bitmask);
  *ddr  |= bitmask;

  return cycles;
}

class ReadCapacitivePin : public SingleOutputComponent {
public:
    virtual void process(Packet in, MicroFlo::PortId port) {
        using namespace ReadCapacitivePinPorts;

        if (in.isSetup()) {
            currentState = false;
            pin = 1;
            touchedThreshold = 60;
        } else if (port == InPorts::trigger) {
            checkPin();
        } else if (port == InPorts::pin && in.isInteger()) {
            pin = in.asInteger();
            checkPin();
        } else if (port == InPorts::threshold && in.isInteger()) {
            touchedThreshold = in.asInteger();
            checkPin();
        }
    }
private:
    void checkPin() {
        currentState = (readCapacitivePin(pin) > touchedThreshold);
        send(Packet(currentState));
    }

private:
    bool currentState;
    int pin;
    int touchedThreshold;
};

#else
class ReadCapacitivePin : public DummyComponent {};
#endif
<commit_msg>ReadCapacitivePin: Check for AVR, return 255 if not<commit_after>/* microflo_component yaml
name: ReadCapacitivePin
description: "Emits true if measured capacitance (in iterations) on @pin exceeds @threshold"
inports:
  trigger:
    type: all
    description: ""
  pin:
    type: all
    description: ""
  threshold:
    type: all
    description: ""
outports:
  out:
    type: all
    description: ""
microflo_component */

#ifdef __AVR__
// readCapacitivePin
//  Input: Arduino pin number
//  Output: A number, from 0 to 17 expressing
//          how much capacitance is on the pin
//  When you touch the pin, or whatever you have
//  attached to it, the number will get higher
//  In order for this to work now,
// The pin should have a 1+Megaohm resistor pulling
//  it up to +5v.
uint8_t readCapacitivePin(int pinToMeasure) {
  // This is how you declare a variable which
  //  will hold the PORT, PIN, and DDR registers
  //  on an AVR
  volatile uint8_t* port;
  volatile uint8_t* ddr;
  volatile uint8_t* pin;
  // Here we translate the input pin number from
  //  Arduino pin number to the AVR PORT, PIN, DDR,
  //  and which bit of those registers we care about.
  byte bitmask;
  if ((pinToMeasure >= 0) && (pinToMeasure <= 7)){
    port = &PORTD;
    ddr = &DDRD;
    bitmask = 1 << pinToMeasure;
    pin = &PIND;
  }
  if ((pinToMeasure > 7) && (pinToMeasure <= 13)){
    port = &PORTB;
    ddr = &DDRB;
    bitmask = 1 << (pinToMeasure - 8);
    pin = &PINB;
  }
  if ((pinToMeasure > 13) && (pinToMeasure <= 19)){
    port = &PORTC;
    ddr = &DDRC;
    bitmask = 1 << (pinToMeasure - 13);
    pin = &PINC;
  }
  // Discharge the pin first by setting it low and output
  *port &= ~(bitmask);
  *ddr  |= bitmask;
  delay(1);
  // Make the pin an input WITHOUT the internal pull-up on
  *ddr &= ~(bitmask);
  // Now see how long the pin to get pulled up
  int cycles = 16000;
  for(int i = 0; i < cycles; i++){
    if (*pin & bitmask){
      cycles = i;
      break;
    }
  }
  // Discharge the pin again by setting it low and output
  //  It's important to leave the pins low if you want to
  //  be able to touch more than 1 sensor at a time - if
  //  the sensor is left pulled high, when you touch
  //  two sensors, your body will transfer the charge between
  //  sensors.
  *port &= ~(bitmask);
  *ddr  |= bitmask;

  return cycles;
}

#else
#warning "ReadCapacitivePin support compiled out"
uint8_t readCapacitivePin(int pinToMeasure) {
    return 255;
}
#endif

class ReadCapacitivePin : public SingleOutputComponent {
public:
    virtual void process(Packet in, MicroFlo::PortId port) {
        using namespace ReadCapacitivePinPorts;

        if (in.isSetup()) {
            currentState = false;
            pin = 1;
            touchedThreshold = 60;
        } else if (port == InPorts::trigger) {
            checkPin();
        } else if (port == InPorts::pin && in.isInteger()) {
            pin = in.asInteger();
            checkPin();
        } else if (port == InPorts::threshold && in.isInteger()) {
            touchedThreshold = in.asInteger();
            checkPin();
        }
    }
private:
    void checkPin() {
        currentState = (readCapacitivePin(pin) > touchedThreshold);
        send(Packet(currentState));
    }

private:
    bool currentState;
    int pin;
    int touchedThreshold;
};
<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 *
 */

#include <folly/portability/GTest.h>
#include <quic/common/test/TestUtils.h>
#include <quic/congestion_control/QuicCubic.h>

using namespace testing;

namespace quic {
namespace test {

class CubicHystartTest : public Test {};

TEST_F(CubicHystartTest, SendAndAck) {
  QuicConnectionStateBase conn(QuicNodeType::Client);
  conn.udpSendPacketLen = 100;
  Cubic cubic(conn);
  auto initCwnd = cubic.getWritableBytes();
  // Packet 0 is sent:
  conn.lossState.largestSent = 0;

  // Packet 0 is acked:
  conn.lossState.srtt = 100us;
  auto packet = makeTestingWritePacket(0, 1000, 1000);
  cubic.onPacketSent(packet);
  cubic.onPacketAckOrLoss(
      makeAck(0, 1000, Clock::now(), packet.time), folly::none);

  EXPECT_EQ(initCwnd + 1000, cubic.getWritableBytes());
  EXPECT_EQ(CubicStates::Hystart, cubic.state());
}

TEST_F(CubicHystartTest, CwndLargerThanSSThresh) {
  QuicConnectionStateBase conn(QuicNodeType::Client);
  Cubic cubic(conn, 0);
  auto initCwnd = cubic.getWritableBytes();
  // Packet 0 is sent:
  conn.lossState.largestSent = 0;

  // Packet 0 is acked:
  auto packet = makeTestingWritePacket(0, 1000, 1000);
  cubic.onPacketSent(packet);
  cubic.onPacketAckOrLoss(
      makeAck(0, 1000, Clock::now(), packet.time), folly::none);
  EXPECT_EQ(initCwnd + 1000, cubic.getWritableBytes());
  EXPECT_EQ(CubicStates::Steady, cubic.state());
}

TEST_F(CubicHystartTest, NoDelayIncrease) {
  QuicConnectionStateBase conn(QuicNodeType::Client);
  conn.udpSendPacketLen = 100;
  Cubic cubic(conn);
  auto initCwnd = cubic.getWritableBytes();
  // Packet 0 is sent:
  conn.lossState.largestSent = 0;

  // Packet 0 is acked:
  conn.lossState.srtt = 2us;
  auto realNow = quic::Clock::now();
  // One onPacketAcked will not trigger DelayIncrease
  auto packet = makeTestingWritePacket(0, 1000, 1000);
  cubic.onPacketSent(packet);
  cubic.onPacketAckOrLoss(
      makeAck(0, 1000, realNow + 2us, packet.time), folly::none);
  EXPECT_EQ(initCwnd + 1000, cubic.getWritableBytes());
  EXPECT_EQ(CubicStates::Hystart, cubic.state());
}

TEST_F(CubicHystartTest, AckTrain) {
  QuicConnectionStateBase conn(QuicNodeType::Client);
  Cubic cubic(conn, std::numeric_limits<uint64_t>::max(), true, true);
  auto initCwnd = cubic.getWritableBytes();
  // srtt will be assigned to delayMin:
  conn.lossState.srtt = 2us;
  auto realNow = quic::Clock::now();
  conn.lossState.largestSent = 0;
  // Packet 0 is sent:
  auto packetSize = kLowSsthreshInMss * conn.udpSendPacketLen;
  auto packet0 = makeTestingWritePacket(
      0, packetSize, packetSize, realNow - std::chrono::milliseconds(10));
  cubic.onPacketSent(packet0);
  // Packet 1 is sent:
  auto packet1 = makeTestingWritePacket(
      1, packetSize, packetSize * 2, realNow - std::chrono::milliseconds(5));
  cubic.onPacketSent(packet1);
  // Packet 0 is acked:
  cubic.onPacketAckOrLoss(
      makeAck(
          0, kLowSsthreshInMss * conn.udpSendPacketLen, realNow, packet0.time),
      folly::none);
  // Packet 1 is acked:
  cubic.onPacketAckOrLoss(
      makeAck(
          1,
          kLowSsthreshInMss * conn.udpSendPacketLen,
          realNow + 2us,
          packet1.time),
      folly::none);
  EXPECT_EQ(
      initCwnd + kLowSsthreshInMss * conn.udpSendPacketLen * 2,
      cubic.getWritableBytes());
  EXPECT_EQ(CubicStates::Steady, cubic.state());
}

TEST_F(CubicHystartTest, NoAckTrainNoDelayIncrease) {
  QuicConnectionStateBase conn(QuicNodeType::Client);
  Cubic cubic(conn, std::numeric_limits<uint64_t>::max(), true, true);
  auto initCwnd = cubic.getWritableBytes();
  // Packet 0 is sent:
  conn.lossState.largestSent = 0;

  // Packet 0 is acked:
  // make sure AckTrain won't find it
  conn.lossState.srtt = 2us;
  auto realNow = quic::Clock::now();
  // One onPacketAcked will not trigger DelayIncrease
  auto packet = makeTestingWritePacket(0, 1000, 1000, realNow);
  cubic.onPacketSent(packet);
  cubic.onPacketAckOrLoss(
      makeAck(0, 1000, realNow + kAckCountingGap + 2us, packet.time),
      folly::none);
  EXPECT_EQ(initCwnd + 1000, cubic.getWritableBytes());
  EXPECT_EQ(CubicStates::Hystart, cubic.state());
}

TEST_F(CubicHystartTest, DelayIncrease) {
  QuicConnectionStateBase conn(QuicNodeType::Client);
  Cubic cubic(conn);
  PacketNum packetNum = 0;
  // get kLowSsthresh out of the way for exit condition
  uint64_t totalSent = 0;
  while (cubic.getWritableBytes() <=
         kLowSsthreshInMss * conn.udpSendPacketLen) {
    auto fullSize = cubic.getWritableBytes();
    auto packet =
        makeTestingWritePacket(packetNum, fullSize, fullSize + totalSent);
    cubic.onPacketSent(packet);
    cubic.onPacketAckOrLoss(
        makeAck(packetNum++, fullSize, Clock::now(), packet.time), folly::none);
    totalSent += fullSize;
  }

  auto packetsSentTime = Clock::now();
  auto firstPacketNum = packetNum++;
  auto packet0 = makeTestingWritePacket(
      packetNum, 1000, totalSent + 1000, packetsSentTime);
  totalSent += 1000;
  conn.lossState.largestSent = firstPacketNum;
  cubic.onPacketSent(packet0);

  auto secondPacketNum = packetNum++;
  auto packet1 = makeTestingWritePacket(
      secondPacketNum, 1000, totalSent + 1000, packetsSentTime);
  totalSent += 1000;
  // Packet 1 is sent:
  conn.lossState.largestSent = secondPacketNum;
  cubic.onPacketSent(packet1);

  conn.lossState.srtt = 100us;

  // First onPacketAcked will set up lastSampledRtt = 200us for next round. It
  // will also set up rttRoundEndTarget = 1 since largestSent = 1.
  auto ackTime = packetsSentTime + 1us;
  auto ackTimeIncrease = 2us;
  ackTime += ackTimeIncrease;
  cubic.onPacketAckOrLoss(
      makeAck(firstPacketNum, 1000, ackTime, packet0.time), folly::none);
  ackTime += ackTimeIncrease;
  cubic.onPacketAckOrLoss(
      makeAck(secondPacketNum, 1000, ackTime, packet1.time), folly::none);
  auto estimatedRttEndTarget = Clock::now();

  auto packet2 = makeTestingWritePacket(
      packetNum, 1000, totalSent + 1000, estimatedRttEndTarget + 1us);
  totalSent += 1000;
  conn.lossState.largestSent = packetNum;
  cubic.onPacketSent(packet2);
  // This will end current RTT round and start a new one next time Ack happens:
  ackTime = packet2.time + ackTimeIncrease;
  cubic.onPacketAckOrLoss(
      makeAck(packetNum, 1000, ackTime, packet2.time), folly::none);
  packetNum++;
  auto cwndEndRound = cubic.getWritableBytes();

  // New RTT round, give currSampledRtt a value larger than previous RTT:
  conn.lossState.srtt = 200us;
  // onPacketAcked kAckSampling - 1 times:
  std::vector<CongestionController::AckEvent> moreAcks;
  for (size_t i = 0; i < kAckSampling - 1; i++) {
    conn.lossState.largestSent = i + packetNum;
    auto packet = makeTestingWritePacket(
        i + packetNum, 1000, totalSent + 1000, packetsSentTime);
    cubic.onPacketSent(packet);
    totalSent += 1000;
    ackTime += ackTimeIncrease;
    moreAcks.push_back(makeAck(1 + packetNum, 1000, ackTime, packet.time));
    packetNum++;
  }
  for (auto& ack : moreAcks) {
    cubic.onPacketAckOrLoss(ack, folly::none);
  }
  // kAckSampling-th onPacketAcked in this round. This will trigger
  // DelayIncrease:
  conn.lossState.largestSent = packetNum;
  auto packetEnd = makeTestingWritePacket(
      packetNum, 1000, totalSent + 1000, packetsSentTime);
  cubic.onPacketSent(packetEnd);
  totalSent += 1000;
  ackTime += ackTimeIncrease;
  cubic.onPacketAckOrLoss(
      makeAck(packetNum, 1000, ackTime, packetEnd.time), folly::none);

  EXPECT_EQ(cwndEndRound + 1000 * kAckSampling, cubic.getWritableBytes());
  EXPECT_EQ(CubicStates::Steady, cubic.state());
}

TEST_F(CubicHystartTest, DelayIncreaseCwndTooSmall) {
  QuicConnectionStateBase conn(QuicNodeType::Client);
  Cubic cubic(conn);
  auto realNow = quic::Clock::now();

  uint64_t totalSent = 0;
  auto packet0 = makeTestingWritePacket(0, 1, totalSent + 1, realNow);
  // Packet 0 is sent:
  conn.lossState.largestSent = 0;
  cubic.onPacketSent(packet0);
  totalSent += 1;

  auto packet1 = makeTestingWritePacket(1, 1, totalSent + 1, realNow);
  conn.lossState.largestSent = 1;
  cubic.onPacketSent(packet1);
  totalSent += 1;

  conn.lossState.srtt = 100us;

  // First onPacketAcked will set up lastSampledRtt = 200us for next round:
  auto ackTime = realNow;
  auto ackTimeIncrease = 2us;
  ackTime += ackTimeIncrease;
  cubic.onPacketAckOrLoss(makeAck(0, 1, ackTime, packet0.time), folly::none);
  ackTime += ackTimeIncrease;
  cubic.onPacketAckOrLoss(makeAck(1, 1, ackTime, packet1.time), folly::none);

  auto packet2 = makeTestingWritePacket(2, 10, 10 + totalSent, realNow);
  conn.lossState.largestSent = 2;
  cubic.onPacketSent(packet2);
  totalSent += 10;
  ackTime += ackTimeIncrease;
  cubic.onPacketAckOrLoss(makeAck(2, 1, ackTime, packet2.time), folly::none);
  auto cwndEndRound = cubic.getWritableBytes();

  // New RTT round, give currSampledRtt a value larger than previous RTT:
  conn.lossState.srtt = 200us;
  // onPacketAcked kAckSampling - 1 times:
  std::vector<CongestionController::AckEvent> moreAcks;
  for (size_t i = 0; i < kAckSampling - 1; i++) {
    conn.lossState.largestSent = i + 3;
    auto packet = makeTestingWritePacket(i + 3, 1, 1 + totalSent, realNow);
    ackTime += ackTimeIncrease;
    moreAcks.push_back(makeAck(i + 3, 1, ackTime, packet.time));
    cubic.onPacketSent(packet);
    totalSent += 1;
  }
  for (auto& ack : moreAcks) {
    cubic.onPacketAckOrLoss(ack, folly::none);
  }
  // kAckSampling-th onPacketAcked in this round. This will trigger
  // DelayIncrease:
  conn.lossState.largestSent = kAckSampling;
  auto packetEnd =
      makeTestingWritePacket(kAckSampling, 1, 1 + totalSent, realNow);
  cubic.onPacketSent(packetEnd);
  totalSent += 1;
  ackTime += ackTimeIncrease;
  cubic.onPacketAckOrLoss(
      makeAck(kAckSampling, 1, ackTime, packetEnd.time), folly::none);
  auto expectedCwnd = cwndEndRound + 1 * kAckSampling;
  // Cwnd < kLowSsthresh, won't exit Hystart state:
  ASSERT_LT(expectedCwnd, kLowSsthreshInMss * conn.udpSendPacketLen);
  EXPECT_EQ(expectedCwnd, cubic.getWritableBytes());
  EXPECT_EQ(CubicStates::Hystart, cubic.state());
}

TEST_F(CubicHystartTest, ReduceByCubicReductionFactor) {
  QuicConnectionStateBase conn(QuicNodeType::Client);
  Cubic cubic(conn, std::numeric_limits<uint64_t>::max());
  auto initCwnd = cubic.getWritableBytes();
  conn.lossState.largestSent = 0;
  auto packet = makeTestingWritePacket(0, 1000, 1000);
  // this increases inflight by 1000:
  cubic.onPacketSent(packet);
  EXPECT_EQ(initCwnd - 1000, cubic.getWritableBytes());
  // this decreases inflight by 1000, and then decreases cwnd by Cubic:
  CongestionController::LossEvent loss;
  loss.addLostPacket(packet);
  cubic.onPacketAckOrLoss(folly::none, std::move(loss));
  EXPECT_EQ(initCwnd * kDefaultCubicReductionFactor, cubic.getWritableBytes());
  EXPECT_EQ(CubicStates::FastRecovery, cubic.state());
}
} // namespace test
} // namespace quic
<commit_msg>Fix Quic cubic unit test flakiness<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 *
 */

#include <folly/portability/GTest.h>
#include <quic/common/test/TestUtils.h>
#include <quic/congestion_control/QuicCubic.h>

using namespace testing;

namespace quic {
namespace test {

class CubicHystartTest : public Test {};

TEST_F(CubicHystartTest, SendAndAck) {
  QuicConnectionStateBase conn(QuicNodeType::Client);
  conn.udpSendPacketLen = 100;
  Cubic cubic(conn);
  auto initCwnd = cubic.getWritableBytes();
  // Packet 0 is sent:
  conn.lossState.largestSent = 0;

  // Packet 0 is acked:
  conn.lossState.srtt = 100us;
  auto packet = makeTestingWritePacket(0, 1000, 1000);
  cubic.onPacketSent(packet);
  cubic.onPacketAckOrLoss(
      makeAck(0, 1000, Clock::now(), packet.time), folly::none);

  EXPECT_EQ(initCwnd + 1000, cubic.getWritableBytes());
  EXPECT_EQ(CubicStates::Hystart, cubic.state());
}

TEST_F(CubicHystartTest, CwndLargerThanSSThresh) {
  QuicConnectionStateBase conn(QuicNodeType::Client);
  Cubic cubic(conn, 0);
  auto initCwnd = cubic.getWritableBytes();
  // Packet 0 is sent:
  conn.lossState.largestSent = 0;

  // Packet 0 is acked:
  auto packet = makeTestingWritePacket(0, 1000, 1000);
  cubic.onPacketSent(packet);
  cubic.onPacketAckOrLoss(
      makeAck(0, 1000, Clock::now(), packet.time), folly::none);
  EXPECT_EQ(initCwnd + 1000, cubic.getWritableBytes());
  EXPECT_EQ(CubicStates::Steady, cubic.state());
}

TEST_F(CubicHystartTest, NoDelayIncrease) {
  QuicConnectionStateBase conn(QuicNodeType::Client);
  conn.udpSendPacketLen = 100;
  Cubic cubic(conn);
  auto initCwnd = cubic.getWritableBytes();
  // Packet 0 is sent:
  conn.lossState.largestSent = 0;

  // Packet 0 is acked:
  conn.lossState.srtt = 2us;
  auto realNow = quic::Clock::now();
  // One onPacketAcked will not trigger DelayIncrease
  auto packet = makeTestingWritePacket(0, 1000, 1000, realNow);
  cubic.onPacketSent(packet);
  cubic.onPacketAckOrLoss(
      makeAck(0, 1000, realNow + 2us, packet.time), folly::none);
  EXPECT_EQ(initCwnd + 1000, cubic.getWritableBytes());
  EXPECT_EQ(CubicStates::Hystart, cubic.state());
}

TEST_F(CubicHystartTest, AckTrain) {
  QuicConnectionStateBase conn(QuicNodeType::Client);
  Cubic cubic(conn, std::numeric_limits<uint64_t>::max(), true, true);
  auto initCwnd = cubic.getWritableBytes();
  // srtt will be assigned to delayMin:
  conn.lossState.srtt = 2us;
  auto realNow = quic::Clock::now();
  conn.lossState.largestSent = 0;
  // Packet 0 is sent:
  auto packetSize = kLowSsthreshInMss * conn.udpSendPacketLen;
  auto packet0 = makeTestingWritePacket(
      0, packetSize, packetSize, realNow - std::chrono::milliseconds(10));
  cubic.onPacketSent(packet0);
  // Packet 1 is sent:
  auto packet1 = makeTestingWritePacket(
      1, packetSize, packetSize * 2, realNow - std::chrono::milliseconds(5));
  cubic.onPacketSent(packet1);
  // Packet 0 is acked:
  cubic.onPacketAckOrLoss(
      makeAck(
          0, kLowSsthreshInMss * conn.udpSendPacketLen, realNow, packet0.time),
      folly::none);
  // Packet 1 is acked:
  cubic.onPacketAckOrLoss(
      makeAck(
          1,
          kLowSsthreshInMss * conn.udpSendPacketLen,
          realNow + 2us,
          packet1.time),
      folly::none);
  EXPECT_EQ(
      initCwnd + kLowSsthreshInMss * conn.udpSendPacketLen * 2,
      cubic.getWritableBytes());
  EXPECT_EQ(CubicStates::Steady, cubic.state());
}

TEST_F(CubicHystartTest, NoAckTrainNoDelayIncrease) {
  QuicConnectionStateBase conn(QuicNodeType::Client);
  Cubic cubic(conn, std::numeric_limits<uint64_t>::max(), true, true);
  auto initCwnd = cubic.getWritableBytes();
  // Packet 0 is sent:
  conn.lossState.largestSent = 0;

  // Packet 0 is acked:
  // make sure AckTrain won't find it
  conn.lossState.srtt = 2us;
  auto realNow = quic::Clock::now();
  // One onPacketAcked will not trigger DelayIncrease
  auto packet = makeTestingWritePacket(0, 1000, 1000, realNow);
  cubic.onPacketSent(packet);
  cubic.onPacketAckOrLoss(
      makeAck(0, 1000, realNow + kAckCountingGap + 2us, packet.time),
      folly::none);
  EXPECT_EQ(initCwnd + 1000, cubic.getWritableBytes());
  EXPECT_EQ(CubicStates::Hystart, cubic.state());
}

TEST_F(CubicHystartTest, DelayIncrease) {
  QuicConnectionStateBase conn(QuicNodeType::Client);
  Cubic cubic(conn);
  PacketNum packetNum = 0;
  // get kLowSsthresh out of the way for exit condition
  uint64_t totalSent = 0;
  while (cubic.getWritableBytes() <=
         kLowSsthreshInMss * conn.udpSendPacketLen) {
    auto fullSize = cubic.getWritableBytes();
    auto packet =
        makeTestingWritePacket(packetNum, fullSize, fullSize + totalSent);
    cubic.onPacketSent(packet);
    cubic.onPacketAckOrLoss(
        makeAck(packetNum++, fullSize, Clock::now(), packet.time), folly::none);
    totalSent += fullSize;
  }

  auto packetsSentTime = Clock::now();
  auto firstPacketNum = packetNum++;
  auto packet0 = makeTestingWritePacket(
      packetNum, 1000, totalSent + 1000, packetsSentTime);
  totalSent += 1000;
  conn.lossState.largestSent = firstPacketNum;
  cubic.onPacketSent(packet0);

  auto secondPacketNum = packetNum++;
  auto packet1 = makeTestingWritePacket(
      secondPacketNum, 1000, totalSent + 1000, packetsSentTime);
  totalSent += 1000;
  // Packet 1 is sent:
  conn.lossState.largestSent = secondPacketNum;
  cubic.onPacketSent(packet1);

  conn.lossState.srtt = 100us;

  // First onPacketAcked will set up lastSampledRtt = 200us for next round. It
  // will also set up rttRoundEndTarget = 1 since largestSent = 1.
  auto ackTime = packetsSentTime + 1us;
  auto ackTimeIncrease = 2us;
  ackTime += ackTimeIncrease;
  cubic.onPacketAckOrLoss(
      makeAck(firstPacketNum, 1000, ackTime, packet0.time), folly::none);
  ackTime += ackTimeIncrease;
  cubic.onPacketAckOrLoss(
      makeAck(secondPacketNum, 1000, ackTime, packet1.time), folly::none);
  auto estimatedRttEndTarget = Clock::now();

  auto packet2 = makeTestingWritePacket(
      packetNum, 1000, totalSent + 1000, estimatedRttEndTarget + 1us);
  totalSent += 1000;
  conn.lossState.largestSent = packetNum;
  cubic.onPacketSent(packet2);
  // This will end current RTT round and start a new one next time Ack happens:
  ackTime = packet2.time + ackTimeIncrease;
  cubic.onPacketAckOrLoss(
      makeAck(packetNum, 1000, ackTime, packet2.time), folly::none);
  packetNum++;
  auto cwndEndRound = cubic.getWritableBytes();

  // New RTT round, give currSampledRtt a value larger than previous RTT:
  conn.lossState.srtt = 200us;
  // onPacketAcked kAckSampling - 1 times:
  std::vector<CongestionController::AckEvent> moreAcks;
  for (size_t i = 0; i < kAckSampling - 1; i++) {
    conn.lossState.largestSent = i + packetNum;
    auto packet = makeTestingWritePacket(
        i + packetNum, 1000, totalSent + 1000, packetsSentTime);
    cubic.onPacketSent(packet);
    totalSent += 1000;
    ackTime += ackTimeIncrease;
    moreAcks.push_back(makeAck(1 + packetNum, 1000, ackTime, packet.time));
    packetNum++;
  }
  for (auto& ack : moreAcks) {
    cubic.onPacketAckOrLoss(ack, folly::none);
  }
  // kAckSampling-th onPacketAcked in this round. This will trigger
  // DelayIncrease:
  conn.lossState.largestSent = packetNum;
  auto packetEnd = makeTestingWritePacket(
      packetNum, 1000, totalSent + 1000, packetsSentTime);
  cubic.onPacketSent(packetEnd);
  totalSent += 1000;
  ackTime += ackTimeIncrease;
  cubic.onPacketAckOrLoss(
      makeAck(packetNum, 1000, ackTime, packetEnd.time), folly::none);

  EXPECT_EQ(cwndEndRound + 1000 * kAckSampling, cubic.getWritableBytes());
  EXPECT_EQ(CubicStates::Steady, cubic.state());
}

TEST_F(CubicHystartTest, DelayIncreaseCwndTooSmall) {
  QuicConnectionStateBase conn(QuicNodeType::Client);
  Cubic cubic(conn);
  auto realNow = quic::Clock::now();

  uint64_t totalSent = 0;
  auto packet0 = makeTestingWritePacket(0, 1, totalSent + 1, realNow);
  // Packet 0 is sent:
  conn.lossState.largestSent = 0;
  cubic.onPacketSent(packet0);
  totalSent += 1;

  auto packet1 = makeTestingWritePacket(1, 1, totalSent + 1, realNow);
  conn.lossState.largestSent = 1;
  cubic.onPacketSent(packet1);
  totalSent += 1;

  conn.lossState.srtt = 100us;

  // First onPacketAcked will set up lastSampledRtt = 200us for next round:
  auto ackTime = realNow;
  auto ackTimeIncrease = 2us;
  ackTime += ackTimeIncrease;
  cubic.onPacketAckOrLoss(makeAck(0, 1, ackTime, packet0.time), folly::none);
  ackTime += ackTimeIncrease;
  cubic.onPacketAckOrLoss(makeAck(1, 1, ackTime, packet1.time), folly::none);

  auto packet2 = makeTestingWritePacket(2, 10, 10 + totalSent, realNow);
  conn.lossState.largestSent = 2;
  cubic.onPacketSent(packet2);
  totalSent += 10;
  ackTime += ackTimeIncrease;
  cubic.onPacketAckOrLoss(makeAck(2, 1, ackTime, packet2.time), folly::none);
  auto cwndEndRound = cubic.getWritableBytes();

  // New RTT round, give currSampledRtt a value larger than previous RTT:
  conn.lossState.srtt = 200us;
  // onPacketAcked kAckSampling - 1 times:
  std::vector<CongestionController::AckEvent> moreAcks;
  for (size_t i = 0; i < kAckSampling - 1; i++) {
    conn.lossState.largestSent = i + 3;
    auto packet = makeTestingWritePacket(i + 3, 1, 1 + totalSent, realNow);
    ackTime += ackTimeIncrease;
    moreAcks.push_back(makeAck(i + 3, 1, ackTime, packet.time));
    cubic.onPacketSent(packet);
    totalSent += 1;
  }
  for (auto& ack : moreAcks) {
    cubic.onPacketAckOrLoss(ack, folly::none);
  }
  // kAckSampling-th onPacketAcked in this round. This will trigger
  // DelayIncrease:
  conn.lossState.largestSent = kAckSampling;
  auto packetEnd =
      makeTestingWritePacket(kAckSampling, 1, 1 + totalSent, realNow);
  cubic.onPacketSent(packetEnd);
  totalSent += 1;
  ackTime += ackTimeIncrease;
  cubic.onPacketAckOrLoss(
      makeAck(kAckSampling, 1, ackTime, packetEnd.time), folly::none);
  auto expectedCwnd = cwndEndRound + 1 * kAckSampling;
  // Cwnd < kLowSsthresh, won't exit Hystart state:
  ASSERT_LT(expectedCwnd, kLowSsthreshInMss * conn.udpSendPacketLen);
  EXPECT_EQ(expectedCwnd, cubic.getWritableBytes());
  EXPECT_EQ(CubicStates::Hystart, cubic.state());
}

TEST_F(CubicHystartTest, ReduceByCubicReductionFactor) {
  QuicConnectionStateBase conn(QuicNodeType::Client);
  Cubic cubic(conn, std::numeric_limits<uint64_t>::max());
  auto initCwnd = cubic.getWritableBytes();
  conn.lossState.largestSent = 0;
  auto packet = makeTestingWritePacket(0, 1000, 1000);
  // this increases inflight by 1000:
  cubic.onPacketSent(packet);
  EXPECT_EQ(initCwnd - 1000, cubic.getWritableBytes());
  // this decreases inflight by 1000, and then decreases cwnd by Cubic:
  CongestionController::LossEvent loss;
  loss.addLostPacket(packet);
  cubic.onPacketAckOrLoss(folly::none, std::move(loss));
  EXPECT_EQ(initCwnd * kDefaultCubicReductionFactor, cubic.getWritableBytes());
  EXPECT_EQ(CubicStates::FastRecovery, cubic.state());
}
} // namespace test
} // namespace quic
<|endoftext|>


<commit_before>// =============================================================================
// PROJECT CHRONO - http://projectchrono.org
//
// Copyright (c) 2014 projectchrono.org
// All right reserved.
//
// Use of this source code is governed by a BSD-style license that can be found
// in the LICENSE file at the top level of the distribution and at
// http://projectchrono.org/license-chrono.txt.
//
// =============================================================================
// Authors: Radu Serban, Justin Madsen
// =============================================================================
//
// Main driver function for the HMMWV 9-body model, using rigid tire-terrain
// contact.
//
// If using the Irrlicht interface, river inputs are obtained from the keyboard.
//
// The global reference frame has Z up, X towards the back of the vehicle, and
// Y pointing to the right.
//
// =============================================================================

#include "core/ChFileutils.h"
#include "core/ChStream.h"
#include "physics/ChApidll.h"
#include "physics/ChSystem.h"
#include "physics/ChLinkDistance.h"

#include "utils/ChUtilsInputOutput.h"

#include "HMMWV_9body_config.h"
#include "HMMWV9_Vehicle.h"
#include "HMMWV9_RigidTerrain.h"

// If Irrlicht support is available...
#if IRRLICHT_ENABLED
  // ...include additional headers
# include "unit_IRRLICHT/ChIrrApp.h"
# include <irrlicht.h>
# include "HMMWV9_IrrGuiDriver.h"
# include "HMMWV9_IrrRenderer.h"
  // ...and specify whether the demo should actually use Irrlicht
# define USE_IRRLICHT
#endif

using namespace chrono;

// =============================================================================

// Initial vehicle position
ChVector<>     initLoc(0, 0, 1.7);    // sprung mass height at design = 49.68 in
ChQuaternion<> initRot(1,0,0,0);      // forward is the positive x-direction

// Rigid terrain dimensions
double terrainHeight = 0;
double terrainLength = 100.0;  // size in X direction
double terrainWidth  = 100.0;   // size in Y directoin

// Camera offset relative to vehicle
ChVector<> cameraOffset(2.5, 0, 2.0);

// Simulation parameters
double tend = 10.0;
double step_size = 0.001;
int out_fps = 30;

// Output directories
const std::string out_dir = "../HMMWV9";
const std::string pov_dir = out_dir + "/POVRAY";

// =============================================================================

int main(int argc, char* argv[])
{
  DLL_CreateGlobals();

  // -----------------
  // Create the system
  // -----------------

  ChSystem m_system;

  m_system.Set_G_acc(ChVector<>(0, 0, -9.81));

  // Integration and Solver settings
  m_system.SetLcpSolverType(ChSystem::LCP_ITERATIVE_SOR);
  m_system.SetIterLCPmaxItersSpeed(150);
  m_system.SetIterLCPmaxItersStab(150);
  m_system.SetMaxPenetrationRecoverySpeed(4.0);
  m_system.SetStep(step_size);

  // Create the HMMWV vehicle
  HMMWV9_Vehicle vehicle(m_system,
                         ChCoordsys<>(initLoc, initRot),
                         false,
                         HMMWV9_Vehicle::NONE,
                         HMMWV9_Vehicle::PRIMITIVES);

  // Create the ground
  HMMWV9_RigidTerrain terrain(m_system, terrainHeight, terrainLength, terrainWidth, 0.8);

#ifdef USE_IRRLICHT
  irr::ChIrrApp application(&m_system,
                            L"HMMWV 9-body demo",
                            irr::core::dimension2d<irr::u32>(1000, 800),
                            false,
                            true);

  application.AddTypicalSky();
  application.AddTypicalLights();

  application.SetTimestep(step_size);

  HMMWV9_IrrGuiDriver driver(application, vehicle);

  driver.CreateCamera(cameraOffset);

  // Set up the assets for rendering
  application.AssetBindAll();
  application.AssetUpdateAll();
#endif

  // ---------------
  // Simulation loop
  // ---------------

  int out_steps = std::ceil((1 / step_size) / out_fps);

  double time = 0;
  int frame = 0;
  int out_frame = 0;

#ifdef USE_IRRLICHT
  while (application.GetDevice()->run())
  {
    bool render = (frame % out_steps == 0);

    // Render scene
    if (render) {
      application.GetVideoDriver()->beginScene(true, true, irr::video::SColor(255, 140, 161, 192));

      application.DrawAll();
      driver.DrawAll();

      out_frame++;
    }

    // Update subsystems and advance simulation by one time step
    driver.Update(time);
    vehicle.Update(time, driver.getThrottle(), driver.getSteering());

    m_system.DoStepDynamics(step_size);
    time += step_size;
    frame++;

    // Complete scene
    if (render)
      application.GetVideoDriver()->endScene();
  }
#else
  if(ChFileutils::MakeDirectory(out_dir.c_str()) < 0) {
    std::cout << "Error creating directory " << out_dir << std::endl;
    return 1;
  }
  if(ChFileutils::MakeDirectory(pov_dir.c_str()) < 0) {
    std::cout << "Error creating directory " << pov_dir << std::endl;
    return 1;
  }

  utils::WriteMeshPovray(HMMWV9_Vehicle::ChassisMeshFile(), HMMWV9_Vehicle::ChassisMeshName(), out_dir);
  utils::WriteMeshPovray(HMMWV9_Wheel::MeshFile(), HMMWV9_Wheel::MeshName(), out_dir);

  char filename[100];

  while (time < tend)
  {
    if (frame % out_steps == 0) {
      // Output render data
      sprintf(filename, "%s/data_%03d.dat", pov_dir.c_str(), out_frame + 1);
      utils::WriteShapesPovray(&m_system, filename);
      std::cout << "------------ Output frame:   " << out_frame << std::endl;
      std::cout << "             Sim frame:      " << frame << std::endl;
      std::cout << "             Time:           " << time << std::endl;
      out_frame++;

      break;

    }

    // Update subsystems and advance simulation by one time step
    ////driver.Update(time);
    ////vehicle.Update(time, driver.getThrottle(), driver.getSteering());

    m_system.DoStepDynamics(step_size);
    time += step_size;
    frame++;
  }
#endif

#ifdef USE_IRRLICHT
  application.GetDevice()->drop();
#endif

  DLL_DeleteGlobals();

  return 0;
}
<commit_msg>Remove obsolete header.<commit_after>// =============================================================================
// PROJECT CHRONO - http://projectchrono.org
//
// Copyright (c) 2014 projectchrono.org
// All right reserved.
//
// Use of this source code is governed by a BSD-style license that can be found
// in the LICENSE file at the top level of the distribution and at
// http://projectchrono.org/license-chrono.txt.
//
// =============================================================================
// Authors: Radu Serban, Justin Madsen
// =============================================================================
//
// Main driver function for the HMMWV 9-body model, using rigid tire-terrain
// contact.
//
// If using the Irrlicht interface, river inputs are obtained from the keyboard.
//
// The global reference frame has Z up, X towards the back of the vehicle, and
// Y pointing to the right.
//
// =============================================================================

#include "core/ChFileutils.h"
#include "core/ChStream.h"
#include "physics/ChApidll.h"
#include "physics/ChSystem.h"
#include "physics/ChLinkDistance.h"

#include "utils/ChUtilsInputOutput.h"

#include "HMMWV_9body_config.h"
#include "HMMWV9_Vehicle.h"
#include "HMMWV9_RigidTerrain.h"

// If Irrlicht support is available...
#if IRRLICHT_ENABLED
  // ...include additional headers
# include "unit_IRRLICHT/ChIrrApp.h"
# include <irrlicht.h>
# include "HMMWV9_IrrGuiDriver.h"
  // ...and specify whether the demo should actually use Irrlicht
# define USE_IRRLICHT
#endif

using namespace chrono;

// =============================================================================

// Initial vehicle position
ChVector<>     initLoc(0, 0, 1.7);    // sprung mass height at design = 49.68 in
ChQuaternion<> initRot(1,0,0,0);      // forward is the positive x-direction

// Rigid terrain dimensions
double terrainHeight = 0;
double terrainLength = 100.0;  // size in X direction
double terrainWidth  = 100.0;   // size in Y directoin

// Camera offset relative to vehicle
ChVector<> cameraOffset(2.5, 0, 2.0);

// Simulation parameters
double tend = 10.0;
double step_size = 0.001;
int out_fps = 30;

// Output directories
const std::string out_dir = "../HMMWV9";
const std::string pov_dir = out_dir + "/POVRAY";

// =============================================================================

int main(int argc, char* argv[])
{
  DLL_CreateGlobals();

  // -----------------
  // Create the system
  // -----------------

  ChSystem m_system;

  m_system.Set_G_acc(ChVector<>(0, 0, -9.81));

  // Integration and Solver settings
  m_system.SetLcpSolverType(ChSystem::LCP_ITERATIVE_SOR);
  m_system.SetIterLCPmaxItersSpeed(150);
  m_system.SetIterLCPmaxItersStab(150);
  m_system.SetMaxPenetrationRecoverySpeed(4.0);
  m_system.SetStep(step_size);

  // Create the HMMWV vehicle
  HMMWV9_Vehicle vehicle(m_system,
                         ChCoordsys<>(initLoc, initRot),
                         false,
                         HMMWV9_Vehicle::NONE,
                         HMMWV9_Vehicle::PRIMITIVES);

  // Create the ground
  HMMWV9_RigidTerrain terrain(m_system, terrainHeight, terrainLength, terrainWidth, 0.8);

#ifdef USE_IRRLICHT
  irr::ChIrrApp application(&m_system,
                            L"HMMWV 9-body demo",
                            irr::core::dimension2d<irr::u32>(1000, 800),
                            false,
                            true);

  application.AddTypicalSky();
  application.AddTypicalLights();

  application.SetTimestep(step_size);

  HMMWV9_IrrGuiDriver driver(application, vehicle);

  driver.CreateCamera(cameraOffset);

  // Set up the assets for rendering
  application.AssetBindAll();
  application.AssetUpdateAll();
#endif

  // ---------------
  // Simulation loop
  // ---------------

  int out_steps = std::ceil((1 / step_size) / out_fps);

  double time = 0;
  int frame = 0;
  int out_frame = 0;

#ifdef USE_IRRLICHT
  while (application.GetDevice()->run())
  {
    bool render = (frame % out_steps == 0);

    // Render scene
    if (render) {
      application.GetVideoDriver()->beginScene(true, true, irr::video::SColor(255, 140, 161, 192));

      application.DrawAll();
      driver.DrawAll();

      out_frame++;
    }

    // Update subsystems and advance simulation by one time step
    driver.Update(time);
    vehicle.Update(time, driver.getThrottle(), driver.getSteering());

    m_system.DoStepDynamics(step_size);
    time += step_size;
    frame++;

    // Complete scene
    if (render)
      application.GetVideoDriver()->endScene();
  }
#else
  if(ChFileutils::MakeDirectory(out_dir.c_str()) < 0) {
    std::cout << "Error creating directory " << out_dir << std::endl;
    return 1;
  }
  if(ChFileutils::MakeDirectory(pov_dir.c_str()) < 0) {
    std::cout << "Error creating directory " << pov_dir << std::endl;
    return 1;
  }

  utils::WriteMeshPovray(HMMWV9_Vehicle::ChassisMeshFile(), HMMWV9_Vehicle::ChassisMeshName(), out_dir);
  utils::WriteMeshPovray(HMMWV9_Wheel::MeshFile(), HMMWV9_Wheel::MeshName(), out_dir);

  char filename[100];

  while (time < tend)
  {
    if (frame % out_steps == 0) {
      // Output render data
      sprintf(filename, "%s/data_%03d.dat", pov_dir.c_str(), out_frame + 1);
      utils::WriteShapesPovray(&m_system, filename);
      std::cout << "------------ Output frame:   " << out_frame << std::endl;
      std::cout << "             Sim frame:      " << frame << std::endl;
      std::cout << "             Time:           " << time << std::endl;
      out_frame++;

      break;

    }

    // Update subsystems and advance simulation by one time step
    ////driver.Update(time);
    ////vehicle.Update(time, driver.getThrottle(), driver.getSteering());

    m_system.DoStepDynamics(step_size);
    time += step_size;
    frame++;
  }
#endif

#ifdef USE_IRRLICHT
  application.GetDevice()->drop();
#endif

  DLL_DeleteGlobals();

  return 0;
}
<|endoftext|>


<commit_before>// =============================================================================
// PROJECT CHRONO - http://projectchrono.org
//
// Copyright (c) 2014 projectchrono.org
// All right reserved.
//
// Use of this source code is governed by a BSD-style license that can be found
// in the LICENSE file at the top level of the distribution and at
// http://projectchrono.org/license-chrono.txt.
//
// =============================================================================
// Author: Antonio Recuero
// =============================================================================
//
// Unit test for ComputeContactForces utility.
// Method system->GetContactContainer()->ComputeContactForces(); iterates over
// all bodies into contact and stores resultant contact force in an unordered map.
// Upon invocation of myBody->GetContactForce(), the user can retrieve the resultant
// of all (!) contact forces acting on the body. In this unit test, the overall
// contact force applied to a contact container is compared to the total weight
// of a number of balls.
//
// =============================================================================

#include "chrono/ChConfig.h"
#include "chrono/lcp/ChLcpIterativeMINRES.h"
#include "chrono/lcp/ChLcpSolverDEM.h"
#include "chrono/physics/ChContactContainerDEM.h"
#include "chrono/physics/ChSystemDEM.h"
#include "chrono/utils/ChUtilsCreators.h"

using namespace chrono;

// ====================================================================================
// Parameters for simulations

double endtime = 3.0;

// Relative tolerance
double rtol = 0.01;

// ---------------------
// Simulation parameters
// ---------------------

double gravity = -9.81;   // gravitational acceleration
double time_step = 5e-3;  // integration step size

enum SolverType { DEFAULT_SOLVER, MINRES_SOLVER, MKL_SOLVER };
SolverType solver_type = MINRES_SOLVER;

bool stiff_contact = true;

// ---------------------------
// Contact material properties
// ---------------------------

bool use_mat_properties = false;
bool use_history = false;
ChSystemDEM::ContactForceModel force_model = ChSystemDEM::Hooke;

float young_modulus = 2e4f;
float friction = 0.4f;
float restitution = 0.99f;
float adhesion = 0.0f;

float kn = 2e4;
float gn = 5e2;
float kt = 0;
float gt = 0;

// -------------------------------
// Parameters for the falling ball
// -------------------------------

int ballId = 100;
double radius = 0.05;
double mass = 5;
ChVector<> pos(0, 0.051, 0);
ChQuaternion<> rot(1, 0, 0, 0);
ChVector<> init_vel(0, 0, 0);
ChVector<> init_omg(0, 0, 0);

// ---------------------------------
// Parameters for the containing bin
// ---------------------------------

int binId = 200;
double width_ = 20;
double length_ = 20;
double thickness = 0.1;

// Forward declaration
bool test_computecontact(ChMaterialSurfaceBase::ContactMethod method, std::shared_ptr<ChMaterialSurfaceBase> mat);

int main(int argc, char* argv[]) {

    auto matDEM = std::make_shared<ChMaterialSurfaceDEM>();
    auto matDVI = std::make_shared<ChMaterialSurface>();

    bool DEM_passed = test_computecontact(ChMaterialSurfaceBase::DEM, matDEM);
    bool DVI_passed = test_computecontact(ChMaterialSurfaceBase::DVI, matDVI);
    if (DEM_passed && DVI_passed) {
        return 0;
    } else
        return 1;
}

// ====================================================================================

bool test_computecontact(ChMaterialSurfaceBase::ContactMethod method, std::shared_ptr<ChMaterialSurfaceBase> mat) {
    ChSystem* system =
        (method == ChMaterialSurfaceBase::DVI) ? new ChSystem : new ChSystemDEM(use_mat_properties, use_history);

    if (method == ChMaterialSurfaceBase::DEM) {
        if (auto material = std::dynamic_pointer_cast<ChMaterialSurfaceDEM>(mat)) {
            // Set the DEM contact force model
            dynamic_cast<ChSystemDEM*>(system)->SetContactForceModel(force_model);
            // Set contact forces as stiff (to force Jacobian computation) or non-stiff
            dynamic_cast<ChSystemDEM*>(system)->SetStiffContact(stiff_contact);
            // Create a material (will be used by both objects)
            material->SetYoungModulus(young_modulus);
            material->SetRestitution(restitution);
            material->SetFriction(friction);
            material->SetAdhesion(adhesion);
            material->SetKn(kn);
            material->SetGn(gn);
            material->SetKt(kt);
            material->SetGt(gt);
        }
    } else if (method == ChMaterialSurfaceBase::DVI) {
        if (auto material = std::dynamic_pointer_cast<ChMaterialSurface>(mat)) {
            // Create a material (will be used by both objects)
            material->SetRestitution(restitution);
            material->SetFriction(friction);
        }
    }
    system->Set_G_acc(ChVector<>(0, gravity, 0));

    // Establish number of falling balls
    unsigned int noBalls = 8;

    // Create the falling balls

    for (unsigned int i = 0; i < noBalls; i++) {
        auto ball = std::make_shared<ChBody>(ChMaterialSurfaceBase::DEM);  // use also complementarity

        ball->SetIdentifier(ballId);
        ball->SetMass(mass);
        ball->SetInertiaXX(0.4 * mass * radius * radius * ChVector<>(1, 1, 1));
        ball->SetPos(pos + ChVector<>(double(i) * 0.1, 0, double(i) * 0.1));
        ball->SetRot(rot);
        ball->SetPos_dt(init_vel);
        ball->SetWvel_par(init_omg);
        ball->SetCollide(true);
        ball->SetBodyFixed(false);
        ball->SetMaterialSurface(mat);

        ball->GetCollisionModel()->ClearModel();
        ball->GetCollisionModel()->AddSphere(radius);
        ball->GetCollisionModel()->BuildModel();

        system->AddBody(ball);
    }

    // Create container box
    auto ground = utils::CreateBoxContainer(system, binId, mat, ChVector<>(width_, length_, 0), thickness,
                                            ChVector<>(0, 0, 0), ChQuaternion<>(1, 0, 0, 0), true, true, false, false);
    // -------------------
    // Setup linear solver
    // -------------------

    switch (solver_type) {
        case DEFAULT_SOLVER: {
            GetLog() << "Using DEFAULT solver.\n";
            system->SetIterLCPmaxItersSpeed(100);
            system->SetTolForce(1e-6);
            break;
        }
        case MINRES_SOLVER: {
            GetLog() << "Using MINRES solver.\n";
            ChLcpIterativeMINRES* minres_solver = new ChLcpIterativeMINRES;
            minres_solver->SetDiagonalPreconditioning(true);
            system->ChangeLcpSolverSpeed(minres_solver);
            system->SetIterLCPmaxItersSpeed(100);
            system->SetTolForce(1e-6);
            break;
        }
    }

    // ----------------
    // Setup integrator
    // ----------------

    system->SetIntegrationType(ChSystem::INT_HHT);
    auto integrator = std::static_pointer_cast<ChTimestepperHHT>(system->GetTimestepper());
    integrator->SetAlpha(0.0);
    integrator->SetMaxiters(100);
    integrator->SetAbsTolerances(1e-08);
    integrator->SetScaling(false);

    // ---------------
    // Simulation loop
    double simtime = 0.0;
    // ---------------
    while (simtime < endtime) {
        system->DoStepDynamics(time_step);

        GetLog() << "t = " << system->GetChTime() << "  NR iters. = " << integrator->GetNumIterations() << "\n";
        system->GetContactContainer()->ComputeContactForces();
        GetLog() << "Total force on ground:  " << ground->GetContactForce() << "\n";

        if (simtime > 0.25) {
            if (std::abs(noBalls * mass * gravity - ground->GetContactForce().y) <
                rtol * std::abs(noBalls * mass * gravity))
                return true;
            else
                return false;
        }
        simtime += time_step;
    }
    return false;
}
<commit_msg>Remove DVI test from unit test<commit_after>// =============================================================================
// PROJECT CHRONO - http://projectchrono.org
//
// Copyright (c) 2014 projectchrono.org
// All right reserved.
//
// Use of this source code is governed by a BSD-style license that can be found
// in the LICENSE file at the top level of the distribution and at
// http://projectchrono.org/license-chrono.txt.
//
// =============================================================================
// Author: Antonio Recuero
// =============================================================================
//
// Unit test for ComputeContactForces utility.
// Method system->GetContactContainer()->ComputeContactForces(); iterates over
// all bodies into contact and stores resultant contact force in an unordered map.
// Upon invocation of myBody->GetContactForce(), the user can retrieve the resultant
// of all (!) contact forces acting on the body. In this unit test, the overall
// contact force applied to a contact container is compared to the total weight
// of a number of balls.
//
// =============================================================================

#include "chrono/ChConfig.h"
#include "chrono/lcp/ChLcpIterativeMINRES.h"
#include "chrono/lcp/ChLcpSolverDEM.h"
#include "chrono/physics/ChContactContainerDEM.h"
#include "chrono/physics/ChSystemDEM.h"
#include "chrono/utils/ChUtilsCreators.h"

using namespace chrono;

// ====================================================================================
// Parameters for simulations

double endtime = 3.0;

// Relative tolerance
double rtol = 0.01;

// ---------------------
// Simulation parameters
// ---------------------

double gravity = -9.81;   // gravitational acceleration
double time_step = 5e-3;  // integration step size

enum SolverType { DEFAULT_SOLVER, MINRES_SOLVER, MKL_SOLVER };
SolverType solver_type = MINRES_SOLVER;

bool stiff_contact = true;

// ---------------------------
// Contact material properties
// ---------------------------

bool use_mat_properties = false;
bool use_history = false;
ChSystemDEM::ContactForceModel force_model = ChSystemDEM::Hooke;

float young_modulus = 2e4f;
float friction = 0.4f;
float restitution = 0.99f;
float adhesion = 0.0f;

float kn = 2e4;
float gn = 5e2;
float kt = 0;
float gt = 0;

// -------------------------------
// Parameters for the falling ball
// -------------------------------

int ballId = 100;
double radius = 0.05;
double mass = 5;
ChVector<> pos(0, 0.051, 0);
ChQuaternion<> rot(1, 0, 0, 0);
ChVector<> init_vel(0, 0, 0);
ChVector<> init_omg(0, 0, 0);

// ---------------------------------
// Parameters for the containing bin
// ---------------------------------

int binId = 200;
double width_ = 20;
double length_ = 20;
double thickness = 0.1;

// Forward declaration
bool test_computecontact(ChMaterialSurfaceBase::ContactMethod method, std::shared_ptr<ChMaterialSurfaceBase> mat);

int main(int argc, char* argv[]) {

    auto matDEM = std::make_shared<ChMaterialSurfaceDEM>();
    auto matDVI = std::make_shared<ChMaterialSurface>();

    bool DEM_passed = test_computecontact(ChMaterialSurfaceBase::DEM, matDEM);
    // bool DVI_passed = test_computecontact(ChMaterialSurfaceBase::DVI, matDVI);
    if (DEM_passed) {
        return 0;
    } else
        return 1;
}

// ====================================================================================

bool test_computecontact(ChMaterialSurfaceBase::ContactMethod method, std::shared_ptr<ChMaterialSurfaceBase> mat) {
    ChSystem* system =
        (method == ChMaterialSurfaceBase::DVI) ? new ChSystem : new ChSystemDEM(use_mat_properties, use_history);

    if (method == ChMaterialSurfaceBase::DEM) {
        if (auto material = std::dynamic_pointer_cast<ChMaterialSurfaceDEM>(mat)) {
            // Set the DEM contact force model
            dynamic_cast<ChSystemDEM*>(system)->SetContactForceModel(force_model);
            // Set contact forces as stiff (to force Jacobian computation) or non-stiff
            dynamic_cast<ChSystemDEM*>(system)->SetStiffContact(stiff_contact);
            // Create a material (will be used by both objects)
            material->SetYoungModulus(young_modulus);
            material->SetRestitution(restitution);
            material->SetFriction(friction);
            material->SetAdhesion(adhesion);
            material->SetKn(kn);
            material->SetGn(gn);
            material->SetKt(kt);
            material->SetGt(gt);
        }
    } else if (method == ChMaterialSurfaceBase::DVI) {
        if (auto material = std::dynamic_pointer_cast<ChMaterialSurface>(mat)) {
            // Create a material (will be used by both objects)
            material->SetRestitution(restitution);
            material->SetFriction(friction);
        }
    }
    system->Set_G_acc(ChVector<>(0, gravity, 0));

    // Establish number of falling balls
    unsigned int noBalls = 8;

    // Create the falling balls

    for (unsigned int i = 0; i < noBalls; i++) {
        auto ball = std::make_shared<ChBody>(ChMaterialSurfaceBase::DEM);  // use also complementarity

        ball->SetIdentifier(ballId);
        ball->SetMass(mass);
        ball->SetInertiaXX(0.4 * mass * radius * radius * ChVector<>(1, 1, 1));
        ball->SetPos(pos + ChVector<>(double(i) * 0.1, 0, double(i) * 0.1));
        ball->SetRot(rot);
        ball->SetPos_dt(init_vel);
        ball->SetWvel_par(init_omg);
        ball->SetCollide(true);
        ball->SetBodyFixed(false);
        ball->SetMaterialSurface(mat);

        ball->GetCollisionModel()->ClearModel();
        ball->GetCollisionModel()->AddSphere(radius);
        ball->GetCollisionModel()->BuildModel();

        system->AddBody(ball);
    }

    // Create container box
    auto ground = utils::CreateBoxContainer(system, binId, mat, ChVector<>(width_, length_, 0), thickness,
                                            ChVector<>(0, 0, 0), ChQuaternion<>(1, 0, 0, 0), true, true, false, false);
    // -------------------
    // Setup linear solver
    // -------------------

    switch (solver_type) {
        case DEFAULT_SOLVER: {
            GetLog() << "Using DEFAULT solver.\n";
            system->SetIterLCPmaxItersSpeed(100);
            system->SetTolForce(1e-6);
            break;
        }
        case MINRES_SOLVER: {
            GetLog() << "Using MINRES solver.\n";
            ChLcpIterativeMINRES* minres_solver = new ChLcpIterativeMINRES;
            minres_solver->SetDiagonalPreconditioning(true);
            system->ChangeLcpSolverSpeed(minres_solver);
            system->SetIterLCPmaxItersSpeed(100);
            system->SetTolForce(1e-6);
            break;
        }
    }

    // ----------------
    // Setup integrator
    // ----------------

    system->SetIntegrationType(ChSystem::INT_HHT);
    auto integrator = std::static_pointer_cast<ChTimestepperHHT>(system->GetTimestepper());
    integrator->SetAlpha(0.0);
    integrator->SetMaxiters(100);
    integrator->SetAbsTolerances(1e-08);
    integrator->SetScaling(false);

    // ---------------
    // Simulation loop
    double simtime = 0.0;
    // ---------------
    while (simtime < endtime) {
        system->DoStepDynamics(time_step);

        // GetLog() << "t = " << system->GetChTime() << "  NR iters. = " << integrator->GetNumIterations() << "\n";
        system->GetContactContainer()->ComputeContactForces();
        // GetLog() << "Total force on ground:  " << ground->GetContactForce() << "\n";

        if (simtime > 0.25) {
            if (std::abs(noBalls * mass * gravity - ground->GetContactForce().y) <
                rtol * std::abs(noBalls * mass * gravity))
                return true;
            else
                return false;
        }
        simtime += time_step;
    }
    return false;
}
<|endoftext|>