[llvm-branch-commits] [mlir] f5d8eb0 - [mlir][Linalg] NFC - getAssumedNonShapedOperands now returns OperandRange

[Nicolas Vasilache via llvm-branch-commits]

Author: Nicolas Vasilache
Date: 2021-01-20T19:23:26Z
New Revision: f5d8eb085af97c6d873edf3ca16d85b8a97c67e6

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

LOG: [mlir][Linalg] NFC - getAssumedNonShapedOperands now returns OperandRange

Also adds a isInput interface method.

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td
index 85133604cda0..b8009a818aa0 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td
@@ -609,6 +609,22 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
         return payloadUsesValueFromOpOperand(&getOutputOpOperands()[index]);
       }]
     >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Return true if `opOperand` is an input tensor.
+      }],
+      /*retTy=*/"bool",
+      /*methodName=*/"isInputTensor",
+      /*args=*/(ins "OpOperand *":$opOperand),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        if (!opOperand->get().getType().template isa<RankedTensorType>())
+          return false;
+        if (opOperand->getOperandNumber() < $_op.getNumInputs())
+          return true;
+        return false;
+      }]
+    >,
     InterfaceMethod<
       /*desc=*/[{
         Return true if `opOperand` is an init tensor. This is true when it is
@@ -1063,18 +1079,13 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
     /// init_tensors operands. Asserts that these operands are value types to
     /// allow transformations like tiling to just use the values when cloning
     /// `linalgOp`.
-    SmallVector<Value, 4> getAssumedNonShapedOperands() {
-      unsigned numShapedOperands = getNumShapedOperands();
-      unsigned nExtraOperands =
-        getOperation()->getNumOperands() - numShapedOperands;
-      SmallVector<Value, 4> res;
-      res.reserve(nExtraOperands);
-      for (unsigned i = 0; i < nExtraOperands; ++i) {
-        res.push_back(getOperation()->getOperand(numShapedOperands + i));
-        assert((res.back().getType().isSignlessIntOrIndexOrFloat()
-                || res.back().getType().template isa<VectorType>()) &&
-               "expected scalar or vector type");
-      }
+    Operation::operand_range getAssumedNonShapedOperands() {
+      Operation::operand_range res{
+        getOperation()->getOperands().begin() + getNumShapedOperands(),
+        getOperation()->getOperands().end()};
+      for (Type t : TypeRange{res})
+        assert((t.isSignlessIntOrIndexOrFloat() || t.template isa<VectorType>())
+               &&"expected scalar or vector type");
       return res;
     }