message: [XLA:GPU] Get rid of unnecessary urem in IrArray with linear index.
diff: @@ -39,13 +39,27 @@ IrArray::Index::Index(llvm::Value* linear, const Shape& shape,
  << "Shape " << ShapeUtil::HumanStringWithLayout(shape)
  << " should have a layout.";
  int64 divisor = 1;
  - for (int64 dimension : LayoutUtil::MinorToMajor(layout_)) {
  + for (int64 i = 0; i < layout_.minor_to_major_size(); ++i) {
  + int64 dimension = layout_.minor_to_major(i);
  int64 size_of_current_dimension = shape.dimensions(dimension);
  - // Emit IR instructions that compute
  - // (linear_index / divisor) % current_dimension
  +
  + // If i is not the last dimension, compute
  + // (linear_index / divisor) % current_dimension.
  + // If i is the last dimension, we can skip the mod, because we assume that
  + // linear is in bounds.
  + //
  + // TODO(jlebar): We could add bounds checks here and elsewhere in this file,
  + // guarded under some sort of xla-memcheck flag. This might be particularly
  + // useful because cuda-memcheck can't help us much in XLA: Most of our
  + // memory lives in one big allocation, so cuda-memcheck can't detect
  + // out-of-bounds accesses.
  + auto* quot = ir_builder->CreateUDiv(linear, ir_builder->getInt64(divisor));
  + if (i < layout_.minor_to_major_size() - 1) {
  multidim_[dimension] = ir_builder->CreateURem(
  - ir_builder->CreateUDiv(linear, ir_builder->getInt64(divisor)),
  - ir_builder->getInt64(size_of_current_dimension));
  + quot, ir_builder->getInt64(size_of_current_dimension));
  + } else {
  + multidim_[dimension] = quot;
  + }
  divisor *= size_of_current_dimension;
  }
  }
  