Do not use ReduceRankOp in makeSubview (#151)

Author: Hardcode84

mlir/lib/dialect/plier_util/dialect.cpp

@@ -79,6 +79,10 @@ mlir::Operation *PlierUtilDialect::materializeConstant(mlir::OpBuilder &builder,
   if (mlir::arith::ConstantOp::isBuildableWith(value, type))
     return builder.create<mlir::arith::ConstantOp>(loc, type, value);
 
+  if (type.isa<mlir::IndexType>())
+    if (auto val = mlir::getConstantIntValue(value))
+      return builder.create<mlir::arith::ConstantIndexOp>(loc, *val);
+
   return nullptr;
 }
 

numba_dpcomp/numba_dpcomp/mlir_compiler/lib/pipelines/lower_to_gpu.cpp

@@ -744,13 +744,7 @@ struct FlattenSubview : public mlir::OpRewritePattern<mlir::memref::SubViewOp> {
       flatSubview = rewriter.createOrFold<mlir::memref::CastOp>(
           loc, dstFlatType, flatSubview);
 
-    // TODO: bug in ReinterpretCastOp::verify
-    auto offset =
-        (mlir::ShapedType::isDynamicStrideOrOffset(resultOffset)
-             ? mlir::OpFoldResult(
-                   rewriter.create<mlir::arith::ConstantIndexOp>(loc, 0)
-                       .getResult())
-             : mlir::OpFoldResult(rewriter.getIndexAttr(0)));
+    auto offset = rewriter.getIndexAttr(0);
 
     for (auto i : llvm::seq<size_t>(0, strides.size())) {
       if (mlir::ShapedType::isDynamicStrideOrOffset(resultStrides[i])) {
@@ -789,7 +783,7 @@ struct FlattenSubview : public mlir::OpRewritePattern<mlir::memref::SubViewOp> {
 
       auto droppedDims = op.getDroppedDims();
       for (auto i : llvm::seq(0u, srcRank)) {
-        if (droppedDims[i]) {
+        if (!droppedDims[i]) {
           filteredSizes.emplace_back(sizes[i]);
           filteredStrides.emplace_back(strides[i]);
         }

numba_dpcomp/numba_dpcomp/mlir_compiler/lib/pipelines/plier_to_linalg.cpp

@@ -744,51 +744,59 @@ static mlir::Value makeSubview(mlir::OpBuilder &builder, mlir::Location loc,
   assert(dstRank > 0);
   assert(dstRank <= srcRank);
 
-  bool useReduceRank = true;
-
   mlir::Value view;
   if (isTensor) {
-    auto resType =
-        [&]() {
-          auto tensorType = srcType.cast<mlir::RankedTensorType>();
-          if (srcRank == dstRank || useReduceRank)
-            return mlir::tensor::ExtractSliceOp::inferResultType(
-                tensorType, offsets, sizes, strides);
-
-          return mlir::tensor::ExtractSliceOp::inferRankReducedResultType(
-              dstRank, tensorType, offsets, sizes, strides);
-        }()
-            .cast<mlir::RankedTensorType>();
+    auto tensorType = srcType.cast<mlir::RankedTensorType>();
+    auto resType = mlir::tensor::ExtractSliceOp::inferResultType(
+                       tensorType, offsets, sizes, strides)
+                       .cast<mlir::RankedTensorType>();
 
     view = builder.create<mlir::tensor::ExtractSliceOp>(
         loc, resType, src, offsets, sizes, strides);
+
+    if (srcRank != dstRank) {
+      llvm::SmallVector<mlir::OpFoldResult> newOfsets(srcRank,
+                                                      builder.getIndexAttr(0));
+      llvm::SmallVector<mlir::OpFoldResult> newStrides(srcRank,
+                                                       builder.getIndexAttr(1));
+      auto viewType = view.getType().cast<mlir::RankedTensorType>();
+      auto reducedType =
+          mlir::tensor::ExtractSliceOp::inferRankReducedResultType(
+              dstRank, viewType, newOfsets, sizes, newStrides)
+              .cast<mlir::RankedTensorType>();
+      view = builder.create<mlir::tensor::ExtractSliceOp>(
+          loc, reducedType, view, newOfsets, sizes, newStrides);
+    }
   } else {
-    auto resType =
-        [&]() {
-          auto memrefType = srcType.cast<mlir::MemRefType>();
-          if (srcRank == dstRank || useReduceRank)
-            return mlir::memref::SubViewOp::inferResultType(memrefType, offsets,
-                                                            sizes, strides);
-
-          return mlir::memref::SubViewOp::inferRankReducedResultType(
-              dstRank, memrefType, offsets, sizes, strides);
-        }()
-            .cast<mlir::MemRefType>();
+    auto memrefType = srcType.cast<mlir::MemRefType>();
+    auto resType = mlir::memref::SubViewOp::inferResultType(memrefType, offsets,
+                                                            sizes, strides)
+                       .cast<mlir::MemRefType>();
 
     view = builder.create<mlir::memref::SubViewOp>(loc, resType, src, offsets,
                                                    sizes, strides);
 
+    if (srcRank != dstRank) {
+      llvm::SmallVector<mlir::OpFoldResult> newOfsets(srcRank,
+                                                      builder.getIndexAttr(0));
+      llvm::SmallVector<mlir::OpFoldResult> newStrides(srcRank,
+                                                       builder.getIndexAttr(1));
+      auto viewType = view.getType().cast<mlir::MemRefType>();
+      auto reducedType = mlir::memref::SubViewOp::inferRankReducedResultType(
+                             dstRank, viewType, newOfsets, sizes, newStrides)
+                             .cast<mlir::MemRefType>();
+      view = builder.create<mlir::memref::SubViewOp>(
+          loc, reducedType, view, newOfsets, sizes, newStrides);
+      resType = reducedType;
+    }
+
     auto flatMemrefType =
         mlir::MemRefType::get(resType.getShape(), resType.getElementType());
+
     if (resType != flatMemrefType)
       view = builder.create<plier::ChangeLayoutOp>(loc, flatMemrefType, view);
   }
 
-  if (srcRank != dstRank && useReduceRank) {
-    llvm::SmallVector<int32_t> mapping(dimIndices.begin(), dimIndices.end());
-    view = builder.createOrFold<plier::ReduceRankOp>(loc, view, mapping);
-  }
-
   return view;
 }