Tensor FFModel::create_tensor()

in src/runtime/model.cc [831:925]

• 87 lines of code
• 17 McCabe index (conditional complexity)

Tensor FFModel::create_tensor(const int dims[],
                              DataType data_type,
                              const Op* owner_op,
                              bool create_grad)
{
  Tensor tensor;
  tensor.data_type = data_type;
  Context ctx = config.lg_ctx;
  Runtime* runtime = config.lg_hlr;

  std::string name = "";
  if (owner_op != NULL)
    name = std::string(owner_op->name);
  IndexSpaceT<NDIM> part_is = (IndexSpaceT<NDIM>) get_or_create_task_is(NDIM, name);
  // Step 1: create regions
  FieldSpace fs = runtime->create_field_space(ctx);
  FieldAllocator allocator= runtime->create_field_allocator(ctx, fs);
  switch (data_type)
  {
    case DT_FLOAT:
      allocator.allocate_field(sizeof(float), FID_DATA);
      break;
    case DT_DOUBLE:
      allocator.allocate_field(sizeof(double), FID_DATA);
      break;
    case DT_INT32:
      allocator.allocate_field(sizeof(int32_t), FID_DATA);
      break;
    case DT_INT64:
      allocator.allocate_field(sizeof(int64_t), FID_DATA);
      break;
    default:
      assert(false);
  }
  Point<NDIM> hi;
  for (int i = 0; i < NDIM; i++)
    hi[i] = dims[NDIM-1-i]-1;
  Rect<NDIM> rect(Point<NDIM>::ZEROES(), hi);
  IndexSpaceT<NDIM> is = runtime->create_index_space(ctx, rect);
  tensor.region = runtime->create_logical_region(ctx, is, fs);
  if (create_grad && config.computationMode == COMP_MODE_TRAINING) {
    tensor.region_grad = runtime->create_logical_region(ctx, is, fs);
  }

  // Step 2: create partitions
  Rect<NDIM> part_rect = runtime->get_index_space_domain(ctx, part_is);

  Transform<NDIM, NDIM> transform;
  Point<NDIM> ext_hi;
  for (int i = 0; i < NDIM; i++) {
    int nparts = part_rect.hi[i] - part_rect.lo[i] + 1;
    ext_hi[i] = (rect.hi[i] - rect.lo[i] + nparts) / nparts - 1;
  }
  Rect<NDIM> extent(Point<NDIM>::ZEROES(), ext_hi);
  for (int i = 0; i < NDIM; i++)
    for (int j = 0; j < NDIM; j++)
      if (i == j)
        transform[i][j] = extent.hi[i] - extent.lo[i] + 1;
      else
        transform[i][j] = 0;
  IndexPartition ip = runtime->create_partition_by_restriction(
      ctx, is, part_is, transform, extent);
  assert(runtime->is_index_partition_disjoint(ctx, ip));
  assert(runtime->is_index_partition_complete(ctx, ip));
  tensor.part = runtime->get_logical_partition(ctx, tensor.region, ip);
  if (create_grad && config.computationMode == COMP_MODE_TRAINING) {
    tensor.part_grad = runtime->get_logical_partition(ctx, tensor.region_grad, ip);
  }
  tensor.numDim = NDIM;
  for (int i = 0; i < NDIM; i++) {
    tensor.adim[i] = rect.hi[i] - rect.lo[i] + 1;
    //tensor.pdim[i] = extent.hi[i] - extent.lo[i] + 1;
  }

#ifdef DEADCODE
  // Initialize tensor with zero
  ArgumentMap argmap;
  IndexLauncher launcher(ZERO_INIT_TASK_ID, part_is,
                         TaskArgument(NULL, 0), argmap,
                         Predicate::TRUE_PRED, false, 0,
                         FFConfig::get_hash_id(name));
  launcher.add_region_requirement(
      RegionRequirement(tensor.part, 0/*projection id*/,
                        WRITE_ONLY, EXCLUSIVE, tensor.region));
  launcher.add_field(0, FID_DATA);
  if (create_grad) {
    launcher.add_region_requirement(
        RegionRequirement(tensor.part_grad, 0/*projection id*/,
                          WRITE_ONLY, EXCLUSIVE, tensor.region_grad));
    launcher.add_field(1, FID_DATA);
  }
  runtime->execute_index_space(ctx, launcher);
#endif
  return tensor;
}