Add verifiers for ToBuiltinTensorOp and FromBuiltinTensorOp (#1089)

This commit adds verifiers to the ops `ToBuiltinTensorOp` and
`FromBuiltinTensorOp` that make sure that the input and output have
the same shape and data type.

Author: ramiro050

include/torch-mlir/Dialect/TorchConversion/IR/TorchConversionOps.td

@@ -42,6 +42,7 @@ def TorchConversion_ToBuiltinTensorOp : TorchConversion_Op<"to_builtin_tensor",
   let assemblyFormat = [{
     $operand attr-dict `:` qualified(type($operand)) `->` qualified(type($result))
   }];
+  let hasVerifier = 1;
 }
 
 def TorchConversion_FromBuiltinTensorOp : TorchConversion_Op<"from_builtin_tensor">
@@ -60,6 +61,7 @@ def TorchConversion_FromBuiltinTensorOp : TorchConversion_Op<"from_builtin_tenso
   let assemblyFormat = [{
     $operand attr-dict `:` qualified(type($operand)) `->` qualified(type($result))
   }];
+  let hasVerifier = 1;
 }
 
 def TorchConversion_ToI1Op : TorchConversion_Op<"to_i1", [

lib/Dialect/TorchConversion/IR/TorchConversionOps.cpp

@@ -21,10 +21,29 @@ using namespace mlir::torch;
 using namespace mlir::torch::TorchConversion;
 using namespace mlir::torch;
 
+static bool haveSameSizeAndElementType(TensorType lhs, TensorType rhs) {
+  if (lhs.hasRank() != rhs.hasRank())
+    return false;
+  bool sameSize = lhs.hasRank() ? lhs.getShape().equals(rhs.getShape()) : true;
+  bool sameElementType = lhs.getElementType() == rhs.getElementType();
+  return sameElementType && sameSize;
+}
+
 //===----------------------------------------------------------------------===//
 // ToBuiltinTensorOp
 //===----------------------------------------------------------------------===//
 
+LogicalResult ToBuiltinTensorOp::verify() {
+  auto resultType = getResult().getType().cast<TensorType>();
+  auto operandType =
+      getOperand().getType().cast<Torch::ValueTensorType>().toBuiltinTensor();
+  if (!haveSameSizeAndElementType(resultType, operandType)) {
+    return emitError()
+           << "operand and result must have the same size and dtype";
+  }
+  return success();
+}
+
 LogicalResult ToBuiltinTensorOp::inferReturnTypes(
     MLIRContext *context, Optional<Location> location, ValueRange operands,
     DictionaryAttr attributes, RegionRange regions,
@@ -37,6 +56,25 @@ LogicalResult ToBuiltinTensorOp::inferReturnTypes(
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// FromBuiltinTensorOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult FromBuiltinTensorOp::verify() {
+  auto resultType =
+      getResult().getType().cast<Torch::ValueTensorType>().toBuiltinTensor();
+  auto operandType = getOperand().getType().cast<TensorType>();
+  if (!haveSameSizeAndElementType(resultType, operandType)) {
+    return emitError()
+           << "operand and result must have the same size and dtype";
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// FromI64Op
+//===----------------------------------------------------------------------===//
+
 OpFoldResult FromI64Op::fold(llvm::ArrayRef<mlir::Attribute> operands) {
   auto attr = operands[0].dyn_cast_or_null<mlir::IntegerAttr>();
   if (attr) {
@@ -46,6 +84,10 @@ OpFoldResult FromI64Op::fold(llvm::ArrayRef<mlir::Attribute> operands) {
   }
 }
 
+//===----------------------------------------------------------------------===//
+// ToI64Op
+//===----------------------------------------------------------------------===//
+
 OpFoldResult ToI64Op::fold(llvm::ArrayRef<mlir::Attribute> operands) {
   auto attr = operands[0].dyn_cast_or_null<mlir::IntegerAttr>();
   if (attr) {

test/Dialect/TorchConversion/ops.mlir

@@ -1,4 +1,4 @@
-// RUN: torch-mlir-opt %s | torch-mlir-opt | FileCheck %s
+// RUN: torch-mlir-opt %s -split-input-file -verify-diagnostics | FileCheck %s
 
 // CHECK-LABEL: func.func @builtin_tensor_interop(
 func.func @builtin_tensor_interop(%arg0: tensor<*xf32>, %arg1: tensor<3x?xi8>, %arg2: !torch.vtensor<*,f32>, %arg3: !torch.vtensor<[3,?],si8>) {
@@ -14,3 +14,35 @@ func.func @builtin_tensor_interop(%arg0: tensor<*xf32>, %arg1: tensor<3x?xi8>, %
   %4 = torch_c.to_builtin_tensor %arg3 : !torch.vtensor<[3,?],si8> -> tensor<3x?xi8>
   return
 }
+
+// -----
+
+func.func @to_builtin_tensor_invalid_size(%arg0: !torch.vtensor<[3,?],si8>) {
+  // expected-error @+1 {{operand and result must have the same size and dtype}}
+  %1 = torch_c.to_builtin_tensor %arg0 : !torch.vtensor<[3,?],si8> -> tensor<?x?xi8>
+  return
+}
+
+// -----
+
+func.func @to_builtin_tensor_invalid_dtype(%arg0: !torch.vtensor<*,si8>) {
+  // expected-error @+1 {{operand and result must have the same size and dtype}}
+  %1 = torch_c.to_builtin_tensor %arg0 : !torch.vtensor<*,si8> -> tensor<*xi64>
+  return
+}
+
+// -----
+
+func.func @from_builtin_tensor_invalid_size(%arg0: tensor<3x?xi8>) {
+  // expected-error @+1 {{operand and result must have the same size and dtype}}
+  %1 = torch_c.from_builtin_tensor %arg0 : tensor<3x?xi8> -> !torch.vtensor<[?,?],si8>
+  return
+}
+
+// -----
+
+func.func @from_builtin_tensor_invalid_dtype(%arg0: tensor<*xi8>) {
+  // expected-error @+1 {{operand and result must have the same size and dtype}}
+  %1 = torch_c.from_builtin_tensor %arg0 : tensor<*xi8> -> !torch.vtensor<*,si64>
+  return
+}