From 39de4d62656b1d07844ab433a80376f7d9f93997 Mon Sep 17 00:00:00 2001
From: Sean Silva <silvasean@google.com>
Date: Tue, 15 Nov 2022 15:25:39 +0000
Subject: [PATCH] [cleanup] Make diagnostics better

Also remove some unused imports.
---
 lib/RefBackend/RefBackend.cpp                            | 9 ++++++---
 python/torch_mlir_e2e_test/configs/native_torch.py       | 3 ---
 .../linalg_on_tensors_backends/refbackend.py             | 2 +-
 test/RefBackend/munge-calling-conventions.mlir           | 9 ++++++++-
 4 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/lib/RefBackend/RefBackend.cpp b/lib/RefBackend/RefBackend.cpp
index 9d3cd786e..981263025 100644
--- a/lib/RefBackend/RefBackend.cpp
+++ b/lib/RefBackend/RefBackend.cpp
@@ -142,9 +142,12 @@ static LogicalResult mungeFunction(
   SmallVector<Type> newArgTypes;
   for (auto arg : func.getArguments()) {
     auto type = arg.getType();
-    if (!isArgMemRefTypeValid(type))
-      return emitError(arg.getLoc(),
-                       "argument must be a memref of f32, f64, i32, i64, i8, i1");
+    if (!isArgMemRefTypeValid(type)) {
+      return emitError(arg.getLoc())
+          .append("argument must be a memref of f32, f64, i32, i64, i8, i1 but "
+                  "got ",
+                  type);
+    }
     auto cast = b.create<memref::CastOp>(arg.getLoc(), type, arg);
     arg.replaceAllUsesExcept(cast, cast);
     arg.setType(getAbiTypeForMemRef(type));
diff --git a/python/torch_mlir_e2e_test/configs/native_torch.py b/python/torch_mlir_e2e_test/configs/native_torch.py
index b85353f65..76c349bc3 100644
--- a/python/torch_mlir_e2e_test/configs/native_torch.py
+++ b/python/torch_mlir_e2e_test/configs/native_torch.py
@@ -3,9 +3,6 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 # Also available under a BSD-style license. See LICENSE.
 
-import copy
-from typing import Any
-
 import torch
 
 from torch_mlir_e2e_test.framework import TestConfig, Trace, TraceItem
diff --git a/python/torch_mlir_e2e_test/linalg_on_tensors_backends/refbackend.py b/python/torch_mlir_e2e_test/linalg_on_tensors_backends/refbackend.py
index be59ee907..d3c5788f0 100644
--- a/python/torch_mlir_e2e_test/linalg_on_tensors_backends/refbackend.py
+++ b/python/torch_mlir_e2e_test/linalg_on_tensors_backends/refbackend.py
@@ -22,7 +22,7 @@ __all__ = [
 
 def assert_arg_type_is_supported(ty):
     SUPPORTED = [np.float32, np.float64, np.uint8, np.int8, np.int32, np.int64, np.bool_]
-    assert ty in SUPPORTED, f"Only numpy arrays with dtypes in {SUPPORTED} are supported"
+    assert ty in SUPPORTED, f"Only numpy arrays with dtypes in {SUPPORTED} are supported, but got {ty}"
 
 
 memref_type_to_np_dtype = {
diff --git a/test/RefBackend/munge-calling-conventions.mlir b/test/RefBackend/munge-calling-conventions.mlir
index c55e43b3b..12fd639db 100644
--- a/test/RefBackend/munge-calling-conventions.mlir
+++ b/test/RefBackend/munge-calling-conventions.mlir
@@ -1,4 +1,4 @@
-// RUN: torch-mlir-opt %s -refback-munge-calling-conventions -split-input-file | FileCheck %s
+// RUN: torch-mlir-opt %s -refback-munge-calling-conventions -split-input-file -verify-diagnostics | FileCheck %s
 
 // CHECK-LABEL:   func.func @f(
 // CHECK-SAME:            %[[ARG0:.*]]: memref<*xf32>) attributes {llvm.emit_c_interface} {
@@ -70,3 +70,10 @@ func.func @multiple_return_values(%arg0: memref<?xf32>, %arg1: memref<?xf32>, %a
 func.func @two_return_values(%arg0: memref<?xf32>, %arg1: memref<?xi64>) -> (memref<?xf32>, memref<?xi64>) {
   return %arg0 ,%arg1 : memref<?xf32>, memref<?xi64>
 }
+
+// -----
+
+// expected-error-re @+1 {{argument must be a memref of {{.*}} but got 'tensor<?xf32>'}}
+func.func @f(%arg0: tensor<?xf32>) {
+  return
+}