[mlir][Linalg] NFC - Post-commit cleanups

[nicolasvasilache]

Summary

This diff addresses leftovers from https://reviews.llvm.org/D72022.

[merge_guards_bot]

Unit tests: pass. 61175 tests passed, 0 failed and 729 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[rriddle]

mlir/include/mlir/Support/Functional.h

50

This seems hardcoded for a very specific type of iterator element, i.e. one where the result is nullable.

[nicolasvasilache]

mlir/include/mlir/Support/Functional.h

50

I can rename to map_if_non_null.
I was reluctant on forcing a second lambda and turn into a map_if but I could do that too and impl map_if_non_null in terms of that.
Thoughts?

[nicolasvasilache]

Addressing review comments.

[nicolasvasilache]

Update comments.

[merge_guards_bot]

Unit tests: pass. 61237 tests passed, 0 failed and 729 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[merge_guards_bot]

Unit tests: pass. 61237 tests passed, 0 failed and 729 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[rriddle]

mlir/include/mlir/Support/Functional.h

50

I'm not convinced that these methods are worth it. This looks like a case that is already covered by a combination of map+filter ranges.

[rriddle]

mlir/include/mlir/Support/Functional.h

50

FWIW, this also applies to other things within this file; e.g. I don't really see the benefit of 'map' when you can use llvm::map_range instead. At that point, you can opt into vector materialization if you actually need it. This file seems to be reimplementing many of the iterator ranges, albeit using SmallVector instead of using iterator ranges.

[nicolasvasilache]

mlir/include/mlir/Support/Functional.h

50

fair enough, a lot of this predates my realizing that LLVM had STLExtras, I'll rework to cleanup.

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

```mlir
         tensor<?x?xf32>
         -> (tensor<?x?xf32>)
     ```
 
     In this case, the number of return values must match the number of output
     tensor arguments. The semantics is that the `linalg.generic` op
     produces (i.e. allocates and fills) its return values.
     Tensor values must be legalized by a buffer allocation pass before most
-    transformations can be applied. In particular, transformations that create
+    transformations can be applied. In particular, transformations which create
     control flow around linalg.generic operations are not expected to mix with
     tensors because SSA values do not escape naturally. Still, transformations
     and rewrites that take advantage of tensor SSA values are expected to be
     useful and will be added in the near future.
   }];
   let verifier = [{ return ::verify(*this); }];
 }
 

mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h

 //===- LinalgTraits.h - Linalg Traits ---------------------------*- C++ -*-===//
 //
 // Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LINALG_LINALGTRAITS_H_
 #define MLIR_DIALECT_LINALG_LINALGTRAITS_H_
 
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/StandardTypes.h"
+#include "mlir/Support/Functional.h"
 #include "mlir/Support/LLVM.h"
 
 namespace mlir {
 namespace OpTrait {
 namespace linalg {
 
 /// This class provides the API for ops that are known to have a specified
 /// number of inputs, all passed as operands. Use as a trait as follows:
public:
   /// Query whether the op has only MemRef input and outputs.
   bool hasBufferSemantics() {
     return this->getOperation()->getNumResults() == 0 &&
            llvm::all_of(getInputsAndOutputs(),
                         [](Value v) { return v.getType().isa<MemRefType>(); });
   }
   /// Query the subset of input operands that are of ranked tensor type.
   SmallVector<RankedTensorType, 4> getInputTensorTypes() {
-    SmallVector<RankedTensorType, 4> res;
-    for (Type type : getInputs().getTypes())
-      if (auto t = type.template dyn_cast<RankedTensorType>())
-        res.push_back(t);
-    return res;
+    return functional::map_if_non_null(
+        [](Type t) { return t.dyn_cast<RankedTensorType>(); },
+        getInputs().getTypes());
   }
   /// Query the subset of output operands that are of ranked tensor type.
   SmallVector<RankedTensorType, 4> getOutputTensorTypes() {
-    SmallVector<RankedTensorType, 4> res;
-    for (Type type : getOutputs().getTypes())
-      if (auto t = type.template dyn_cast<RankedTensorType>())
-        res.push_back(t);
-    return res;
+    return functional::map_if_non_null(
+        [](Type t) { return t.dyn_cast<RankedTensorType>(); },
+        getOutputs().getTypes());
   }
   /// Return the range over outputs.
   Operation::operand_range getOutputs() {
     auto range = this->getOperation()->getOperands();
     return {range.begin() + nInputs(),
             range.begin() + getNumInputsAndOutputs()};
   }
   /// Return the number of inputs and outputs.

mlir/include/mlir/Support/Functional.h

 //===- Functional.h - Helpers for functional-style Combinators --*- C++ -*-===//
 //
 // Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_SUPPORT_FUNCTIONAL_H_
 #define MLIR_SUPPORT_FUNCTIONAL_H_
 
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Casting.h"
 
 /// This file provides some simple template functional-style sugar to operate
 /// on **value** types. Make sure when using that the stored type is cheap to
 /// copy!
 ///

 
 /// Map with templated container.
 template <typename Fn, typename ContainerType>
 auto map(Fn fun, ContainerType input)
     -> decltype(map(fun, std::begin(input), std::end(input))) {
   return map(fun, std::begin(input), std::end(input));
 }
 
+/// Map with iterators. `fun` also acts like a filter: only mapped elements that
+/// evaluate to true are copied.
+template <typename Fn, typename IterType, typename PredFn>
+auto map_if(Fn fun, IterType begin, IterType end, PredFn predicate)
+    -> SmallVector<typename std::result_of<Fn(decltype(*begin))>::type, 8> {
+  using R = typename std::result_of<Fn(decltype(*begin))>::type;
+  SmallVector<R, 8> res;
+  // auto i works with both pointer types and value types with an operator*.
+  // auto *i only works for pointer types.
+  for (auto i = begin; i != end; ++i) {
+    auto v = fun(*i);
+    if (predicate(v))
+      res.push_back(v);
+  }
+  return res;
+}
+
+/// Map with templated container. `fun` also acts like a filter: only mapped
+/// elements that evaluate to true are copied.
+template <typename Fn, typename ContainerType, typename PredFn>
+auto map_if(Fn fun, ContainerType input, PredFn predicate)
+    -> decltype(map_if(fun, std::begin(input), std::end(input), predicate)) {
+  return map_if(fun, std::begin(input), std::end(input), predicate);
+}
+
+/// Map. `fun` also acts like a filter: only mapped elements that evaluate to
+/// true are copied.
+template <typename Fn, typename IterType>
+auto map_if_non_null(Fn fun, IterType begin, IterType end)
+    -> SmallVector<typename std::result_of<Fn(decltype(*begin))>::type, 8> {
+  return map_if(fun, begin, end,
+                [](typename std::result_of<Fn(decltype(*begin))>::type v) {
+                  return v !=
+                         typename std::result_of<Fn(decltype(*begin))>::type();
+                });
+}
+
+/// Map with templated container. `fun` also acts like a filter: only mapped
+/// elements that evaluate to true are copied.
+template <typename Fn, typename ContainerType>
+auto map_if_non_null(Fn fun, ContainerType input)
+    -> decltype(map_if_non_null(fun, std::begin(input), std::end(input))) {
+  return map_if_non_null(fun, std::begin(input), std::end(input));
+}
+
 /// Zip map with 2 templated container, iterates to the min of the sizes of
 /// the 2 containers.
 /// TODO(ntv): make variadic when needed.
 template <typename Fn, typename ContainerType1, typename ContainerType2>
 auto zipMap(Fn fun, ContainerType1 input1, ContainerType2 input2)
     -> SmallVector<typename std::result_of<Fn(decltype(*input1.begin()),
                                               decltype(*input2.begin()))>::type,
                    8> {