[mlir][linalg] Match element type of result when doing propagation of unpack through elementwise

[qedawkins]

Summary

When propagating tensor.unpack ops through elementwise generics, a new
output tensor is needed if the element type of the input differs from
that of the output in the elementwise op.

[chelini]

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

437–456

differs

441

Do we have a shorter way to get the element type from a generic?

445

Maybe a bit shorter if we write:

if (producerUnPackOp.getDestType().getElementType() != genericOutElementType)

[qedawkins]

address comments

[qedawkins]

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

441

It's not shorter but perhaps this looks a bit cleaner? Alternatively we can try to infer the element type from the body of the generic (not sure exactly how to write that so it would be better than this).

[chelini]

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

441

Thanks! Just found getElementTypeOrSelf. You can pass to it the result Value.

[qedawkins]

Use getElementTypeOrSelf

[qedawkins]

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

441

Perfect, thanks for the tip!

[chelini]

Thanks

mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp

if (auto initTensor = genericOp.getDpsInitOperand(0)
     if (auto packOp = packedOutOperand.getDefiningOp<tensor::PackOp>())
       dest = packOp.getDest();
   }
 
   // Pack the genericOp.
   GenericOp newGenericOp = packElementWiseOp(rewriter, genericOp, dest,
                                              packedOutIndexingMap, packInfo);
 
-  auto unPackOp = unPackedOperand->get().getDefiningOp<tensor::UnPackOp>();
+  // If the output element type for the generic differs from the source
+  // unpack op, we need to create a new destination tensor.
+  auto loc = genericOp.getLoc();
+  Value unPackDest = producerUnPackOp.getDest();
+  auto genericOutElementType = genericOp.getResultTypes()
+                                   .front()
+                                   .cast<RankedTensorType>()
+                                   .getElementType();
+  if (producerUnPackOp.getDestType().getElementType() !=
+      genericOutElementType) {
+    SmallVector<OpFoldResult> unPackMixedSizes;
+    if (auto unPackEmpty = unPackDest.getDefiningOp<tensor::EmptyOp>())
+      unPackMixedSizes = unPackEmpty.getMixedSizes();
+    else
+      unPackMixedSizes = tensor::getMixedSizes(rewriter, loc, unPackDest);
+
+    unPackDest = rewriter.create<tensor::EmptyOp>(loc, unPackMixedSizes,
+                                                  genericOutElementType);
+  }
+
   // Insert an unPackOp right after the packed generic.
   Value unPackOpRes =
       rewriter
           .create<tensor::UnPackOp>(
-              genericOp.getLoc(),
+              loc,
               newGenericOp.getTiedOpResult(newGenericOp.getDpsInitOperand(0)),
-              unPackOp.getDest(), producerUnPackOp.getInnerDimsPos(),
+              unPackDest, producerUnPackOp.getInnerDimsPos(),
               producerUnPackOp.getMixedTiles(),
               producerUnPackOp.getOuterDimsPerm())
           .getResult();
 
   return std::make_tuple(newGenericOp, unPackOpRes);
 }
 
 // Wrapper pattern that applies pushDownUnPackOpThroughElemGenericOp method.

mlir/test/Dialect/Linalg/data-layout-propagation.mlir

 // CHECK: %[[UNPACK:.+]] = tensor.unpack %[[RES]] 
 // CHECK-SAME:  outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32] 
 // CHECK-SAME:  into %[[ARG0_UNPACK_EMPTY]]
 
 // -----
 
 #map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
 
+func.func @unpack_element_type_change(%arg0: tensor<12x2x56x56x32xf32>, %init: tensor<12x56x56x64xf16>) -> tensor<12x56x56x64xf16> {
+  %0 = tensor.empty() : tensor<12x56x56x64xf32>
+  %1 = tensor.unpack %arg0 outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32] into %0 : tensor<12x2x56x56x32xf32> -> tensor<12x56x56x64xf32>
+  %2 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%1: tensor<12x56x56x64xf32>) outs(%init : tensor<12x56x56x64xf16>) {
+    ^bb0(%in: f32, %out: f16):
+      %3 = arith.truncf %in : f32 to f16
+      linalg.yield %3 : f16
+  } -> tensor<12x56x56x64xf16>
+  return %2 : tensor<12x56x56x64xf16>
+}
+
+// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
+// CHECK: func.func @unpack_element_type_change
+// CHECK-SAME:  %[[ARG0:[a-zA-Z0-9]+]]
+// CHECK-SAME:  %[[ARG1:[a-zA-Z0-9]+]]
+// CHECK: %[[ARG0_UNPACK_EMPTY:.+]] = tensor.empty() : tensor<12x56x56x64xf32>
+// CHECK: %[[UNPACKED_ARG0:.+]] = tensor.unpack %[[ARG0]]
+// CHECK-SAME:  outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32]
+// CHECK-SAME:  into %[[ARG0_UNPACK_EMPTY]]
+// CHECK: %[[ARG1_PACK_EMPTY:.+]] = tensor.empty() : tensor<12x2x56x56x32xf16>
+// CHECK: %[[ARG1_PACK:.+]] = tensor.pack %[[ARG1]]
+// CHECK-SAME:  outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32]
+// CHECK-SAME:  into %[[ARG1_PACK_EMPTY]]
+// CHECK: %[[ARG0_PACK_EMPTY:.+]] = tensor.empty() : tensor<12x2x56x56x32xf32>
+// CHECK: %[[ARG0_PACK:.+]] = tensor.pack %[[UNPACKED_ARG0]]
+// CHECK-SAME:  outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32]
+// CHECK-SAME:  into %[[ARG0_PACK_EMPTY]]
+// CHECK: %[[RES:.+]] = linalg.generic
+// CHECK-SAME:  indexing_maps = [#[[MAP]], #[[MAP]]]
+// CHECK-SAME:  ins(%[[ARG0_PACK]]
+// CHECK-SAME:  outs(%[[ARG1_PACK]]
+// CHECK: %[[ARG0_NEW_EMPTY_UNPACK:.+]] = tensor.empty() : tensor<12x56x56x64xf16>
+// CHECK: %[[UNPACK:.+]] = tensor.unpack %[[RES]]
+// CHECK-SAME:  outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32]
+// CHECK-SAME:  into %[[ARG0_NEW_EMPTY_UNPACK]]
+
+// -----
+
+#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+
 func.func @forward_tensor_empty(%arg0: tensor<12x2x56x56x32xf32>) -> tensor<12x56x56x64xf32> {
   %init = tensor.empty() : tensor<12x56x56x64xf32>
   %0 = tensor.empty() : tensor<12x56x56x64xf32>
   %1 = tensor.unpack %arg0 outer_dims_perm = [0, 3, 1, 2] inner_dims_pos = [3] inner_tiles = [32] into %0 : tensor<12x2x56x56x32xf32> -> tensor<12x56x56x64xf32>
   %2 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%1: tensor<12x56x56x64xf32>) outs(%init : tensor<12x56x56x64xf32>) {
     ^bb0(%in: f32, %out: f32):
       %3 = arith.addf %in, %in : f32
       linalg.yield %3 : f32