Skip to content

[mlir][spirv] Add conversion pass to rewrite splat constant composite… #148910

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 12 commits into from
Jul 18, 2025
Merged

[mlir][spirv] Add conversion pass to rewrite splat constant composite… #148910

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
936435e
[mlir][spirv] Add conversion pass to rewrite splat constant composite…
mahabadm Jul 2, 2025
14b70ce
Minor bug fix and revision
mahabadm Jul 16, 2025
bee1e45
Merge branch 'main' into constant_composite_replicates_conversion_pass
mahabadm Jul 16, 2025
51ce7ac
Addressing code review comments
mahabadm Jul 16, 2025
d065a5a
Addressing further code review comments
mahabadm Jul 17, 2025
b718e34
Addressing further code review comments
mahabadm Jul 17, 2025
8234e68
Minor typo fix
mahabadm Jul 17, 2025
568bccb
Addressing further code review comments
mahabadm Jul 17, 2025
ae1eb18
Addressing further code review comments
mahabadm Jul 17, 2025
6ea73e3
Slight change of logic for value type detection
mahabadm Jul 18, 2025
a368eb7
Merge branch 'main' into constant_composite_replicates_conversion_pass
mahabadm Jul 18, 2025
1584f58
Addressing further code review comments
mahabadm Jul 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion mlir/include/mlir/Dialect/SPIRV/Transforms/Passes.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@ def SPIRVWebGPUPreparePass : Pass<"spirv-webgpu-prepare", "spirv::ModuleOp"> {
}

def SPIRVReplicatedConstantCompositePass
: Pass<"spirv-replicated-const-composite", "spirv::ModuleOp"> {
: Pass<"spirv-convert-to-replicated-const-composite", "spirv::ModuleOp"> {
let summary = "Convert splat composite constants and spec constants to"
"corresponding replicated constant composite ops defined by"
"SPV_EXT_replicated_composites";
Expand Down
10 changes: 6 additions & 4 deletions mlir/lib/Dialect/SPIRV/Transforms/ConversionToReplicatedConstantCompositePass.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
//===- ConversionToReplicatedConstantCompositePass.cpp
//---------------------------===//
//===- ConversionToReplicatedConstantCompositePass.cpp --------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
Expand Down Expand Up @@ -83,7 +82,7 @@ class ConstantOpConversion : public OpRewritePattern<spirv::ConstantOp> {

if (splatCount == 1)
return rewriter.notifyMatchFailure(op,
"composite has only one consituent");
"composite has only one constituent");

rewriter.replaceOpWithNewOp<spirv::EXTConstantCompositeReplicateOp>(
op, op.getType(), splatAttr);
Expand All @@ -102,7 +101,7 @@ class SpecConstantCompositeOpConversion
if (!compositeType)
return rewriter.notifyMatchFailure(op, "not a composite constant");

auto constituents = op.getConstituents();
ArrayAttr constituents = op.getConstituents();
if (constituents.size() == 1)
return rewriter.notifyMatchFailure(op,
"composite has only one consituent");
Expand All @@ -113,6 +112,9 @@ class SpecConstantCompositeOpConversion

auto splatConstituent =
dyn_cast<FlatSymbolRefAttr>(op.getConstituents()[0]);
if (!splatConstituent)
return rewriter.notifyMatchFailure(
op, "expected flat symbol reference for splat constituent");

rewriter.replaceOpWithNewOp<spirv::EXTSpecConstantCompositeReplicateOp>(
op, TypeAttr::get(op.getType()), op.getSymNameAttr(), splatConstituent);
Expand Down
145 changes: 128 additions & 17 deletions mlir/test/Dialect/SPIRV/Transforms/replicated-const-composites.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
// RUN: mlir-opt -spirv-replicated-const-composite -split-input-file -verify-diagnostics %s -o - | FileCheck %s
// RUN: mlir-opt --spirv-convert-to-replicated-const-composite --split-input-file --verify-diagnostics %s -o - | FileCheck %s

spirv.module Logical GLSL450 {
spirv.func @splat_vector_of_i32() -> (vector<3xi32>) "None" {
%0 = spirv.Constant dense<2> : vector<3xi32>
// CHECK: {{%.*}} = spirv.EXT.ConstantCompositeReplicate [2 : i32] : vector<3xi32>
%0 = spirv.Constant dense<2> : vector<3xi32>
spirv.ReturnValue %0 : vector<3xi32>
}
}
Expand All @@ -12,8 +12,8 @@ spirv.module Logical GLSL450 {

spirv.module Logical GLSL450 {
spirv.func @splat_array_of_i32() -> (!spirv.array<3 x i32>) "None" {
%0 = spirv.Constant [1 : i32, 1 : i32, 1 : i32] : !spirv.array<3 x i32>
// CHECK: {{%.*}} = spirv.EXT.ConstantCompositeReplicate [1 : i32] : !spirv.array<3 x i32>
%0 = spirv.Constant [1 : i32, 1 : i32, 1 : i32] : !spirv.array<3 x i32>
spirv.ReturnValue %0 : !spirv.array<3 x i32>
}
}
Expand All @@ -22,7 +22,7 @@ spirv.module Logical GLSL450 {

spirv.module Logical GLSL450 {
spirv.func @splat_array_of_splat_array_of_i32() -> (!spirv.array<2 x !spirv.array<3 x i32>>) "None" {
// CHECK: spirv.EXT.ConstantCompositeReplicate [3 : i32] : !spirv.array<2 x !spirv.array<3 x i32>>
// CHECK: {{%.*}} = spirv.EXT.ConstantCompositeReplicate [3 : i32] : !spirv.array<2 x !spirv.array<3 x i32>>
%0 = spirv.Constant [[3 : i32, 3 : i32, 3 : i32], [3 : i32, 3 : i32, 3 : i32]] : !spirv.array<2 x !spirv.array<3 x i32>>
spirv.ReturnValue %0 : !spirv.array<2 x !spirv.array<3 x i32>>
}
Expand All @@ -32,7 +32,7 @@ spirv.module Logical GLSL450 {

spirv.module Logical GLSL450 {
spirv.func @splat_array_of_non_splat_array_of_i32() -> (!spirv.array<2 x !spirv.array<3 x i32>>) "None" {
// CHECK: spirv.EXT.ConstantCompositeReplicate {{\[}}[1 : i32, 2 : i32, 3 : i32]] : !spirv.array<2 x !spirv.array<3 x i32>>
// CHECK: {{%.*}} = spirv.EXT.ConstantCompositeReplicate {{\[}}[1 : i32, 2 : i32, 3 : i32]] : !spirv.array<2 x !spirv.array<3 x i32>>
%0 = spirv.Constant [[1 : i32, 2 : i32, 3 : i32], [1 : i32, 2 : i32, 3 : i32]] : !spirv.array<2 x !spirv.array<3 x i32>>
spirv.ReturnValue %0 : !spirv.array<2 x !spirv.array<3 x i32>>
}
Expand All @@ -42,8 +42,8 @@ spirv.module Logical GLSL450 {

spirv.module Logical GLSL450 {
spirv.func @splat_array_of_vectors_of_i32() -> (!spirv.array<2xvector<2xi32>>) "None" {
%0 = spirv.Constant [dense<[1, 2]> : vector<2xi32>, dense<[1, 2]> : vector<2xi32>] : !spirv.array<2 x vector<2xi32>>
// CHECK: {{%.*}} = spirv.EXT.ConstantCompositeReplicate [dense<[1, 2]> : vector<2xi32>] : !spirv.array<2 x vector<2xi32>>
%0 = spirv.Constant [dense<[1, 2]> : vector<2xi32>, dense<[1, 2]> : vector<2xi32>] : !spirv.array<2 x vector<2xi32>>
spirv.ReturnValue %0 : !spirv.array<2 x vector<2xi32>>
}
}
Expand All @@ -52,8 +52,8 @@ spirv.module Logical GLSL450 {

spirv.module Logical GLSL450 {
spirv.func @splat_array_of_splat_vectors_of_i32() -> (!spirv.array<2 x vector<2xi32>>) "None" {
%0 = spirv.Constant [dense<2> : vector<2xi32>, dense<2> : vector<2xi32>] : !spirv.array<2 x vector<2xi32>>
// CHECK: {{%.*}} = spirv.EXT.ConstantCompositeReplicate [2 : i32] : !spirv.array<2 x vector<2xi32>>
%0 = spirv.Constant [dense<2> : vector<2xi32>, dense<2> : vector<2xi32>] : !spirv.array<2 x vector<2xi32>>
spirv.ReturnValue %0 : !spirv.array<2 x vector<2xi32>>
}
}
Expand All @@ -62,7 +62,7 @@ spirv.module Logical GLSL450 {

spirv.module Logical GLSL450 {
spirv.func @splat_tensor_of_i32() -> (!spirv.array<2 x !spirv.array<3 x i32>>) "None" {
// CHECK: spirv.EXT.ConstantCompositeReplicate [3 : i32] : !spirv.array<2 x !spirv.array<3 x i32>>
// CHECK: {{%.*}} = spirv.EXT.ConstantCompositeReplicate [3 : i32] : !spirv.array<2 x !spirv.array<3 x i32>>
%0 = spirv.Constant dense<3> : tensor<2x3xi32> : !spirv.array<2 x !spirv.array<3 x i32>>
spirv.ReturnValue %0 : !spirv.array<2 x !spirv.array<3 x i32>>
}
Expand All @@ -72,7 +72,7 @@ spirv.module Logical GLSL450 {

spirv.module Logical GLSL450 {
spirv.func @splat_arm_tensor_of_i32() -> (!spirv.arm.tensor<2x3xi32>) "None" {
// CHECK: spirv.EXT.ConstantCompositeReplicate [2 : i32] : !spirv.arm.tensor<2x3xi32>
// CHECK: {{%.*}} = spirv.EXT.ConstantCompositeReplicate [2 : i32] : !spirv.arm.tensor<2x3xi32>
%0 = spirv.Constant dense<2> : !spirv.arm.tensor<2x3xi32>
spirv.ReturnValue %0 : !spirv.arm.tensor<2x3xi32>
}
Expand All @@ -82,8 +82,8 @@ spirv.module Logical GLSL450 {

spirv.module Logical GLSL450 {
spirv.func @splat_vector_of_f32() -> (vector<3xf32>) "None" {
%0 = spirv.Constant dense<2.0> : vector<3xf32>
// CHECK: {{%.*}} = spirv.EXT.ConstantCompositeReplicate [2.000000e+00 : f32] : vector<3xf32>
%0 = spirv.Constant dense<2.0> : vector<3xf32>
spirv.ReturnValue %0 : vector<3xf32>
}
}
Expand All @@ -92,8 +92,8 @@ spirv.module Logical GLSL450 {

spirv.module Logical GLSL450 {
spirv.func @splat_array_of_f32() -> (!spirv.array<3 x f32>) "None" {
%0 = spirv.Constant [1.0 : f32, 1.0 : f32, 1.0 : f32] : !spirv.array<3 x f32>
// CHECK: {{%.*}} = spirv.EXT.ConstantCompositeReplicate [1.000000e+00 : f32] : !spirv.array<3 x f32>
%0 = spirv.Constant [1.0 : f32, 1.0 : f32, 1.0 : f32] : !spirv.array<3 x f32>
spirv.ReturnValue %0 : !spirv.array<3 x f32>
}
}
Expand All @@ -102,7 +102,7 @@ spirv.module Logical GLSL450 {

spirv.module Logical GLSL450 {
spirv.func @splat_array_of_splat_array_of_f32() -> (!spirv.array<2 x !spirv.array<3 x f32>>) "None" {
// CHECK: spirv.EXT.ConstantCompositeReplicate [3.000000e+00 : f32] : !spirv.array<2 x !spirv.array<3 x f32>>
// CHECK: {{%.*}} = spirv.EXT.ConstantCompositeReplicate [3.000000e+00 : f32] : !spirv.array<2 x !spirv.array<3 x f32>>
%0 = spirv.Constant [[3.0 : f32, 3.0 : f32, 3.0 : f32], [3.0 : f32, 3.0 : f32, 3.0 : f32]] : !spirv.array<2 x !spirv.array<3 x f32>>
spirv.ReturnValue %0 : !spirv.array<2 x !spirv.array<3 x f32>>
}
Expand All @@ -112,7 +112,7 @@ spirv.module Logical GLSL450 {

spirv.module Logical GLSL450 {
spirv.func @splat_array_of_non_splat_array_of_f32() -> (!spirv.array<2 x !spirv.array<3 x f32>>) "None" {
// CHECK: spirv.EXT.ConstantCompositeReplicate {{\[}}[1.000000e+00 : f32, 2.000000e+00 : f32, 3.000000e+00 : f32]] : !spirv.array<2 x !spirv.array<3 x f32>>
// CHECK: {{%.*}} = spirv.EXT.ConstantCompositeReplicate {{\[}}[1.000000e+00 : f32, 2.000000e+00 : f32, 3.000000e+00 : f32]] : !spirv.array<2 x !spirv.array<3 x f32>>
%0 = spirv.Constant [[1.0 : f32, 2.0 : f32, 3.0 : f32], [1.0 : f32, 2.0 : f32, 3.0 : f32]] : !spirv.array<2 x !spirv.array<3 x f32>>
spirv.ReturnValue %0 : !spirv.array<2 x !spirv.array<3 x f32>>
}
Expand All @@ -122,8 +122,8 @@ spirv.module Logical GLSL450 {

spirv.module Logical GLSL450 {
spirv.func @splat_array_of_vectors_of_f32() -> (!spirv.array<2xvector<2xf32>>) "None" {
%0 = spirv.Constant [dense<[1.0, 2.0]> : vector<2xf32>, dense<[1.0, 2.0]> : vector<2xf32>] : !spirv.array<2 x vector<2xf32>>
// CHECK: {{%.*}} = spirv.EXT.ConstantCompositeReplicate [dense<[1.000000e+00, 2.000000e+00]> : vector<2xf32>] : !spirv.array<2 x vector<2xf32>>
%0 = spirv.Constant [dense<[1.0, 2.0]> : vector<2xf32>, dense<[1.0, 2.0]> : vector<2xf32>] : !spirv.array<2 x vector<2xf32>>
spirv.ReturnValue %0 : !spirv.array<2 x vector<2xf32>>
}
}
Expand All @@ -132,8 +132,8 @@ spirv.module Logical GLSL450 {

spirv.module Logical GLSL450 {
spirv.func @splat_array_of_splat_vectors_of_f32() -> (!spirv.array<2 x vector<2xf32>>) "None" {
%0 = spirv.Constant [dense<2.0> : vector<2xf32>, dense<2.0> : vector<2xf32>] : !spirv.array<2 x vector<2xf32>>
// CHECK: {{%.*}} = spirv.EXT.ConstantCompositeReplicate [2.000000e+00 : f32] : !spirv.array<2 x vector<2xf32>>
%0 = spirv.Constant [dense<2.0> : vector<2xf32>, dense<2.0> : vector<2xf32>] : !spirv.array<2 x vector<2xf32>>
spirv.ReturnValue %0 : !spirv.array<2 x vector<2xf32>>
}
}
Expand All @@ -142,7 +142,7 @@ spirv.module Logical GLSL450 {

spirv.module Logical GLSL450 {
spirv.func @splat_tensor_of_f32() -> (!spirv.array<2 x !spirv.array<3 x f32>>) "None" {
// CHECK: spirv.EXT.ConstantCompositeReplicate [3.000000e+00 : f32] : !spirv.array<2 x !spirv.array<3 x f32>>
// CHECK: {{%.*}} = spirv.EXT.ConstantCompositeReplicate [3.000000e+00 : f32] : !spirv.array<2 x !spirv.array<3 x f32>>
%0 = spirv.Constant dense<3.0> : tensor<2x3xf32> : !spirv.array<2 x !spirv.array<3 x f32>>
spirv.ReturnValue %0 : !spirv.array<2 x !spirv.array<3 x f32>>
}
Expand All @@ -152,14 +152,94 @@ spirv.module Logical GLSL450 {

spirv.module Logical GLSL450 {
spirv.func @splat_arm_tensor_of_f32() -> (!spirv.arm.tensor<2x3xf32>) "None" {
// CHECK: spirv.EXT.ConstantCompositeReplicate [2.000000e+00 : f32] : !spirv.arm.tensor<2x3xf32>
// CHECK: {{%.*}} = spirv.EXT.ConstantCompositeReplicate [2.000000e+00 : f32] : !spirv.arm.tensor<2x3xf32>
%0 = spirv.Constant dense<2.0> : !spirv.arm.tensor<2x3xf32>
spirv.ReturnValue %0 : !spirv.arm.tensor<2x3xf32>
}
}

// -----

spirv.module Logical GLSL450 {
spirv.func @array_of_one_i32() -> (!spirv.array<1 x i32>) "None" {
// CHECK-NOT: spirv.EXT.ConstantCompositeReplicate
%0 = spirv.Constant [1 : i32] : !spirv.array<1 x i32>
spirv.ReturnValue %0 : !spirv.array<1 x i32>
}
}

// -----

spirv.module Logical GLSL450 {
spirv.func @arm_tensor_of_one_i32() -> (!spirv.arm.tensor<1xi32>) "None" {
// CHECK-NOT: spirv.EXT.ConstantCompositeReplicate
%0 = spirv.Constant dense<1> : !spirv.arm.tensor<1xi32>
spirv.ReturnValue %0 : !spirv.arm.tensor<1xi32>
}
}

// -----

spirv.module Logical GLSL450 {
spirv.func @non_splat_vector_of_i32() -> (vector<3xi32>) "None" {
// CHECK-NOT: spirv.EXT.ConstantCompositeReplicate
%0 = spirv.Constant dense<[0, 1, 2]> : vector<3xi32>
spirv.ReturnValue %0 : vector<3xi32>
}
}

// -----

spirv.module Logical GLSL450 {
spirv.func @non_splat_array_of_vectors_of_i32() -> (!spirv.array<2xvector<2xi32>>) "None" {
// CHECK-NOT: spirv.EXT.ConstantCompositeReplicate
%0 = spirv.Constant [dense<[1, 2]> : vector<2xi32>, dense<[1, 3]> : vector<2xi32>] : !spirv.array<2 x vector<2xi32>>
spirv.ReturnValue %0 : !spirv.array<2 x vector<2xi32>>
}
}

// -----

spirv.module Logical GLSL450 {
spirv.func @array_of_one_f32() -> (!spirv.array<1 x f32>) "None" {
// CHECK-NOT: spirv.EXT.ConstantCompositeReplicate
%0 = spirv.Constant [1.0 : f32] : !spirv.array<1 x f32>
spirv.ReturnValue %0 : !spirv.array<1 x f32>
}
}

// -----

spirv.module Logical GLSL450 {
spirv.func @arm_tensor_of_one_f32() -> (!spirv.arm.tensor<1xf32>) "None" {
// CHECK-NOT: spirv.EXT.ConstantCompositeReplicate
%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>>
spirv.ReturnValue %0 : !spirv.array<2 x vector<2xf32>>
}
}

// -----

spirv.module Logical GLSL450 {

spirv.SpecConstant @sc_i32_1 = 1 : i32
Expand Down Expand Up @@ -189,4 +269,35 @@ spirv.module Logical GLSL450 {

// CHECK: spirv.EXT.SpecConstantCompositeReplicate @scc_splat_arm_tensor_of_f32 (@sc_f32_1) : !spirv.arm.tensor<3xf32>
spirv.SpecConstantComposite @scc_splat_arm_tensor_of_f32 (@sc_f32_1, @sc_f32_1, @sc_f32_1) : !spirv.arm.tensor<3xf32>

spirv.SpecConstant @sc_i32_2 = 2 : i32

// CHECK-NOT: spirv.EXT.SpecConstantCompositeReplicate
spirv.SpecConstantComposite @scc_array_of_one_i32 (@sc_i32_1) : !spirv.array<1 x i32>

// CHECK-NOT: spirv.EXT.SpecConstantCompositeReplicate
spirv.SpecConstantComposite @scc_arm_tensor_of_one_i32 (@sc_i32_1) : !spirv.arm.tensor<1xi32>

// CHECK-NOT: spirv.EXT.SpecConstantCompositeReplicate
spirv.SpecConstantComposite @scc_non_splat_vector_of_i32 (@sc_i32_1, @sc_i32_1, @sc_i32_2) : vector<3 x i32>

// CHECK-NOT: spirv.EXT.SpecConstantCompositeReplicate
spirv.SpecConstantComposite @scc_non_splat_arm_tensor_of_i32 (@sc_i32_2, @sc_i32_1, @sc_i32_1) : !spirv.arm.tensor<3xi32>

spirv.SpecConstant @sc_f32_2 = 2.0 : f32

// CHECK-NOT: spirv.EXT.SpecConstantCompositeReplicate
spirv.SpecConstantComposite @scc_array_of_one_f32 (@sc_f32_1) : !spirv.array<1 x f32>

// CHECK-NOT: spirv.EXT.SpecConstantCompositeReplicate
spirv.SpecConstantComposite @scc_arm_tensor_of_one_f32 (@sc_f32_1) : !spirv.arm.tensor<1xf32>

// CHECK-NOT: spirv.EXT.SpecConstantCompositeReplicate
spirv.SpecConstantComposite @scc_non_splat_vector_of_f32 (@sc_f32_1, @sc_f32_1, @sc_f32_2) : vector<3 x f32>

// CHECK-NOT: spirv.EXT.SpecConstantCompositeReplicate
spirv.SpecConstantComposite @scc_non_splat_arm_tensor_of_f32 (@sc_f32_2, @sc_f32_1, @sc_f32_1) : !spirv.arm.tensor<3xf32>

// 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)>
}