[mlir] Add helper to check elementwise-mappable ops with tensors and scalars

[snarang181]

This patch introduces a more general helper for identifying elementwise-mappable operations. The existing utility, isElementwiseMappableOpOnRankedTensors, only accepted operations when all operands were ranked tensors. In practice, many elementwise operations in MLIR allow mixing tensor operands with scalars.
The new helper relaxes the restriction by accepting operands that are either ranked tensors or “scalar-like” types.

[llvmbot]

@llvm/pr-subscribers-mlir-linalg

Author: Samarth Narang (snarang181)

Changes

This patch introduces a more general helper for identifying elementwise-mappable operations. The existing utility, isElementwiseMappableOpOnRankedTensors, only accepted operations when all operands were ranked tensors. In practice, many elementwise operations in MLIR allow mixing tensor operands with scalars.
The new helper relaxes the restriction by accepting operands that are either ranked tensors or “scalar-like” types.

---

Full diff: https://github.com/llvm/llvm-project/pull/154872.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp (+27-3)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp
index c52315333c5b3..87e6ff2fa13c6 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp
@@ -20,13 +20,37 @@ namespace mlir {
 
 using namespace mlir;
 
+// Treats primitive scalars and 0-D tensors as "scalar-like" for broadcasting.
+static inline bool isScalarLike(Type t) {
+  if (llvm::isa<IntegerType, FloatType, IndexType, ComplexType>(t))
+    return true;
+  if (auto rt = dyn_cast<RankedTensorType>(t))
+    return rt.getRank() == 0; // 0-D tensors are scalar-like
+  return false;
+}
+
 static bool isElementwiseMappableOpOnRankedTensors(Operation *op) {
   if (!OpTrait::hasElementwiseMappableTraits(op))
     return false;
 
-  // TODO: The conversion pattern can be made to work for `any_of` here, but
-  // it's more complex as it requires tracking which operands are scalars.
-  return llvm::all_of(op->getOperandTypes(), llvm::IsaPred<RankedTensorType>);
+  auto types = op->getOperandTypes();
+
+  // We want at least one ranked tensor.
+  bool anyRankedTensor = llvm::any_of(
+      types, [](Type type) { return isa<RankedTensorType>(type); });
+
+  // No invalid operands (i.e., every operand is a ranked tensor or
+  // scalar-like).
+  bool noneInvalid = llvm::none_of(types, [](Type t) {
+    // Invalid if neither ranked tensor nor scalar-like.
+    if (llvm::isa<RankedTensorType>(t))
+      return false;
+    if (isScalarLike(t))
+      return false;
+    return true; // Could be a memref, unranked tensor, vector, etc.
+  });
+
+  return anyRankedTensor && noneInvalid;
 }
 
 /// Given `op` assumed `isElementwiseMappableOpOnRankedTensors`, iterate over

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Samarth Narang (snarang181)

Changes

This patch introduces a more general helper for identifying elementwise-mappable operations. The existing utility, isElementwiseMappableOpOnRankedTensors, only accepted operations when all operands were ranked tensors. In practice, many elementwise operations in MLIR allow mixing tensor operands with scalars.
The new helper relaxes the restriction by accepting operands that are either ranked tensors or “scalar-like” types.

---

Full diff: https://github.com/llvm/llvm-project/pull/154872.diff

1 Files Affected:

• (modified) mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp (+27-3)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp
index c52315333c5b3..87e6ff2fa13c6 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp
@@ -20,13 +20,37 @@ namespace mlir {
 
 using namespace mlir;
 
+// Treats primitive scalars and 0-D tensors as "scalar-like" for broadcasting.
+static inline bool isScalarLike(Type t) {
+  if (llvm::isa<IntegerType, FloatType, IndexType, ComplexType>(t))
+    return true;
+  if (auto rt = dyn_cast<RankedTensorType>(t))
+    return rt.getRank() == 0; // 0-D tensors are scalar-like
+  return false;
+}
+
 static bool isElementwiseMappableOpOnRankedTensors(Operation *op) {
   if (!OpTrait::hasElementwiseMappableTraits(op))
     return false;
 
-  // TODO: The conversion pattern can be made to work for `any_of` here, but
-  // it's more complex as it requires tracking which operands are scalars.
-  return llvm::all_of(op->getOperandTypes(), llvm::IsaPred<RankedTensorType>);
+  auto types = op->getOperandTypes();
+
+  // We want at least one ranked tensor.
+  bool anyRankedTensor = llvm::any_of(
+      types, [](Type type) { return isa<RankedTensorType>(type); });
+
+  // No invalid operands (i.e., every operand is a ranked tensor or
+  // scalar-like).
+  bool noneInvalid = llvm::none_of(types, [](Type t) {
+    // Invalid if neither ranked tensor nor scalar-like.
+    if (llvm::isa<RankedTensorType>(t))
+      return false;
+    if (isScalarLike(t))
+      return false;
+    return true; // Could be a memref, unranked tensor, vector, etc.
+  });
+
+  return anyRankedTensor && noneInvalid;
 }
 
 /// Given `op` assumed `isElementwiseMappableOpOnRankedTensors`, iterate over

[rengolin]

Please expand the test coverage. Should be as simple as add new cases to the existing files.

[adam-smnk]

These tests still don't exercise the new logic. The rewrite occurs on arith ops that have tensor only operands.
I'm not sure if there are any upstream ops which are both ElementwiseMappable and accept mixed scalar/shape type operands.

[snarang181]

Would you have any advise on how to proceed?

[adam-smnk]

After a bit of digging, I've found there's this test op that should also work with this conversion pass:

func.func @test(%arg0: f32, %arg1: tensor<?x?xf32>) {
  %0 = "test.elementwise_mappable"(%arg0, %arg1) : (f32, tensor<?x?xf32>) -> tensor<?x?xf32>
  return
}