[DT] Support partial load/store for GPU padding encoding resolver. (iree-org#22372)

I don't see an actual use case, and the revision is mostly for
completeness. By definition, the padding resolver behaves like an
identity resolver. The main difference is that the DispatchTensorType
types get padded. The offsets, sizes and strides do not get transformed
in this case.

Signed-off-by: hanhanW <[email protected]>

Author: hanhanW

compiler/src/iree/compiler/Codegen/Common/test/materialize_encoding_into_padding.mlir

@@ -34,6 +34,42 @@ func.func @set_pad_encoding_and_store() {
 
 // -----
 
+// This tests that the padding resolver can handle partial loads/stores. The
+// offsets, sizes and strides are arbitrarily chosen in the test.
+// TODO(#20160): Move the test case to materialize_encoding_for_iree_ops.mlir.
+
+#binding_ro = #hal.pipeline.binding<storage_buffer, "ReadOnly|Indirect">
+#binding = #hal.pipeline.binding<storage_buffer, Indirect>
+#encoding_mmt = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f16, f16, f16]>
+#pad_encoding = #iree_encoding.layout<[#iree_encoding.padding<[0, 64]>]>
+func.func @set_pad_encoding_and_partial_load_store() {
+  %c0 = arith.constant 0 : index
+  %0 = hal.interface.constant.load layout(<constants = 1, bindings = [#binding_ro, #binding], flags = Indirect>) ordinal(0) : i32
+  %1 = arith.index_castui %0 : i32 to index
+  %3 = hal.interface.binding.subspan layout(<constants = 1, bindings = [#binding_ro, #binding], flags = Indirect>) binding(0) alignment(64) offset(%1) flags("ReadOnly|Indirect")
+    : !iree_tensor_ext.dispatch.tensor<readonly:tensor<2048x2048xf16>>
+  %4 = hal.interface.binding.subspan layout(<constants = 1, bindings = [#binding_ro, #binding], flags = Indirect>) binding(1) alignment(64) offset(%c0) flags(Indirect)
+    : !iree_tensor_ext.dispatch.tensor<writeonly:tensor<2048x2048xf16, #pad_encoding>>
+  %5 = iree_tensor_ext.dispatch.tensor.load %3, offsets = [0, 0], sizes = [1024, 1024], strides = [2, 2]
+    : !iree_tensor_ext.dispatch.tensor<readonly:tensor<2048x2048xf16>> -> tensor<1024x1024xf16>
+  %6 = iree_encoding.set_encoding %5 : tensor<1024x1024xf16> -> tensor<1024x1024xf16, #encoding_mmt>
+  iree_tensor_ext.dispatch.tensor.store %6, %4, offsets = [0, 0], sizes = [1024, 1024], strides = [2, 2]
+    : tensor<1024x1024xf16, #encoding_mmt> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<2048x2048xf16, #pad_encoding>>
+  return
+}
+
+// CHECK-LABEL: @set_pad_encoding_and_partial_load_store
+// CHECK:         %[[A:.+]] = hal.interface.binding.subspan layout({{.+}}) binding(0)
+// CHECK-SAME:                  !iree_tensor_ext.dispatch.tensor<readonly:tensor<2048x2048xf16>>
+// CHECK:         %[[B:.+]] = hal.interface.binding.subspan layout({{.+}}) binding(1)
+// CHECK-SAME:                  !iree_tensor_ext.dispatch.tensor<writeonly:tensor<2048x2112xf16>>
+// CHECK:         %[[LD:.+]] = iree_tensor_ext.dispatch.tensor.load %[[A]], offsets = [0, 0], sizes = [1024, 1024], strides = [2, 2]
+// CHECK-SAME:                  !iree_tensor_ext.dispatch.tensor<readonly:tensor<2048x2048xf16>> -> tensor<1024x1024xf16>
+// CHECK:         iree_tensor_ext.dispatch.tensor.store %[[LD]], %[[B]], offsets = [0, 0], sizes = [1024, 1024], strides = [2, 2]
+// CHECK-SAME:                  tensor<1024x1024xf16> -> !iree_tensor_ext.dispatch.tensor<writeonly:tensor<2048x2112xf16>>
+
+// -----
+
 #binding_ro = #hal.pipeline.binding<storage_buffer, "ReadOnly|Indirect">
 #binding = #hal.pipeline.binding<storage_buffer, Indirect>
 #encoding_mmt = #iree_encoding.encoding<operand_index = 0 : index, op_type = matmul, element_types = [f16, f16, f16]>

compiler/src/iree/compiler/Codegen/ExternalInterfaces/GPUEncodingExternalModels.cpp

@@ -593,17 +593,9 @@ struct GPUPadEncodingLayoutMaterializerAttr final
     if (!boundType || !boundType.getEncoding()) {
       return failure();
     }
-    // Only handle cases where the slice spans the whole
-    // `!iree_tensor_ext.dispatch.tensor` type.
-    // TODO(hanchung): Enable partial slices. It was copied from pattern's
-    // implementaion, i.e., the users, and it can be dropped after we move the
-    // checks to the interface implementations.
-    if (!type.doesSliceSpanWholeTensor(dynamicDims, offsets, sizes, strides)) {
-      return failure();
-    }
-    newSizes = getMixedValues(boundType.getShape(), dynamicDims, builder);
-    newOffsets.resize(newSizes.size(), builder.getIndexAttr(0));
-    newStrides.resize(newSizes.size(), builder.getIndexAttr(1));
+    newSizes.assign(sizes.begin(), sizes.end());
+    newOffsets.assign(offsets.begin(), offsets.end());
+    newStrides.assign(strides.begin(), strides.end());
     return success();
   }