From b2bc3ef80b3c6f14acc5ee2d456e7f6d4f851b65 Mon Sep 17 00:00:00 2001
From: Siddhesh Deodhar <siddheshdeodhar95@gmail.com>
Date: Sat, 2 Nov 2024 19:17:05 -0400
Subject: [PATCH] Fix assertion when tiling linalg.generic

Fix assertion during tiling due to unsupported linalg.generic containing negative coefficients
---
 mlir/include/mlir/IR/AffineMap.h              |  3 ++
 .../Linalg/Transforms/TilingInterfaceImpl.cpp |  8 +++++
 mlir/lib/IR/AffineMap.cpp                     | 23 +++++++++++++++
 mlir/test/Dialect/Linalg/tile-invalid.mlir    | 29 +++++++++++++++++++
 4 files changed, 63 insertions(+)
 create mode 100644 mlir/test/Dialect/Linalg/tile-invalid.mlir

diff --git a/mlir/include/mlir/IR/AffineMap.h b/mlir/include/mlir/IR/AffineMap.h
index e30950bbf292d..c2bf0d6b91fed 100644
--- a/mlir/include/mlir/IR/AffineMap.h
+++ b/mlir/include/mlir/IR/AffineMap.h
@@ -382,6 +382,9 @@ class AffineMap {
   /// Returns true if the AffineMap represents a symbol-less permutation map.
   bool isPermutation() const;
 
+  /// Returns true if the AffineMap contains non-positive coefficients
+  bool isNonPositiveCoefficients() const;
+
   /// Returns the map consisting of the `resultPos` subset.
   AffineMap getSubMap(ArrayRef<unsigned> resultPos) const;
 
diff --git a/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp b/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
index f86715a94b268..b1e1bf9f721fe 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
@@ -119,6 +119,14 @@ struct LinalgOpTilingInterface
     // specified could lead to out of bounds accesses.
     Location loc = op->getLoc();
     LinalgOp linalgOp = cast<LinalgOp>(op);
+    SmallVector<AffineMap> indexingMaps = linalgOp.getIndexingMapsArray();
+    if (llvm::any_of(linalgOp.getIndexingMapsArray(), [](AffineMap m) {
+          return m.isNonPositiveCoefficients();
+        })) {
+      return linalgOp.emitOpError(
+          "tiling not supported because op has a non positive coefficient");
+    }
+
     SmallVector<Value> valuesToTile = linalgOp->getOperands();
     SmallVector<Value> tiledOperands = makeTiledShapes(
         b, loc, linalgOp, valuesToTile, offsets, sizes, {}, true);
diff --git a/mlir/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp
index ea3c0723b0775..63a2506552b28 100644
--- a/mlir/lib/IR/AffineMap.cpp
+++ b/mlir/lib/IR/AffineMap.cpp
@@ -9,6 +9,7 @@
 #include "mlir/IR/AffineMap.h"
 #include "AffineMapDetail.h"
 #include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
@@ -651,6 +652,28 @@ bool AffineMap::isPermutation() const {
   return isProjectedPermutation();
 }
 
+struct CheckCoefficients : public AffineExprVisitor<CheckCoefficients> {
+  CheckCoefficients() {}
+
+  void visitAffineBinaryOpExpr(AffineBinaryOpExpr expr) {
+    visit(expr.getLHS());
+    visit(expr.getRHS());
+    if (expr.getKind() == mlir::AffineExprKind::Mul)
+      isNonPositiveCoefficients |=
+          cast<AffineConstantExpr>(expr.getRHS()).getValue() <= 0;
+  }
+  bool isNonPositiveCoefficients = false;
+};
+
+bool AffineMap::isNonPositiveCoefficients() const {
+
+  return llvm::any_of(getResults(), [](AffineExpr e) {
+    CheckCoefficients t;
+    t.visit(e);
+    return t.isNonPositiveCoefficients;
+  });
+}
+
 AffineMap AffineMap::getSubMap(ArrayRef<unsigned> resultPos) const {
   SmallVector<AffineExpr, 4> exprs;
   exprs.reserve(resultPos.size());
diff --git a/mlir/test/Dialect/Linalg/tile-invalid.mlir b/mlir/test/Dialect/Linalg/tile-invalid.mlir
new file mode 100644
index 0000000000000..1684b1de72c26
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/tile-invalid.mlir
@@ -0,0 +1,29 @@
+// RUN: mlir-opt %s -transform-interpreter -split-input-file -verify-diagnostics
+
+#map1 = affine_map<(d0) -> (d0)>
+#map2 = affine_map<(d0) -> (-d0 + 7)>
+
+func.func @test(%arg0: tensor<8xi8>, %arg1: tensor<8xi8>) -> tensor<8xi8> {
+   // expected-error @below{{tiling not supported because op has a non positive coefficient}}
+   // expected-error @below{{op faild to tile operation}}
+   // expected-error @below{{op failed to generate tiling loops}}
+  %0 = linalg.generic {
+    indexing_maps = [#map2, #map1],
+    iterator_types = ["parallel"]}
+  ins(%arg0 : tensor<8xi8>)
+  outs(%arg1 : tensor<8xi8>) {
+  ^bb0(%in: i8, %out: i8):
+    linalg.yield %in : i8
+  } -> tensor<8xi8>
+  return %0 : tensor<8xi8>
+}
+
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+    %1:2 = transform.structured.tile_using_for
+      %0 tile_sizes [2] : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+     transform.yield
+  }
+}