[mlir][Linalg] Add missing support for memory space to DropUnitDims

[nicolasvasilache]

Summary

Previously, the memory space would not be propagated properly, creating invalid IR.

[pifon2a]

Thank you, Nicolas!

mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir

861–862

nit: wrap IR to 80 chars

873

nit: split the line, that way you won't need to do %{{.*}}

// CHECK: memref.collapse_shape 
// CHECK-SAME: [] memref<1x1xf32, 3> into memref<f32, 3>

[nicolasvasilache]

mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir

873

Thanks, addressed in 153709c0934dd615cda3c9e496bd04beae8da595

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

replaceUnitExtents(GenericOp genericOp, OpOperand *opOperand,
 
   // Compute the tensor or scalar replacement type.
   Type elementType = getElementTypeOrSelf(opOperand->get());
   Type replacementType;
   if (elementType == opOperand->get().getType()) {
     replacementType = elementType;
   } else if (actualType.isa<RankedTensorType>()) {
     replacementType = RankedTensorType::get(newShape, elementType);
-  } else if (actualType.isa<MemRefType>()) {
-    replacementType = MemRefType::get(newShape, elementType);
+  } else {
+    auto memrefType = actualType.cast<MemRefType>();
+    replacementType = MemRefType::get(newShape, elementType, {},
+                                      memrefType.getMemorySpaceAsInt());
   }
-  assert(replacementType && "unsupported shaped type");
   UnitExtentReplacementInfo info = {replacementType,
                                     AffineMap::get(indexingMap.getNumDims(),
                                                    indexingMap.getNumSymbols(),
                                                    newIndexExprs, context),
                                     ArrayAttr::get(context, reassociationMaps)};
   return info;
 }
 

mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir

scf.foreach_thread.perform_concurrently {
       //      CHECK: tensor.parallel_insert_slice %{{[0-9a-z]*}} into %{{[0-9a-z]*}}
       // CHECK-SAME: [%{{.*}}, %{{.*}}] [1, 1] [1, 1] : tensor<f32> into tensor<4x2xf32>
       tensor.parallel_insert_slice %2 into %0[%arg0, %arg1] [1, 1] [1, 1] :
         tensor<1x1xf32> into tensor<4x2xf32>
     }
   }  
   return %res: tensor<4x2xf32>
 }
+
+// -----
+
+#map0 = affine_map<(i, j) -> (i, j)>
+#access = [#map0, #map0]
+#trait = {
+  iterator_types = ["parallel", "parallel"],
+  indexing_maps = #access,
+  library_call = "some_external_func"
+}
+
+func.func @drop_all_loops(%arg0 : memref<1x1xf32, 3>) -> memref<1x1xf32, 3>
+{
+  linalg.generic #trait
+     ins(%arg0 : memref<1x1xf32, 3>)
+    outs(%arg0 : memref<1x1xf32, 3>) {
+       ^bb0(%arg1: f32, %arg2: f32) :
+         linalg.yield %arg1 : f32
+       }
+  return %arg0 : memref<1x1xf32, 3>
+}
+
+// CHECK-LABEL: func @drop_all_loops
+//       CHECK:   memref.collapse_shape %{{.*}} [] memref<1x1xf32, 3> into memref<f32, 3>
+//       CHECK:   linalg.generic{{.*}}memref<f32, 3>