[mlir][spirv] Fix a crash in spirv::ISubOp::fold

[CoTinker]

This PR fixes a crash if spirv.ISub is not integer type. Fixes #131283.

[llvmbot]

@llvm/pr-subscribers-mlir-spirv

@llvm/pr-subscribers-mlir

Author: Longsheng Mou (CoTinker)

Changes

This PR fixes a crash if spirv.ISub is not integer type. Fixes #131283.

---

Full diff: https://github.com/llvm/llvm-project/pull/131570.diff

2 Files Affected:

• (modified) mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp (+1-1)
• (modified) mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir (+7)

diff --git a/mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp b/mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp
index b1acfd1a2abed..e36d4b910193e 100644
--- a/mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/SPIRVCanonicalization.cpp
@@ -476,7 +476,7 @@ OpFoldResult spirv::IMulOp::fold(FoldAdaptor adaptor) {
 OpFoldResult spirv::ISubOp::fold(FoldAdaptor adaptor) {
   // x - x = 0
   if (getOperand1() == getOperand2())
-    return Builder(getContext()).getIntegerAttr(getType(), 0);
+    return Builder(getContext()).getZeroAttr(getType());
 
   // According to the SPIR-V spec:
   //
diff --git a/mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir b/mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir
index 3a775e209903c..0fd6c18a6c241 100644
--- a/mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir
+++ b/mlir/test/Dialect/SPIRV/Transforms/canonicalize.mlir
@@ -593,6 +593,13 @@ func.func @isub_x_x(%arg0: i32) -> i32 {
   return %0: i32
 }
 
+// CHECK-LABEL: @isub_vector_x_x
+func.func @isub_vector_x_x(%arg0: vector<3xi32>) -> vector<3xi32> {
+  // CHECK: spirv.Constant dense<0>
+  %0 = spirv.ISub %arg0, %arg0: vector<3xi32>
+  return %0: vector<3xi32>
+}
+
 // CHECK-LABEL: @const_fold_scalar_isub_normal
 func.func @const_fold_scalar_isub_normal() -> (i32, i32, i32) {
   %c5 = spirv.Constant 5 : i32