[xla:gpu] Add IsContiguousSlice on Shapes instead of HloInstruction

PiperOrigin-RevId: 615608156

Author: tyb0807

xla/service/gpu/ir_emission_utils.cc

@@ -216,13 +216,17 @@ bool IsContiguousSlice(const HloInstruction& instr) {
   // src and dst dimensions match.
   const Shape& src_shape = slice->operand(0)->shape();
   const Shape& dst_shape = slice->shape();
+  return IsContiguousSlice(src_shape, dst_shape);
+}
+
+bool IsContiguousSlice(const Shape& orig, const Shape& sliced) {
   bool sliced_dim_found = false;
-  for (auto dim : src_shape.layout().minor_to_major()) {
+  for (auto dim : orig.layout().minor_to_major()) {
     if (!sliced_dim_found) {
-      sliced_dim_found = dst_shape.dimensions(dim) < src_shape.dimensions(dim);
+      sliced_dim_found = sliced.dimensions(dim) < orig.dimensions(dim);
       continue;
     }
-    if (dst_shape.dimensions(dim) != 1) return false;
+    if (sliced.dimensions(dim) != 1) return false;
   }
   return true;
 }

xla/service/gpu/ir_emission_utils.h

@@ -101,6 +101,9 @@ bool IsSliceWithUnitStrides(const HloInstruction* instr);
 // slice.
 bool IsContiguousSlice(const HloInstruction& instr);
 
+// Returns true if `sliced` is a contiguous slice of `orig`.
+bool IsContiguousSlice(const Shape& orig, const Shape& sliced);
+
 // Emits code to shuffle data between threads of a warp. This has the same
 // semantics as the PTX "shfl.sync.down" instruction but works for values that
 // aren't 32 bits in size. The last operand of the emitted "shfl" is