Addressing further code review comments

Signed-off-by: Mohammadreza Ameri Mahabadian <[email protected]>

Author: mahabadm

mlir/include/mlir/Dialect/SPIRV/Transforms/Passes.td

@@ -78,7 +78,7 @@ def SPIRVWebGPUPreparePass : Pass<"spirv-webgpu-prepare", "spirv::ModuleOp"> {
 }
 
 def SPIRVReplicatedConstantCompositePass
-    : Pass<"spirv-convert-to-replicated-const-composite", "spirv::ModuleOp"> {
+    : Pass<"spirv-promote-to-replicated-constants", "spirv::ModuleOp"> {
   let summary = "Convert splat composite constants and spec constants to "
                 "corresponding replicated constant composite ops defined by "
                 "SPV_EXT_replicated_composites";

mlir/lib/Dialect/SPIRV/Transforms/ConvertToReplicatedConstantCompositePass.cpp

@@ -23,28 +23,22 @@ namespace mlir::spirv {
 namespace {
 
 static std::pair<Attribute, uint32_t>
-getSplatAttributeAndCount(Attribute valueAttr) {
+getSplatAttrAndNumElements(Attribute valueAttr) {
   Attribute attr;
   uint32_t splatCount = 0;
-  if (auto denseAttr = dyn_cast<DenseElementsAttr>(valueAttr)) {
-    if (denseAttr.isSplat()) {
-      attr = denseAttr.getSplatValue<Attribute>();
-      splatCount = denseAttr.size();
-    }
-  } else if (auto arrayAttr = dyn_cast<ArrayAttr>(valueAttr)) {
-    if (std::adjacent_find(arrayAttr.begin(), arrayAttr.end(),
-                           std::not_equal_to<>()) == arrayAttr.end()) {
+  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) {
-    auto typedAttr = dyn_cast<TypedAttr>(attr);
-    if ((typedAttr && isa<spirv::CompositeType>(typedAttr.getType())) ||
-        isa<ArrayAttr>(attr)) {
+    if (attr) {
+      // Find the inner-most splat value for array of composites
       std::pair<Attribute, uint32_t> newSplatAttrAndCount =
-          getSplatAttributeAndCount(attr);
+          getSplatAttrAndNumElements(attr);
       if (newSplatAttrAndCount.first) {
         return newSplatAttrAndCount;
       }
@@ -63,7 +57,7 @@ struct ConstantOpConversion final : OpRewritePattern<spirv::ConstantOp> {
     if (!compositeType)
       return rewriter.notifyMatchFailure(op, "not a composite constant");
 
-    auto [splatAttr, splatCount] = getSplatAttributeAndCount(op.getValue());
+    auto [splatAttr, splatCount] = getSplatAttrAndNumElements(op.getValue());
     if (!splatAttr)
       return rewriter.notifyMatchFailure(op, "composite is not splat");
 
@@ -73,7 +67,6 @@ struct ConstantOpConversion final : OpRewritePattern<spirv::ConstantOp> {
 
     rewriter.replaceOpWithNewOp<spirv::EXTConstantCompositeReplicateOp>(
         op, op.getType(), splatAttr);
-
     return success();
   }
 };
@@ -93,8 +86,7 @@ struct SpecConstantCompositeOpConversion final
       return rewriter.notifyMatchFailure(op,
                                          "composite has only one consituent");
 
-    if (!(std::adjacent_find(constituents.begin(), constituents.end(),
-                             std::not_equal_to<>()) == constituents.end()))
+    if (!(llvm::all_equal(constituents)))
       return rewriter.notifyMatchFailure(op, "composite is not splat");
 
     auto splatConstituent = dyn_cast<FlatSymbolRefAttr>(constituents[0]);

mlir/test/Dialect/SPIRV/Transforms/replicated-const-composites.mlir

@@ -1,6 +1,6 @@
-// RUN: mlir-opt --spirv-convert-to-replicated-const-composite --split-input-file --verify-diagnostics %s | FileCheck %s
+// RUN: mlir-opt --spirv-promote-to-replicated-constants --split-input-file %s | FileCheck %s
 
-spirv.module Logical GLSL450 {
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, ReplicatedCompositesEXT], [SPV_EXT_replicated_composites]> {
   spirv.func @splat_vector_of_i32() -> (vector<3xi32>) "None" {
     // CHECK: {{%.*}} = spirv.EXT.ConstantCompositeReplicate [2 : i32] : vector<3xi32>
     %0 = spirv.Constant dense<2> : vector<3xi32>
@@ -132,21 +132,13 @@ spirv.module Logical GLSL450 {
     %0 = spirv.Constant dense<1.0> : !spirv.arm.tensor<1xf32>
     spirv.ReturnValue %0 : !spirv.arm.tensor<1xf32>
   }
-}
-
-// -----
 
-spirv.module Logical GLSL450 {
   spirv.func @non_splat_vector_of_f32() -> (vector<3xf32>) "None" {
     // CHECK-NOT: spirv.EXT.ConstantCompositeReplicate
     %0 = spirv.Constant dense<[0.0, 1.0, 2.0]> : vector<3xf32>
     spirv.ReturnValue %0 : vector<3xf32>
   }
-}
 
-// -----
-
-spirv.module Logical GLSL450 {
   spirv.func @non_splat_array_of_vectors_of_f32() -> (!spirv.array<2xvector<2xf32>>) "None" {
     // CHECK-NOT: spirv.EXT.ConstantCompositeReplicate
     %0 = spirv.Constant [dense<[1.0, 2.0]> : vector<2xf32>, dense<[1.0, 3.0]> : vector<2xf32>] : !spirv.array<2 x vector<2xf32>>
@@ -156,7 +148,7 @@ spirv.module Logical GLSL450 {
 
 // -----
 
-spirv.module Logical GLSL450 {
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, ReplicatedCompositesEXT], [SPV_EXT_replicated_composites]> {
 
   spirv.SpecConstant @sc_i32_1 = 1 : i32
 
@@ -216,4 +208,4 @@ spirv.module Logical GLSL450 {
 
   // CHECK-NOT: spirv.EXT.SpecConstantCompositeReplicate
   spirv.SpecConstantComposite @scc_struct_of_i32_and_f32 (@sc_i32_1, @sc_i32_1, @sc_f32_1) : !spirv.struct<(i32, i32, f32)>
-}
+}