int8 test worked

apache · Jan 11, 2022 · ad3faf6 · ad3faf6
1 parent e9b0287
commit ad3faf6
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/src/relay/backend/contrib/cutlass/codegen.cc b/src/relay/backend/contrib/cutlass/codegen.cc
@@ -43,7 +43,11 @@ namespace contrib {
 using namespace backend;
 using Str2StrMap = std::unordered_map<std::string, std::string>;
 
-static Str2StrMap dtype_map = {{"float16", "cutlass::half_t"}, {"float32", "float"}};
+static Str2StrMap dtype_map = {{"float16", "cutlass::half_t"},
+                               {"float32", "float"},
+                               {"int8", "int8_t"},
+                               {"uint8", "uint8_t"},
+                               {"int32", "int32_t"}};
 
 constexpr const char* kAnyDim = "Any";
 

diff --git a/tests/python/contrib/test_cutlass.py b/tests/python/contrib/test_cutlass.py
@@ -392,6 +392,7 @@ def verify_conv2d(
     use_fast_math=False,
     data_dtype="float16",
     weight_dtype="float16",
+    ref_target="cuda",
 ):
     if not has_cutlass():
         return
@@ -436,7 +437,7 @@ def verify_conv2d(
         rt_mod_ref, dev = get_ref_rt_mod(
             convert_conv2d_layout(mod_ref, {"nn.conv2d": ["NHWC", "HWIO"]}),
             params,
-            target="cuda",
+            target=ref_target,
         )
 
     ref_out = get_output(rt_mod_ref, ["data"], [np_data])
@@ -575,6 +576,7 @@ def test_int8():
         run_benchmark=False,
         data_dtype="int8",
         weight_dtype="int8",
+        ref_target="llvm",
     )
 
 
@@ -610,4 +612,5 @@ def test_3xtf32():
 
 
 if __name__ == "__main__":
-    pytest.main([__file__])
+    # pytest.main([__file__])
+    test_int8()