[mlir][arith] Refine the verifier for arith.constant (#87999)

Disallows initialization of scalable vectors with an attribute of
arbitrary values, e.g.:
```mlir
  %c = arith.constant dense<[0, 1]> : vector<[2] x i32>
```

Initialization using vector splats remains allowed (i.e. when all the
init values are identical):
```mlir
  %c = arith.constant dense<[1, 1]> : vector<[2] x i32>
```

Note: This is a re-upload of #86178

Author: banach-space

mlir/lib/Dialect/Arith/IR/ArithOps.cpp

@@ -213,6 +213,15 @@ LogicalResult arith::ConstantOp::verify() {
     return emitOpError(
         "value must be an integer, float, or elements attribute");
   }
+
+  // Note, we could relax this for vectors with 1 scalable dim, e.g.:
+  //  * arith.constant dense<[[3, 3], [1, 1]]> : vector<2 x [2] x i32>
+  // However, this would most likely require updating the lowerings to LLVM.
+  auto vecType = dyn_cast<VectorType>(type);
+  if (vecType && vecType.isScalable() && !isa<SplatElementsAttr>(getValue()))
+    return emitOpError(
+        "intializing scalable vectors with elements attribute is not supported"
+        " unless it's a vector splat");
   return success();
 }
 

mlir/test/Dialect/Arith/invalid.mlir

@@ -64,6 +64,24 @@ func.func @constant_out_of_range() {
 
 // -----
 
+func.func @constant_invalid_scalable_1d_vec_initialization() {
+^bb0:
+  // expected-error@+1 {{'arith.constant' op intializing scalable vectors with elements attribute is not supported unless it's a vector splat}}
+  %c = arith.constant dense<[0, 1]> : vector<[2] x i32>
+  return
+}
+
+// -----
+
+func.func @constant_invalid_scalable_2d_vec_initialization() {
+^bb0:
+  // expected-error@+1 {{'arith.constant' op intializing scalable vectors with elements attribute is not supported unless it's a vector splat}}
+  %c = arith.constant dense<[[3, 3], [1, 1]]> : vector<2 x [2] x i32>
+  return
+}
+
+// -----
+
 func.func @constant_wrong_type() {
 ^bb:
   %x = "arith.constant"(){value = 10.} : () -> f32 // expected-error {{'arith.constant' op failed to verify that all of {value, result} have same type}}

mlir/test/Dialect/Vector/linearize.mlir

@@ -153,14 +153,3 @@ func.func @test_0d_vector() -> vector<f32> {
   // ALL: return %[[CST]]
   return %0 : vector<f32>
 }
-
-// -----
-
-func.func @test_scalable_no_linearize(%arg0: vector<2x[2]xf32>) -> vector<2x[2]xf32> {
-  // expected-error@+1 {{failed to legalize operation 'arith.constant' that was explicitly marked illegal}}
-  %0 = arith.constant dense<[[1., 1.], [3., 3.]]> : vector<2x[2]xf32>
-  %1 = math.sin %arg0 : vector<2x[2]xf32>
-  %2 = arith.addf %0, %1 : vector<2x[2]xf32>
-
-  return %2 : vector<2x[2]xf32>
-}