[Mlir-commits] [mlir] daf8f9f - [Linalg] Add rank zero operand support to push down extract slice pattern (#157532)

[llvmlistbot at llvm.org]

Author: Nirvedh Meshram
Date: 2025-09-08T14:53:02-05:00
New Revision: daf8f9fc1ccc6c5679bc89058fd66d8ea4da9d59

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

LOG: [Linalg] Add rank zero operand support to push down extract slice pattern (#157532)

Currently the pattern would crash for rank 0 operand as it decides the
padding based on affine results, but for rank 0 there are no affine
results in the operand affine map

Signed-off-by: Nirvedh Meshram <nirvedh at gmail.com>

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
    mlir/test/Dialect/Linalg/data-layout-propagation.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp b/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
index 40085a2368009..ed2efd6fea5f7 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/DataLayoutPropagation.cpp
@@ -1399,6 +1399,10 @@ pushDownExtractSliceOpThroughGenericOp(RewriterBase &rewriter,
       continue;
     }
     AffineMap IndexingMap = genericOp.getMatchingIndexingMap(operand);
+    if (IndexingMap.getNumResults() == 0) {
+      paddedInputs.push_back(operand->get());
+      continue;
+    }
     SmallVector<OpFoldResult> operandLowPads(IndexingMap.getNumResults(),
                                              getAsIndexOpFoldResult(ctx, 0));
     SmallVector<OpFoldResult> operandHighPads(IndexingMap.getNumResults(),

diff  --git a/mlir/test/Dialect/Linalg/data-layout-propagation.mlir b/mlir/test/Dialect/Linalg/data-layout-propagation.mlir
index 0e42027644797..fb16e1e7dcda4 100644
--- a/mlir/test/Dialect/Linalg/data-layout-propagation.mlir
+++ b/mlir/test/Dialect/Linalg/data-layout-propagation.mlir
@@ -1559,4 +1559,21 @@ func.func @nopush_rankreducingextract(%arg0: tensor<128x128x128xf32>, %arg1: ten
 
 // CHECK-LABEL: func.func @nopush_rankreducingextract
 // CHECK:         %[[GENERIC:.+]] = linalg.generic
-// CHECK:         return %[[GENERIC]]   
+// CHECK:         return %[[GENERIC]]
+
+// -----
+
+func.func @push_extract_through_generic_rank0_operand(%arg0: tensor<128x128xf32>, %arg1: tensor<?x?xbf16>, %arg2: index, %arg3 : f32) -> tensor<?x?xbf16> {
+  %extracted_slice = tensor.extract_slice %arg0[%arg2, %arg2] [%arg2, %arg2] [1, 1] : tensor<128x128xf32> to tensor<?x?xf32>
+  %0 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,affine_map<(d0, d1) -> ()> ,affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} ins(%extracted_slice,  %arg3 : tensor<?x?xf32>, f32) outs(%arg1 : tensor<?x?xbf16>) {
+  ^bb0(%in: f32, %in_1 : f32, %out: bf16):
+    %1 = arith.truncf %in : f32 to bf16
+    linalg.yield %1 : bf16
+  } -> tensor<?x?xbf16>
+  return %0 : tensor<?x?xbf16>
+}
+
+// CHECK-LABEL: func.func @push_extract_through_generic_rank0_operand
+// CHECK:         %[[GENERIC:.+]] = linalg.generic
+// CHECK:         %[[EXTRACT:.+]] = tensor.extract_slice %[[GENERIC]]         
+// CHECK:         return %[[EXTRACT]]