distance op verify + distance op tests

Author: mishaobu

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

@@ -2104,20 +2104,29 @@ LogicalResult spirv::GLDistanceOp::verify() {
   auto p1Type = getP1().getType();
   auto resultType = getResult().getType();
 
-  auto p0VectorType = p0Type.dyn_cast<VectorType>();
-  auto p1VectorType = p1Type.dyn_cast<VectorType>();
-  if (!p0VectorType || !p1VectorType)
-    return emitOpError("operands must be vectors");
+  auto getFloatType = [](Type type) -> FloatType {
+    if (auto vectorType = llvm::dyn_cast<VectorType>(type))
+      return llvm::dyn_cast<FloatType>(vectorType.getElementType());
+    return llvm::dyn_cast<FloatType>(type);
+  };
 
-  if (p0VectorType.getShape() != p1VectorType.getShape())
-    return emitOpError("operands must have same shape");
+  FloatType p0FloatType = getFloatType(p0Type);
+  FloatType p1FloatType = getFloatType(p1Type);
+  FloatType resultFloatType = llvm::dyn_cast<FloatType>(resultType);
 
-  if (!resultType.isa<FloatType>())
-    return emitOpError("result must be scalar float");
+  if (!p0FloatType || !p1FloatType || !resultFloatType)
+    return emitOpError("operands and result must be float scalar or vector of float"); 
 
-  if (p0VectorType.getElementType() != resultType || 
-      p1VectorType.getElementType() != resultType)
-    return emitOpError("operand vector elements must match result type");
+  if (p0FloatType != resultFloatType || p1FloatType != resultFloatType)
+    return emitOpError("operand and result element types must match");
+
+  if (auto p0Vec = llvm::dyn_cast<VectorType>(p0Type)) {
+    if (!llvm::dyn_cast<VectorType>(p1Type) || 
+        p0Vec.getShape() != llvm::dyn_cast<VectorType>(p1Type).getShape())
+      return emitOpError("vector operands must have same shape");
+  } else if (llvm::isa<VectorType>(p1Type)) {
+    return emitOpError("expected both operands to be scalars or both to be vectors");
+  }
 
   return success();
 }

mlir/test/Dialect/SPIRV/IR/gl-ops.mlir

@@ -541,3 +541,45 @@ func.func @findumsb(%arg0 : i64) -> () {
   %2 = spirv.GL.FindUMsb %arg0 : i64
   return
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spirv.GL.Distance 
+//===----------------------------------------------------------------------===//
+
+func.func @distance_scalar(%arg0 : f32, %arg1 : f32) {
+  // CHECK: spirv.GL.Distance {{%.*}}, {{%.*}} : f32, f32 -> f32
+  %0 = spirv.GL.Distance %arg0, %arg1 : f32, f32 -> f32
+  return
+}
+
+func.func @distance_vector(%arg0 : vector<3xf32>, %arg1 : vector<3xf32>) {
+  // CHECK: spirv.GL.Distance {{%.*}}, {{%.*}} : vector<3xf32>, vector<3xf32> -> f32
+  %0 = spirv.GL.Distance %arg0, %arg1 : vector<3xf32>, vector<3xf32> -> f32
+  return
+}
+
+// -----
+
+func.func @distance_invalid_type(%arg0 : i32, %arg1 : i32) {
+  // expected-error @+1 {{'spirv.GL.Distance' op operand #0 must be 16/32/64-bit float or vector of 16/32/64-bit float values of length 2/3/4/8/16}}
+  %0 = spirv.GL.Distance %arg0, %arg1 : i32, i32 -> f32
+  return
+}
+
+// -----
+
+func.func @distance_invalid_vector_size(%arg0 : vector<5xf32>, %arg1 : vector<5xf32>) {
+  // expected-error @+1 {{'spirv.GL.Distance' op operand #0 must be 16/32/64-bit float or vector of 16/32/64-bit float values of length 2/3/4/8/16}}
+  %0 = spirv.GL.Distance %arg0, %arg1 : vector<5xf32>, vector<5xf32> -> f32
+  return
+}
+
+// -----
+
+func.func @distance_invalid_result(%arg0 : f32, %arg1 : f32) {
+  // expected-error @+1 {{'spirv.GL.Distance' op result #0 must be 16/32/64-bit float}}
+  %0 = spirv.GL.Distance %arg0, %arg1 : f32, f32 -> i32
+  return
+}