[Mlir-commits] [mlir] 428bc6f - [mlir][Linalg] Fix constant detection in linalg.pad_tensor vectorization.

[Nicolas Vasilache]

Author: Nicolas Vasilache
Date: 2021-02-14T15:53:39Z
New Revision: 428bc6feed088accf549296b77ecf544d54ff1c9

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

LOG: [mlir][Linalg] Fix constant detection in linalg.pad_tensor vectorization.

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
    mlir/test/Dialect/Linalg/vectorization.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index bfd288464c68..e350d8931753 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -429,8 +429,9 @@ LogicalResult PadTensorOpVectorizationPattern::matchAndRewrite(
     if (Attribute attr = ofr.dyn_cast<Attribute>())
       return attr.cast<IntegerAttr>().getInt() != 0;
     Value v = ofr.get<Value>();
-    if (auto constOp = v.getDefiningOp<ConstantIntOp>())
-      return constOp.getValue() != 0;
+    if (auto constOp = v.getDefiningOp<ConstantOp>())
+      if (auto intAttr = constOp.getValue().dyn_cast<IntegerAttr>())
+        return intAttr.getValue().getSExtValue() != 0;
     return true;
   };
 

diff  --git a/mlir/test/Dialect/Linalg/vectorization.mlir b/mlir/test/Dialect/Linalg/vectorization.mlir
index bb532b2a550c..13d2e181e2ea 100644
--- a/mlir/test/Dialect/Linalg/vectorization.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization.mlir
@@ -402,7 +402,8 @@ func @pad_static(%arg0: tensor<?x?x?xf32>, %pad_value: f32) -> tensor<2x3x4xf32>
   //      CHECK: %[[INIT:.*]] = linalg.init_tensor [2, 3, 4] : tensor<2x3x4xf32>
   //      CHECK: %[[WRITTEN:.*]] = vector.transfer_write %[[READ]], %[[INIT]][%[[C0]], %[[C0]], %[[C0]]]
   // CHECK-SAME:   {masked = [false, false, false]} : vector<2x3x4xf32>, tensor<2x3x4xf32>
-  %0 = linalg.pad_tensor %arg0 low[0, 0, 0] high[0, 0, 0] {
+  %c0 = constant 0 : index
+  %0 = linalg.pad_tensor %arg0 low[0, %c0, 0] high[0, 0, %c0] {
     ^bb0(%arg1: index, %arg2: index, %arg3: index):
       linalg.yield %pad_value : f32
     } : tensor<?x?x?xf32> to tensor<2x3x4xf32>