[mlir][spirv][NFC] Clean up SPIR-V canonicalization

[kuhar]

Summary

Follow best practices. Use llvm helper functions for readability.

[antiagainst]

Nice!

mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp

 //===- SPIRVCanonicalization.cpp - MLIR SPIR-V canonicalization patterns --===//
 //
 // Part of the LLVM 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines the folders and canonicalization patterns for SPIR-V ops.
 //
 //===----------------------------------------------------------------------===//
 
-#include <utility>
 #include <optional>
+#include <utility>
 
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 
 #include "mlir/Dialect/CommonFolders.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVectorExtras.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Common utility functions
 //===----------------------------------------------------------------------===//
 
 /// Returns the boolean value under the hood if the given `boolAttr` is a scalar

 //===----------------------------------------------------------------------===//
 // spirv.AccessChainOp
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 /// Combines chained `spirv::AccessChainOp` operations into one
 /// `spirv::AccessChainOp` operation.
-struct CombineChainedAccessChain
-    : public OpRewritePattern<spirv::AccessChainOp> {
-  using OpRewritePattern<spirv::AccessChainOp>::OpRewritePattern;
+struct CombineChainedAccessChain final
+    : OpRewritePattern<spirv::AccessChainOp> {
+  using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(spirv::AccessChainOp accessChainOp,
                                 PatternRewriter &rewriter) const override {
-    auto parentAccessChainOp = dyn_cast_or_null<spirv::AccessChainOp>(
-        accessChainOp.getBasePtr().getDefiningOp());
+    auto parentAccessChainOp =
+        accessChainOp.getBasePtr().getDefiningOp<spirv::AccessChainOp>();
 
     if (!parentAccessChainOp) {
       return failure();
     }
 
     // Combine indices.
     SmallVector<Value, 4> indices(parentAccessChainOp.getIndices());
-    indices.append(accessChainOp.getIndices().begin(),
-                   accessChainOp.getIndices().end());
+    llvm::append_range(indices, accessChainOp.getIndices());
 
     rewriter.replaceOpWithNewOp<spirv::AccessChainOp>(
         accessChainOp, parentAccessChainOp.getBasePtr(), indices);
 
     return success();
   }
 };
 } // namespace
while (auto insertOp =
     compositeOp = insertOp.getComposite();
   }
 
   if (auto constructOp =
           compositeOp.getDefiningOp<spirv::CompositeConstructOp>()) {
     auto type = llvm::cast<spirv::CompositeType>(constructOp.getType());
     if (getIndices().size() == 1 &&
         constructOp.getConstituents().size() == type.getNumElements()) {
-      auto i = getIndices().begin()->cast<IntegerAttr>();
-      if (static_cast<size_t>(i.getValue().getSExtValue()) <
-          constructOp.getConstituents().size())
+      auto i = llvm::cast<IntegerAttr>(*getIndices().begin());
+      if (i.getValue().getSExtValue() <
+          static_cast<int64_t>(constructOp.getConstituents().size()))
         return constructOp.getConstituents()[i.getValue().getSExtValue()];
     }
   }
 
-  auto indexVector =
-      llvm::to_vector<8>(llvm::map_range(getIndices(), [](Attribute attr) {
-        return static_cast<unsigned>(llvm::cast<IntegerAttr>(attr).getInt());
-      }));
+  auto indexVector = llvm::map_to_vector(getIndices(), [](Attribute attr) {
+  });
   return extractCompositeElement(adaptor.getComposite(), indexVector);
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.Constant
 //===----------------------------------------------------------------------===//
 
 OpFoldResult spirv::ConstantOp::fold(FoldAdaptor /*adaptor*/) {

 }
 
 //===----------------------------------------------------------------------===//
 // spirv.LogicalOr
 //===----------------------------------------------------------------------===//
 
 OpFoldResult spirv::LogicalOrOp::fold(FoldAdaptor adaptor) {
   if (auto rhs = getScalarOrSplatBoolAttr(adaptor.getOperand2())) {
-    if (*rhs)
+    if (*rhs) {
       // x || true = true
       return adaptor.getOperand2();
+    }
 
+    if (!*rhs) {
       // x || false = x
-    if (!*rhs)
       return getOperand1();
     }
+  }
 
   return Attribute();
 }
 
 //===----------------------------------------------------------------------===//
 // spirv.mlir.selection
 //===----------------------------------------------------------------------===//
 

 //
 //                             ...
 //                            \   /
 //                              v
 //                       +-------------+
 //                       | merge block |
 //                       +-------------+
 //
-struct ConvertSelectionOpToSelect
-    : public OpRewritePattern<spirv::SelectionOp> {
-  using OpRewritePattern<spirv::SelectionOp>::OpRewritePattern;
+struct ConvertSelectionOpToSelect final : OpRewritePattern<spirv::SelectionOp> {
+  using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(spirv::SelectionOp selectionOp,
                                 PatternRewriter &rewriter) const override {
-    auto *op = selectionOp.getOperation();
-    auto &body = op->getRegion(0);
+    Operation *op = selectionOp.getOperation();
+    Region &body = op->getRegion(0);
     // Verifier allows an empty region for `spirv.mlir.selection`.
     if (body.empty()) {
       return failure();
     }
 
     // Check that region consists of 4 blocks:
     // header block, `true` block, `false` block and merge block.
-    if (std::distance(body.begin(), body.end()) != 4) {
+    if (llvm::range_size(body) != 4) {
       return failure();
     }
 
-    auto *headerBlock = selectionOp.getHeaderBlock();
+    Block *headerBlock = selectionOp.getHeaderBlock();
     if (!onlyContainsBranchConditionalOp(headerBlock)) {
       return failure();
     }
 
     auto brConditionalOp =
         cast<spirv::BranchConditionalOp>(headerBlock->front());
 
-    auto *trueBlock = brConditionalOp.getSuccessor(0);
-    auto *falseBlock = brConditionalOp.getSuccessor(1);
-    auto *mergeBlock = selectionOp.getMergeBlock();
+    Block *trueBlock = brConditionalOp.getSuccessor(0);
+    Block *falseBlock = brConditionalOp.getSuccessor(1);
+    Block *mergeBlock = selectionOp.getMergeBlock();
 
     if (failed(canCanonicalizeSelection(trueBlock, falseBlock, mergeBlock)))
       return failure();
 
-    auto trueValue = getSrcValue(trueBlock);
-    auto falseValue = getSrcValue(falseBlock);
-    auto ptrValue = getDstPtr(trueBlock);
+    Value trueValue = getSrcValue(trueBlock);
+    Value falseValue = getSrcValue(falseBlock);
+    Value ptrValue = getDstPtr(trueBlock);
     auto storeOpAttributes =
         cast<spirv::StoreOp>(trueBlock->front())->getAttrs();
 
     auto selectOp = rewriter.create<spirv::SelectOp>(
         selectionOp.getLoc(), trueValue.getType(),
         brConditionalOp.getCondition(), trueValue, falseValue);
     rewriter.create<spirv::StoreOp>(selectOp.getLoc(), ptrValue,
                                     selectOp.getResult(), storeOpAttributes);
private:
   //    is a `spirv.Store` and the last operation is a `spirv.Branch`.
   // 2. Each `spirv.Store` uses the same pointer and the same memory attributes.
   // 3. A control flow goes into the given merge block from the given
   //    conditional blocks.
   LogicalResult canCanonicalizeSelection(Block *trueBlock, Block *falseBlock,
                                          Block *mergeBlock) const;
 
   bool onlyContainsBranchConditionalOp(Block *block) const {
-    return std::next(block->begin()) == block->end() &&
+    return llvm::hasSingleElement(*block) &&
            isa<spirv::BranchConditionalOp>(block->front());
   }
 
   bool isSameAttrList(spirv::StoreOp lhs, spirv::StoreOp rhs) const {
     return lhs->getDiscardableAttrDictionary() ==
                rhs->getDiscardableAttrDictionary() &&
            lhs.getProperties() == rhs.getProperties();
   }
Value getDstPtr(Block *block) const {
     auto storeOp = cast<spirv::StoreOp>(block->front());
     return storeOp.getPtr();
   }
 };
 
 LogicalResult ConvertSelectionOpToSelect::canCanonicalizeSelection(
     Block *trueBlock, Block *falseBlock, Block *mergeBlock) const {
   // Each block must consists of 2 operations.
-  if ((std::distance(trueBlock->begin(), trueBlock->end()) != 2) ||
-      (std::distance(falseBlock->begin(), falseBlock->end()) != 2)) {
+  if (llvm::range_size(*trueBlock) != 2 || llvm::range_size(*falseBlock) != 2) {
     return failure();
   }
 
   auto trueBrStoreOp = dyn_cast<spirv::StoreOp>(trueBlock->front());
   auto trueBrBranchOp =
       dyn_cast<spirv::BranchOp>(*std::next(trueBlock->begin()));
   auto falseBrStoreOp = dyn_cast<spirv::StoreOp>(falseBlock->front());
   auto falseBrBranchOp =