[Mlir-commits] [mlir] [mlir][spirv] Add conversion pass to rewrite splat constant composite… (PR #148910)

[Jakub Kuderski]

================
@@ -0,0 +1,117 @@
+//===- ConvertToReplicatedConstantCompositePass.cpp -----------------------===//
+//
+// 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 implements a pass to convert a splat composite spirv.Constant and
+// spirv.SpecConstantComposite to spirv.EXT.ConstantCompositeReplicate and
+// spirv.EXT.SpecConstantCompositeReplicate respectively.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
+#include "mlir/Dialect/SPIRV/Transforms/Passes.h"
+#include "mlir/Transforms/WalkPatternRewriteDriver.h"
+
+namespace mlir::spirv {
+#define GEN_PASS_DEF_SPIRVREPLICATEDCONSTANTCOMPOSITEPASS
+#include "mlir/Dialect/SPIRV/Transforms/Passes.h.inc"
+
+namespace {
+
+static std::pair<Attribute, uint32_t>
+getSplatAttrAndNumElements(Attribute valueAttr) {
+  Attribute attr;
+  uint32_t splatCount = 0;
+  if (auto splatAttr = dyn_cast<SplatElementsAttr>(valueAttr)) {
+    return {splatAttr.getSplatValue<Attribute>(), splatAttr.size()};
+  }
+  if (auto arrayAttr = dyn_cast<ArrayAttr>(valueAttr)) {
+    if (llvm::all_equal(arrayAttr)) {
+      attr = arrayAttr[0];
+      splatCount = arrayAttr.size();
+    }
+
+    if (attr) {
+      // Find the inner-most splat value for array of composites
+      std::pair<Attribute, uint32_t> newSplatAttrAndCount =
----------------
kuhar wrote:

Use structured bindings. Also, now that the function is renamed, we can also rename count --> numElements in the variable name too

https://github.com/llvm/llvm-project/pull/148910