[Mlir-commits] [mlir] [mlir][Vector] Add support for 0-d shapes in extract-shape_cast folder (PR #116650)

[Kunwar Grover]

https://github.com/Groverkss updated https://github.com/llvm/llvm-project/pull/116650

>From 1ea6677dfdf8548690ea06135e6b71edcf80321c Mon Sep 17 00:00:00 2001
From: Kunwar Grover <groverkss at gmail.com>
Date: Mon, 18 Nov 2024 16:31:02 +0000
Subject: [PATCH 1/2] [mlir][Vector] Add support for 0-d shapes in
 extract-shape_cast folder

---
 mlir/lib/Dialect/Vector/IR/VectorOps.cpp   |  5 ---
 mlir/test/Dialect/Vector/canonicalize.mlir | 40 ++++++++++++++--------
 2 files changed, 26 insertions(+), 19 deletions(-)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 0c0a7bc98d8b5e..ad709813c6216a 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -1756,11 +1756,6 @@ static Value foldExtractFromShapeCast(ExtractOp extractOp) {
   if (!shapeCastOp)
     return Value();
 
-  // 0-D vectors not supported.
-  assert(!hasZeroDimVectors(extractOp) && "0-D vectors not supported");
-  if (hasZeroDimVectors(shapeCastOp))
-    return Value();
-
   // Get the nth dimension size starting from lowest dimension.
   auto getDimReverse = [](VectorType type, int64_t n) {
     return type.getShape().take_back(n + 1).front();
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 5ae769090dac66..04518a56c3dd20 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -782,23 +782,23 @@ func.func @fold_extract_shapecast(%arg0 : vector<5x1x3x2xf32>,
 
 // -----
 
-// CHECK-LABEL: fold_extract_shapecast_negative
-//       CHECK:   %[[V:.*]] = vector.shape_cast %{{.*}} : vector<16xf32> to vector<2x4x2xf32>
-//       CHECK:   %[[R:.*]] = vector.extract %[[V]][1] : vector<4x2xf32> from vector<2x4x2xf32>
-//       CHECK:   return %[[R]] : vector<4x2xf32>
-func.func @fold_extract_shapecast_negative(%arg0 : vector<16xf32>) -> vector<4x2xf32> {
-  %0 = vector.shape_cast %arg0 : vector<16xf32> to vector<2x4x2xf32>
-  %r = vector.extract %0[1] : vector<4x2xf32> from vector<2x4x2xf32>
-  return %r : vector<4x2xf32>
+// CHECK-LABEL: fold_extract_shapecast_0d_result
+//  CHECK-SAME: %[[IN:.*]]: vector<1x1x1xf32>
+//       CHECK:   %[[R:.*]] = vector.extract %[[IN]][0, 0, 0] : f32 from vector<1x1x1xf32>
+//       CHECK:   return %[[R]] : f32
+func.func @fold_extract_shapecast_0d_result(%arg0 : vector<1x1x1xf32>) -> f32 {
+  %0 = vector.shape_cast %arg0 : vector<1x1x1xf32> to vector<f32>
+  %r = vector.extract %0[] : f32 from vector<f32>
+  return %r : f32
 }
 
 // -----
 
-// CHECK-LABEL: dont_fold_0d_extract_shapecast
-//       CHECK:   %[[V:.*]] = vector.shape_cast %{{.*}} : vector<f32> to vector<1xf32>
-//       CHECK:   %[[R:.*]] = vector.extract %[[V]][0] : f32 from vector<1xf32>
+// CHECK-LABEL: fold_extract_shapecast_0d_source
+//  CHECK-SAME: %[[IN:.*]]: vector<f32>
+//       CHECK:   %[[R:.*]] = vector.extract %[[IN]][] : f32 from vector<f32>
 //       CHECK:   return %[[R]] : f32
-func.func @dont_fold_0d_extract_shapecast(%arg0 : vector<f32>) -> f32 {
+func.func @fold_extract_shapecast_0d_source(%arg0 : vector<f32>) -> f32 {
   %0 = vector.shape_cast %arg0 : vector<f32> to vector<1xf32>
   %r = vector.extract %0[0] : f32 from vector<1xf32>
   return %r : f32
@@ -806,11 +806,23 @@ func.func @dont_fold_0d_extract_shapecast(%arg0 : vector<f32>) -> f32 {
 
 // -----
 
-// CHECK-LABEL: fold_extract_shapecast_to_shapecast
+// CHECK-LABEL: fold_extract_shapecast_negative
+//       CHECK:   %[[V:.*]] = vector.shape_cast %{{.*}} : vector<16xf32> to vector<2x4x2xf32>
+//       CHECK:   %[[R:.*]] = vector.extract %[[V]][1] : vector<4x2xf32> from vector<2x4x2xf32>
+//       CHECK:   return %[[R]] : vector<4x2xf32>
+func.func @fold_extract_shapecast_negative(%arg0 : vector<16xf32>) -> vector<4x2xf32> {
+  %0 = vector.shape_cast %arg0 : vector<16xf32> to vector<2x4x2xf32>
+  %r = vector.extract %0[1] : vector<4x2xf32> from vector<2x4x2xf32>
+  return %r : vector<4x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: canonicalize_extract_shapecast_to_shapecast
 //  CHECK-SAME: (%[[ARG:.+]]: vector<3x4xf32>)
 //       CHECK:   %[[R:.+]] = vector.shape_cast %[[ARG]] : vector<3x4xf32> to vector<12xf32>
 //       CHECK:   return %[[R]]
-func.func @fold_extract_shapecast_to_shapecast(%arg0 : vector<3x4xf32>) -> vector<12xf32> {
+func.func @canonicalize_extract_shapecast_to_shapecast(%arg0 : vector<3x4xf32>) -> vector<12xf32> {
   %0 = vector.shape_cast %arg0 : vector<3x4xf32> to vector<1x12xf32>
   %r = vector.extract %0[0] : vector<12xf32> from vector<1x12xf32>
   return %r : vector<12xf32>

>From e3419a9b4c8276a806b92be0f672ca8ad83d1be1 Mon Sep 17 00:00:00 2001
From: Kunwar Grover <groverkss at gmail.com>
Date: Tue, 3 Dec 2024 15:42:29 +0000
Subject: [PATCH 2/2] canonicalize -> fold

---
 mlir/test/Dialect/Vector/canonicalize.mlir | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 04518a56c3dd20..89af0f7332f5c4 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -818,11 +818,11 @@ func.func @fold_extract_shapecast_negative(%arg0 : vector<16xf32>) -> vector<4x2
 
 // -----
 
-// CHECK-LABEL: canonicalize_extract_shapecast_to_shapecast
+// CHECK-LABEL: fold_extract_shapecast_to_shapecast
 //  CHECK-SAME: (%[[ARG:.+]]: vector<3x4xf32>)
 //       CHECK:   %[[R:.+]] = vector.shape_cast %[[ARG]] : vector<3x4xf32> to vector<12xf32>
 //       CHECK:   return %[[R]]
-func.func @canonicalize_extract_shapecast_to_shapecast(%arg0 : vector<3x4xf32>) -> vector<12xf32> {
+func.func @fold_extract_shapecast_to_shapecast(%arg0 : vector<3x4xf32>) -> vector<12xf32> {
   %0 = vector.shape_cast %arg0 : vector<3x4xf32> to vector<1x12xf32>
   %r = vector.extract %0[0] : vector<12xf32> from vector<1x12xf32>
   return %r : vector<12xf32>