From 307124535f326d75725ecf6d3f4e9e0321162e9a Mon Sep 17 00:00:00 2001
From: Stephan Herhut <herhut@google.com>
Date: Thu, 15 Oct 2020 16:27:31 +0200
Subject: [PATCH] [mlir][standard] Fix parsing of scalar subview and
 canonicalize

Parsing of a scalar subview did not create the required static_offsets attribute.
This also adds support for folding scalar subviews away.

Differential Revision: https://reviews.llvm.org/D89467
---
 mlir/include/mlir/Dialect/StandardOps/IR/Ops.td |  1 +
 mlir/lib/Dialect/StandardOps/IR/Ops.cpp         | 11 ++++++++++-
 mlir/test/IR/core-ops.mlir                      |  3 +++
 mlir/test/Transforms/constant-fold.mlir         |  9 +++++++++
 4 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
index 6b60c2f79b01..c024e19b5009 100644
--- a/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
+++ b/mlir/include/mlir/Dialect/StandardOps/IR/Ops.td
@@ -3097,6 +3097,7 @@ def SubViewOp : BaseOpWithOffsetSizesAndStrides<
   }];
 
   let hasCanonicalizer = 1;
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
index 82058fdcc03c..8fe45cbb1a13 100644
--- a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
+++ b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
@@ -2531,8 +2531,10 @@ parseListOfOperandsOrIntegers(OpAsmParser &parser, OperationState &result,
   if (failed(parser.parseLSquare()))
     return failure();
   // 0-D.
-  if (succeeded(parser.parseOptionalRSquare()))
+  if (succeeded(parser.parseOptionalRSquare())) {
+    result.addAttribute(attrName, parser.getBuilder().getArrayAttr({}));
     return success();
+  }
 
   SmallVector<int64_t, 4> attrVals;
   while (true) {
@@ -3333,6 +3335,13 @@ void SubViewOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                  SubViewOpMemRefCastFolder>(context);
 }
 
+OpFoldResult SubViewOp::fold(ArrayRef<Attribute> operands) {
+  if (getResultRank() == 0 && getSourceRank() == 0)
+    return getViewSource();
+
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // SubTensorOp
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/IR/core-ops.mlir b/mlir/test/IR/core-ops.mlir
index 219c3bc84d57..da7394eae784 100644
--- a/mlir/test/IR/core-ops.mlir
+++ b/mlir/test/IR/core-ops.mlir
@@ -827,6 +827,9 @@ func @memref_subview(%arg0 : index, %arg1 : index, %arg2 : index) {
   %21 = subview %20[0, 0, 0][1, 16, 4][1, 1, 1] : memref<8x16x4xf32> to memref<16x4xf32>
 
   %22 = subview %20[3, 4, 2][1, 6, 3][1, 1, 1] : memref<8x16x4xf32> to memref<6x3xf32, offset: 210, strides: [4, 1]>
+
+  %23 = alloc() : memref<f32>
+  %78 = subview %23[] [] []  : memref<f32> to memref<f32>
   return
 }
 
diff --git a/mlir/test/Transforms/constant-fold.mlir b/mlir/test/Transforms/constant-fold.mlir
index 36fa234213ea..c75c89877830 100644
--- a/mlir/test/Transforms/constant-fold.mlir
+++ b/mlir/test/Transforms/constant-fold.mlir
@@ -744,3 +744,12 @@ func @splat_fold() -> (vector<4xf32>, tensor<4xf32>) {
   // CHECK-NEXT: [[T:%.*]] = constant dense<1.000000e+00> : tensor<4xf32>
   // CHECK-NEXT: return [[V]], [[T]] : vector<4xf32>, tensor<4xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @subview_scalar_fold
+func @subview_scalar_fold(%arg0: memref<f32>) -> memref<f32> {
+  // CHECK-NOT: subview
+  %c = subview %arg0[] [] [] : memref<f32> to memref<f32>
+  return %c : memref<f32>
+}