[Mlir-commits] [mlir] 8f2457c - [mlir][Linalg] Add missing support for memory space to DropUnitDims

[Nicolas Vasilache]

Author: Nicolas Vasilache
Date: 2022-09-01T04:10:27-07:00
New Revision: 8f2457ccf01f00bb83bd391ac04209e75361fb9a

URL: https://github.com/llvm/llvm-project/commit/8f2457ccf01f00bb83bd391ac04209e75361fb9a
DIFF: https://github.com/llvm/llvm-project/commit/8f2457ccf01f00bb83bd391ac04209e75361fb9a.diff

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

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

Differential Revision: https://reviews.llvm.org/D133101

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
    mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
index 0d7e41062541b..cfc50f985a584 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
@@ -293,10 +293,11 @@ replaceUnitExtents(GenericOp genericOp, OpOperand *opOperand,
     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(),

diff  --git a/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir b/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
index a272dea3fa232..27f540c8e7159 100644
--- a/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
+++ b/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
@@ -847,3 +847,28 @@ func.func @reduce_dispatch_0() -> 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>