[ROCm] OffsetCalc Unroll Optimization (#2597)

cherry-pick of pytorch#161700

Our compiler is generating inefficient code for the offsetCalc in
certain situations. The root-cause for this needs to be identified. For
now specialized unrolling based on 'dims' notably helps perf.

Fixes SWDEV-545713, SWDEV-545710

Author: jerrymannil

aten/src/ATen/cuda/detail/OffsetCalculator.cuh

@@ -45,6 +45,24 @@ struct OffsetCalculator {
 
   C10_HOST_DEVICE offset_type get(index_t linear_idx) const {
     offset_type offsets;
+
+#if defined(USE_ROCM)
+    if ((dims > 0) && (dims <= 2)) {
+      auto divmod = sizes_[0].divmod(linear_idx);
+      #pragma unroll
+      for (int arg = 0; arg < NARGS; arg++)
+        offsets[arg] = divmod.mod * strides_[0][arg];
+      if (dims >= 2) {
+        divmod = sizes_[1].divmod(divmod.div);
+        #pragma unroll
+        for (int arg = 0; arg < NARGS; arg++)
+          offsets[arg] += divmod.mod * strides_[1][arg];
+      }
+      // [...]
+      return offsets;
+    }
+#endif
+
     #pragma unroll
     for (int arg = 0; arg < NARGS; arg++) {
       offsets[arg] = 0;