All Downloads are FREE. Search and download functionalities are using the official Maven repository.

com.intel.analytics.bigdl.tensor.DenseTensorApply.scala Maven / Gradle / Ivy

There is a newer version: 0.11.1
Show newest version
/*
 * Copyright 2016 The BigDL Authors.
 *
 * 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.
 */

package com.intel.analytics.bigdl.tensor

object DenseTensorApply {
  /**
   * Iterate through tensor1, and apply func to the elements,
   * set function result to tensor 2
   *
   * @param tensor1 the tensor1
   * @param tensor2 the result tensor
   * @param func    (tensor1Data, tensor1Offset, tensor2Data,
   *                tensor2Offset)
   */
  def apply1[A, B](tensor1: Tensor[A], tensor2: Tensor[B],
    func: TensorDiffTypeFunc4[A, B]): Unit = {

    if (tensor1.isEmpty) {
      return
    }

    // shortcut for scalar
    if (tensor1.isScalar && tensor2.isScalar) {
      val data1 = tensor1.storage().array()
      val index1 = tensor1.storageOffset() - 1
      val data2 = tensor2.storage().array()
      val index2 = tensor2.storageOffset() - 1
      func(data1, index1, data2, index2)
      return
    }

    val stride1 = getStride(tensor1)
    val stride2 = getStride(tensor2)
    val (largestDim1, largestSize1) = getLargestContiguousSize(tensor1)
    val (largestDim2, largestSize2) = getLargestContiguousSize(tensor2)
    val counter1 = getCounter(largestDim1)
    val counter2 = getCounter(largestDim2)
    val data1 = tensor1.storage().array()
    val data2 = tensor2.storage().array()
    var offset1 = tensor1.storageOffset() - 1
    var offset2 = tensor2.storageOffset() - 1
    var hasFinished1 = false
    var hasFinished2 = false
    var i1 = 0
    var i2 = 0
    while (!hasFinished1 && !hasFinished2) {
      while (i1 < largestSize1 && i2 < largestSize2) {
        val index1 = offset1 + i1 * stride1
        val index2 = offset2 + i2 * stride2
        func(data1, index1, data2, index2)
        i1 += 1
        i2 += 1
      }
      val r1 = updateCounter(tensor1, counter1, offset1, largestDim1)
      val r2 = updateCounter(tensor2, counter2, offset2, largestDim2)
      hasFinished1 = r1._1
      hasFinished2 = r2._1
      offset1 = r1._2
      offset2 = r2._2
      i1 = 0
      i2 = 0
    }
  }

  /**
   * Iterate through tensor1, tensor2, and apply func to the elements
   *
   * @param tensor
   * @param func (tensor1Data, tensor1Offset)
   */
  def apply1[@specialized(Float, Double) T](
    tensor: Tensor[T], func: TensorFunc2[T]): Unit = {

    if (tensor.isEmpty) {
      return
    }

    // shortcut for scalar
    if (tensor.isScalar) {
      val data = tensor.storage().array()
      val index = tensor.storageOffset() - 1
      func(data, index)
      return
    }


    val stride = getStride(tensor)
    val (largestDim, largestSize) = getLargestContiguousSize(tensor)
    val counter = getCounter(largestDim)
    val data = tensor.storage().array()
    var offset = tensor.storageOffset() - 1
    var hasFinished = false
    var i = 0
    while (!hasFinished) {
      while (i < largestSize) {
        val index = offset + i * stride
        func(data, index)
        i += 1
      }
      val r = updateCounter(tensor, counter, offset, largestDim)
      hasFinished = r._1
      offset = r._2
      i = 0
    }
  }

  /**
   * Iterate through tensor1, tensor2, and apply func to the elements
   *
   * @param tensor1 the tensor
   * @param tensor2 the tensor
   * @param func    (tensor1Data, tensor1Offset, tensor2Data, tensor2Offset)
   */
  def apply2[T](tensor1: Tensor[T], tensor2: Tensor[T],
    func: TensorFunc4[T]): Unit = {
    require(tensor1.nElement() == tensor2.nElement(),
      s"inconsistent tensor size: ${tensor1.nElement()} == ${tensor2.nElement()}")

    if (tensor1.isEmpty) {
      return
    }

    // shortcut for scalar
    if (tensor1.isScalar && tensor2.isScalar) {
      val tensor1Data = tensor1.storage().array()
      val tensor2Data = tensor2.storage().array()
      val tensor1Index = tensor1.storageOffset() - 1
      val tensor2Index = tensor2.storageOffset() - 1
      func(tensor1Data, tensor1Index, tensor2Data, tensor2Index)
      return
    }

    val tensor1Data = tensor1.storage().array()
    var tensor1Offset = tensor1.storageOffset() - 1
    val tensor2Data = tensor2.storage().array()
    var tensor2Offset = tensor2.storageOffset() - 1

    var adjacent = false
    if (tensor1.nDimension == 1 && tensor2.nDimension == 1 && tensor1.stride(1) == 1 &&
      tensor2.stride(1) == 1) {
      adjacent = true
    }
    if (tensor1.nDimension == 2 && tensor2.nDimension == 2) {
      if (tensor1.stride(2) == 1 && tensor2.stride(2) == 1 && tensor1.stride(1) == tensor1.size(2)
        && tensor2.stride(1) == tensor2.size(2)) {
        adjacent = true
      }

      if (tensor1.stride(1) == 1 && tensor2.stride(1) == 1 && tensor1.stride(2) == tensor1.size(1)
        && tensor2.stride(2) == tensor2.size(1)) {
        adjacent = true
      }
    }
    if (adjacent) {
      var i = 0
      while (i < tensor1.nElement()) {
        func(tensor1Data, tensor1Offset + i, tensor2Data, tensor2Offset + i)
        i += 1
      }
      return
    }

    val tensor1Stride = getStride(tensor1)
    val (largestDim1, largestSize1) = getLargestContiguousSize(tensor1)
    val counter1 = getCounter(largestDim1)
    val tensor2Stride = getStride(tensor2)
    val (largestDim2, largestSize2) = getLargestContiguousSize(tensor2)
    val counter2 = getCounter(largestDim2)

    var hasFinished = false
    var i1 = 0
    var i2 = 0
    while (!hasFinished) {
      while (i1 < largestSize1 && i2 < largestSize2) {
        func(tensor1Data, tensor1Offset + i1 * tensor1Stride, tensor2Data,
          tensor2Offset + i2 * tensor2Stride)
        i1 = i1 + 1
        i2 = i2 + 1
      }

      if (i1 == largestSize1) {
        val r = updateCounter(tensor1, counter1, tensor1Offset, largestDim1)
        hasFinished = r._1
        tensor1Offset = r._2
        i1 = 0
      }

      if (i2 == largestSize2) {
        val r = updateCounter(tensor2, counter2, tensor2Offset, largestDim2)
        hasFinished = r._1
        tensor2Offset = r._2
        i2 = 0
      }
    }
  }

  /**
   * Iterate through tensor1, tensor2, tensor3, and apply func to the elements
   *
   * @param tensor1 the tensor
   * @param tensor2 the tensor
   * @param tensor3 the tensor
   * @param func    (tensor1Data, tensor1Offset, tensor2Data, tensor2Offset, tensor3Data,
   *                tensor3Offset)
   */
  private[bigdl] def apply3[@specialized(Float, Double) T](tensor1: Tensor[T],
    tensor2: Tensor[T], tensor3: Tensor[T],
    func: TensorFunc6[T]): Unit = {

    require(tensor1.nElement() == tensor2.nElement() && tensor2.nElement() == tensor3.nElement(),
      "inconsistent tensor size")

    if (tensor1.isEmpty) {
      return
    }

    // shortcut for scalar
    if (tensor1.isScalar && tensor2.isScalar && tensor3.isScalar) {
      val tensor1Data = tensor1.storage().array()
      val tensor2Data = tensor2.storage().array()
      val tensor3Data = tensor3.storage().array()
      val tensor1Index = tensor1.storageOffset() - 1
      val tensor2Index = tensor2.storageOffset() - 1
      val tensor3Index = tensor3.storageOffset() - 1
      func(tensor1Data, tensor1Index, tensor2Data, tensor2Index, tensor3Data, tensor3Index)
      return
    }

    val tensor1Data = tensor1.storage().array()
    var tensor1Offset = tensor1.storageOffset() - 1
    val tensor1Stride = getStride(tensor1)
    val (tensor1Dim, tensor1Size) = getLargestContiguousSize(tensor1)
    val tensor1Counter = getCounter(tensor1Dim)

    val tensor2Data = tensor2.storage().array()
    var tensor2Offset = tensor2.storageOffset() - 1
    val tensor2Stride = getStride(tensor2)
    val (tensor2Dim, tensor2Size) = getLargestContiguousSize(tensor2)
    val tensor2Counter = getCounter(tensor2Dim)

    val tensor3Data = tensor3.storage().array()
    var tensor3Offset = tensor3.storageOffset() - 1
    val tensor3Stride = getStride(tensor3)
    val (tensor3Dim, tensor3Size) = getLargestContiguousSize(tensor3)
    val tensor3Counter = getCounter(tensor3Dim)

    var hasFinished = false
    var i1 = 0
    var i2 = 0
    var i3 = 0
    while (!hasFinished) {
      while (i1 < tensor1Size && i2 < tensor2Size && i3 < tensor3Size) {
        func(tensor1Data, tensor1Offset + i1 * tensor1Stride, tensor2Data,
          tensor2Offset + i2 * tensor2Stride,
          tensor3Data, tensor3Offset + i3 * tensor3Stride)
        i1 += 1
        i2 += 1
        i3 += 1
      }

      if (i1 == tensor1Size) {
        val r = updateCounter(tensor1, tensor1Counter, tensor1Offset, tensor1Dim)
        hasFinished = r._1
        tensor1Offset = r._2
        i1 = 0
      }

      if (i2 == tensor2Size) {
        val r = updateCounter(tensor2, tensor2Counter, tensor2Offset, tensor2Dim)
        hasFinished = r._1
        tensor2Offset = r._2
        i2 = 0
      }

      if (i3 == tensor3Size) {
        val r = updateCounter(tensor3, tensor3Counter, tensor3Offset, tensor3Dim)
        hasFinished = r._1
        tensor3Offset = r._2
        i3 = 0
      }
    }
  }

  /**
   * Get the stride discard dimensions with size 1
   *
   * @param tensor tensor
   * @return
   */
  def getStride[T](tensor: Tensor[T]): Int = {
    var d = tensor.nDimension()
    while (d > 0) {
      if (tensor.size(d) != 1) {
        return tensor.stride(d)
      }
      d -= 1
    }

    0
  }

  def getLargestContiguousSize[T](tensor: Tensor[T]): (Int, Int) = {
    var largestSize = 1
    var largestDim = tensor.nDimension()
    while (largestDim > 0) {
      if (tensor.size(largestDim) != 1) {
        if (tensor.stride(largestDim) == largestSize) {
          largestSize = largestSize * tensor.size(largestDim)
        } else {
          return (largestDim, largestSize)
        }
      }
      largestDim -= 1
    }
    (largestDim, largestSize)
  }

  def getCounter(largestDim: Int): Array[Int] = {
    val counter = new Array[Int](largestDim)
    var d = 0
    while (d < largestDim) {
      counter(d) = 0
      d += 1
    }
    counter
  }

  def updateCounter[T](tensor: Tensor[T], counter: Array[Int], offset: Int,
    dim: Int): (Boolean, Int) = {
    if (dim == 0) {
      return (true, offset)
    }

    var _offset = offset
    var i = dim
    while (i > 0) {
      counter(i - 1) += 1
      _offset += tensor.stride(i)
      if (counter(i - 1) == tensor.size(i)) {
        if (i == 1) {
          return (true, _offset)
        } else {
          _offset -= counter(i - 1) * tensor.stride(i)
          counter(i - 1) = 0
        }
      } else {
        return (false, _offset)
      }
      i -= 1
    }

    (false, _offset)
  }

  /**
   * Iterate through tensor1, tensor2, and apply func to the elements,
   * set function result to tensor 3
   *
   * @param tensor1 the tensor1
   * @param tensor2 the tensor2
   * @param tensor3 the result tensor
   * @param func    (tensor1Data, tensor1Offset, tensor2Data,
   *                tensor2Offset, tensor3Data, tensor3Offset)
   */
  def apply2[A, B, C](tensor1: Tensor[A], tensor2: Tensor[B], tensor3: Tensor[C],
    func: TensorDiffTypeFunc6[A, B, C])
  : Unit = {
    require(tensor1.nElement() == tensor2.nElement(),
      s"inconsistent tensor size: ${tensor1.nElement()} == ${tensor2.nElement()}")

    if (tensor1.isEmpty) {
      return
    }

    // shortcut for scalar
    if (tensor1.isScalar && tensor2.isScalar) {
      val tensor1Data = tensor1.storage().array()
      val tensor2Data = tensor2.storage().array()
      val tensor3Data = tensor3.storage().array()
      val tensor1Index = tensor1.storageOffset() - 1
      val tensor2Index = tensor2.storageOffset() - 1
      val tensor3Index = tensor3.storageOffset() - 1
      func(tensor1Data, tensor1Index, tensor2Data, tensor2Index, tensor3Data, tensor3Index)
      return
    }

    val tensor1Data = tensor1.storage().array()
    var tensor1Offset = tensor1.storageOffset() - 1
    val tensor2Data = tensor2.storage().array()
    var tensor2Offset = tensor2.storageOffset() - 1
    val tensor3Data = tensor3.storage().array()
    val tensor3Offset = tensor3.storageOffset() - 1

    var adjacent = false
    if (tensor1.nDimension == 1 && tensor2.nDimension == 1 && tensor1.stride(1) == 1 &&
      tensor2.stride(1) == 1) {
      adjacent = true
    }
    if (tensor1.nDimension == 2 && tensor2.nDimension == 2) {
      if (tensor1.stride(2) == 1 && tensor2.stride(2) == 1 && tensor1.stride(1) == tensor1.size(2)
        && tensor2.stride(1) == tensor2.size(2)) {
        adjacent = true
      }

      if (tensor1.stride(1) == 1 && tensor2.stride(1) == 1 && tensor1.stride(2) == tensor1.size(1)
        && tensor2.stride(2) == tensor2.size(1)) {
        adjacent = true
      }
    }
    if (adjacent) {
      var i = 0
      while (i < tensor1.nElement()) {
        func(
          tensor1Data, tensor1Offset + i,
          tensor2Data, tensor2Offset + i,
          tensor3Data, tensor3Offset + i)
        i += 1
      }
      return
    }

    val tensor1Stride = getStride(tensor1)
    val (largestDim1, largestSize1) = getLargestContiguousSize(tensor1)
    val counter1 = getCounter(largestDim1)
    val tensor2Stride = getStride(tensor2)
    val (largestDim2, largestSize2) = getLargestContiguousSize(tensor2)
    val counter2 = getCounter(largestDim2)
    val tensor3Stride = getStride(tensor3)
    val (largestDim3, largestSize3) = getLargestContiguousSize(tensor3)
    val counter3 = getCounter(largestDim3)

    var hasFinished = false
    var i1 = 0
    var i2 = 0
    var i3 = 0
    while (!hasFinished) {
      while (i1 < largestSize1 && i2 < largestSize2) {
        func(
          tensor1Data, tensor1Offset + i1 * tensor1Stride,
          tensor2Data, tensor2Offset + i2 * tensor2Stride,
          tensor3Data, tensor3Offset + i3 * tensor3Stride
        )
        i1 = i1 + 1
        i2 = i2 + 1
        i3 = i3 + 1
      }

      if (i1 == largestSize1) {
        val r = updateCounter(tensor1, counter1, tensor1Offset, largestDim1)
        hasFinished = r._1
        tensor1Offset = r._2
        i1 = 0
      }

      if (i2 == largestSize2) {
        val r = updateCounter(tensor2, counter2, tensor2Offset, largestDim2)
        hasFinished = r._1
        tensor2Offset = r._2
        i2 = 0
      }

      if (i3 == largestSize3) {
        val r = updateCounter(tensor3, counter3, tensor3Offset, largestDim3)
        hasFinished = r._1
        tensor2Offset = r._2
        i3 = 0
      }
    }
  }

}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy