[Mlir-commits] [mlir] [mlir][transform] Add linalg.fold_reshape_ops_by_expansion transform op (PR #233732)

[Tuomas Kärnä]

================
@@ -0,0 +1,92 @@
+// RUN: mlir-opt %s \
+// RUN: -transform-preload-library='transform-library-paths=%p/td/fold-reshape-ops-by-expansion.mlir' \
+// RUN: -transform-interpreter=entry-point=fold_reshape_ops_by_expansion \
+// RUN: | FileCheck %s --check-prefixes=CHECK,UNRESTRICTED
+// RUN: mlir-opt %s \
+// RUN: -transform-preload-library='transform-library-paths=%p/td/fold-reshape-ops-by-expansion.mlir' \
+// RUN: -transform-interpreter=entry-point=fold_reshape_ops_by_expansion_unrestricted \
+// RUN: | FileCheck %s --check-prefixes=CHECK,UNRESTRICTED
+// RUN: mlir-opt %s \
+// RUN: -transform-preload-library='transform-library-paths=%p/td/fold-reshape-ops-by-expansion.mlir' \
+// RUN: -transform-interpreter=entry-point=fold_reshape_ops_by_expansion_single_use \
+// RUN: | FileCheck %s --check-prefixes=CHECK,SINGLE
+
+// CHECK-LABEL: func.func @collapse_single_use
+// CHECK: %[[COLLAPSE_SINGLE_FUSED:.*]] = linalg.generic
+// CHECK-SAME: iterator_types = ["parallel", "parallel"]
+// CHECK: %[[COLLAPSE_SINGLE_RESULT:.*]] = tensor.collapse_shape %[[COLLAPSE_SINGLE_FUSED]]
+// CHECK: return %[[COLLAPSE_SINGLE_RESULT]]
+func.func @collapse_single_use(%arg: tensor<2x3xf32>, %init: tensor<6xf32>) -> tensor<6xf32> {
+  %collapsed = tensor.collapse_shape %arg [[0, 1]] : tensor<2x3xf32> into tensor<6xf32>
+  %result = linalg.generic {
+    indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+    iterator_types = ["parallel"]
+  } ins(%collapsed : tensor<6xf32>) outs(%init : tensor<6xf32>) {
+  ^bb0(%in: f32, %out: f32):
+    %neg = arith.negf %in : f32
+    linalg.yield %neg : f32
+  } -> tensor<6xf32>
+  return %result : tensor<6xf32>
+}
+
+// CHECK-LABEL: func.func @collapse_multi_use
+// CHECK: %[[COLLAPSED:.*]] = tensor.collapse_shape
+// UNRESTRICTED: %[[COLLAPSE_MULTI_FUSED:.*]] = linalg.generic
+// UNRESTRICTED-SAME: iterator_types = ["parallel", "parallel"]
+// UNRESTRICTED: %[[COLLAPSE_MULTI_RESULT:.*]] = tensor.collapse_shape %[[COLLAPSE_MULTI_FUSED]]
+// UNRESTRICTED: return %[[COLLAPSED]], %[[COLLAPSE_MULTI_RESULT]]
+// SINGLE: %[[COLLAPSE_MULTI_ORIGINAL:.*]] = linalg.generic
+// SINGLE-SAME: iterator_types = ["parallel"]
+// SINGLE-SAME: ins(%[[COLLAPSED]] : tensor<6xf32>)
+// SINGLE: return %[[COLLAPSED]], %[[COLLAPSE_MULTI_ORIGINAL]]
+func.func @collapse_multi_use(%arg: tensor<2x3xf32>, %init: tensor<6xf32>) -> (tensor<6xf32>, tensor<6xf32>) {
+  %collapsed = tensor.collapse_shape %arg [[0, 1]] : tensor<2x3xf32> into tensor<6xf32>
+  %result = linalg.generic {
+    indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+    iterator_types = ["parallel"]
+  } ins(%collapsed : tensor<6xf32>) outs(%init : tensor<6xf32>) {
+  ^bb0(%in: f32, %out: f32):
+    %neg = arith.negf %in : f32
+    linalg.yield %neg : f32
+  } -> tensor<6xf32>
+  return %collapsed, %result : tensor<6xf32>, tensor<6xf32>
+}
+
+// CHECK-LABEL: func.func @expand_single_use
+// CHECK: %[[EXPAND_SINGLE_FUSED:.*]] = linalg.generic
+// CHECK-SAME: iterator_types = ["parallel", "parallel"]
+// CHECK: return %[[EXPAND_SINGLE_FUSED]]
----------------
tkarna wrote:

updated

https://github.com/llvm/llvm-project/pull/233732