[clang-tools-extra] [libcxx] [clang] [lld] [lldb] [llvm] [mlir] [mlir][transform] Add elementwise criteria to `match.structured.body` (PR #79626)

[via cfe-commits]

https://github.com/srcarroll updated https://github.com/llvm/llvm-project/pull/79626

>From ab475c9ffb7c3562bad4772389e97b82e9f110c0 Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Fri, 26 Jan 2024 11:55:06 -0600
Subject: [PATCH] Add elementwise criteria to match.structured.body

---
 .../Linalg/TransformOps/LinalgMatchOps.td     |  4 +++
 .../Linalg/TransformOps/LinalgMatchOps.cpp    |  9 ++++-
 .../Dialect/Linalg/match-ops-interpreter.mlir | 34 +++++++++++++++++++
 .../Dialect/Linalg/match-ops-invalid.mlir     |  2 +-
 4 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgMatchOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgMatchOps.td
index 162dd05f93030f..dfeb8ae5d5ddbc 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgMatchOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgMatchOps.td
@@ -106,6 +106,9 @@ def MatchStructuredBodyOp : Op<Transform_Dialect, "match.structured.body", [
       * `passthrough`: the body of the structured payload op only forwards
         inputs to the outputs (copy or broadcast).
 
+      * `elementwise`: the body of the structured payload op represents an
+        elementwise operation.
+
       * `contraction`: the body of the structured payload op is a contraction
         of the form `<red>(<elem>(bbarg0, bbarg1), bbarg2)` where `<elem>` and
         `<red>` are binary operations whose names are specified in the attribute
@@ -123,6 +126,7 @@ def MatchStructuredBodyOp : Op<Transform_Dialect, "match.structured.body", [
   let arguments = (ins TransformHandleTypeInterface:$operand_handle,
                        OptionalAttr<I64Attr>:$reduction_position,
                        UnitAttr:$passthrough,
+                       UnitAttr:$elementwise,
                        OptionalAttr<StrArrayAttr>:$contraction);
   let assemblyFormat = "$operand_handle attr-dict `:` type($operand_handle)";
   let extraClassDeclaration = SingleOpMatcher.extraDeclaration;
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgMatchOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgMatchOps.cpp
index 115da4b90e063a..fb18886c16b16d 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgMatchOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgMatchOps.cpp
@@ -11,6 +11,7 @@
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
 #include "mlir/Dialect/Linalg/TransformOps/Syntax.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/Transform/IR/MatchInterfaces.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/Interfaces/FunctionImplementation.h"
@@ -187,6 +188,11 @@ DiagnosedSilenceableFailure transform::MatchStructuredBodyOp::matchOperation(
     }
     return DiagnosedSilenceableFailure::success();
   }
+  if (getElementwise()) {
+    if (!isElementwise(linalgOp))
+      return emitSilenceableError() << "not elementwise";
+    return DiagnosedSilenceableFailure::success();
+  }
   if (std::optional<ArrayAttr> contractionOps = getContraction()) {
     Block &body = linalgOp->getRegion(0).front();
     std::string message;
@@ -209,13 +215,14 @@ DiagnosedSilenceableFailure transform::MatchStructuredBodyOp::matchOperation(
 
 LogicalResult transform::MatchStructuredBodyOp::verify() {
   int64_t numOptions = getReductionPosition().has_value() + getPassthrough() +
-                       getContraction().has_value();
+                       getElementwise() + getContraction().has_value();
 
   if (numOptions > 1) {
     std::string attributeNames;
     llvm::raw_string_ostream os(attributeNames);
     llvm::interleaveComma(ArrayRef<StringAttr>{getReductionPositionAttrName(),
                                                getPassthroughAttrName(),
+                                               getElementwiseAttrName(),
                                                getContractionAttrName()},
                           os);
     return emitOpError() << "only one of {" << os.str() << "} is allowed";
diff --git a/mlir/test/Dialect/Linalg/match-ops-interpreter.mlir b/mlir/test/Dialect/Linalg/match-ops-interpreter.mlir
index a7353a4c38881e..0efe70a7b9ae1e 100644
--- a/mlir/test/Dialect/Linalg/match-ops-interpreter.mlir
+++ b/mlir/test/Dialect/Linalg/match-ops-interpreter.mlir
@@ -180,6 +180,40 @@ module attributes { transform.with_named_sequence } {
 
 // -----
 
+module attributes { transform.with_named_sequence } {
+  transform.named_sequence @print_elementwise(%arg0: !transform.any_op {transform.readonly}) {
+    transform.test_print_remark_at_operand %arg0, "elementwise" : !transform.any_op
+    transform.yield
+  }
+
+  transform.named_sequence @match_structured_body_elementwise(%arg0: !transform.any_op {transform.readonly}) -> !transform.any_op {
+    %0 = transform.match.structured failures(propagate) %arg0 : (!transform.any_op) -> !transform.any_op {
+    ^bb0(%arg1: !transform.any_op):
+      transform.match.structured.body %arg1 { elementwise } : !transform.any_op
+      transform.match.structured.yield %arg1 : !transform.any_op
+    }
+    transform.yield %0 : !transform.any_op
+  }
+
+  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.consumed}) {
+    transform.foreach_match in %arg0
+        @match_structured_body_elementwise -> @print_elementwise
+        : (!transform.any_op) -> !transform.any_op
+    transform.yield
+  }
+
+  func.func @payload(%in1: tensor<2xf32>, %in2: tensor<2xf32>, %out: tensor<2xf32>) -> tensor<2xf32> attributes { transform.target_tag = "start_here" } {
+    %cst0 = arith.constant 0.0 : f32
+    // expected-remark @below {{elementwise}}
+    %fill = linalg.fill ins(%cst0: f32) outs(%out: tensor<2xf32>) -> tensor<2xf32>
+    // expected-remark @below {{elementwise}}
+    %add = linalg.map {arith.addf} ins(%in1, %in2: tensor<2xf32>, tensor<2xf32>) outs(%fill: tensor<2xf32>)
+    return %add : tensor<2xf32>
+  }
+}
+
+// -----
+
 module attributes { transform.with_named_sequence } {
   transform.named_sequence @print_reduction(%arg0: !transform.any_op {transform.readonly}) {
     transform.debug.emit_remark_at %arg0, "reduction" : !transform.any_op
diff --git a/mlir/test/Dialect/Linalg/match-ops-invalid.mlir b/mlir/test/Dialect/Linalg/match-ops-invalid.mlir
index ec99e205090c4c..9ff430a3503606 100644
--- a/mlir/test/Dialect/Linalg/match-ops-invalid.mlir
+++ b/mlir/test/Dialect/Linalg/match-ops-invalid.mlir
@@ -64,7 +64,7 @@ transform.sequence failures(suppress) {
 ^bb0(%arg0: !transform.any_op):
   transform.match.structured %arg0 : !transform.any_op {
   ^bb1(%arg1: !transform.any_op):
-    // expected-error @below {{only one of {"reduction_position", "passthrough", "contraction"} is allowed}}
+    // expected-error @below {{only one of {"reduction_position", "passthrough", "elementwise", "contraction"} is allowed}}
     transform.match.structured.body %arg1 { passthrough, reduction_position = 0 } : !transform.any_op
     transform.match.structured.yield
   }