rocm: fix dense_rocblas in strategy, topi (apache#5191)

trevor-m · Apr 1, 2020 · 759d901 · 759d901
1 parent c3a64c1
commit 759d901
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/python/tvm/relay/op/strategy/rocm.py b/python/tvm/relay/op/strategy/rocm.py
@@ -129,7 +129,7 @@ def dense_strategy_rocm(attrs, inputs, out_type, target):
         assert out_type.dtype == inputs[0].dtype, "Mixed precision not supported."
         strategy.add_implementation(
             wrap_compute_dense(topi.rocm.dense_rocblas),
-            wrap_topi_schedule(topi.rocm.dense_rocblas),
+            wrap_topi_schedule(topi.rocm.schedule_dense_rocblas),
             name="dense_rocblas.rocm",
             plevel=15)
     return strategy
diff --git a/topi/python/topi/rocm/dense.py b/topi/python/topi/rocm/dense.py
@@ -123,6 +123,8 @@ def dense_rocblas(cfg, data, weight, bias=None, out_dtype=None):
     output : tvm.te.Tensor
         2-D with shape [batch, out_dim]
     """
+    if out_dtype is None:
+        out_dtype = data.dtype
     assert out_dtype == data.dtype, "Mixed precision not supported."
     matmul = rocblas.matmul(data, weight, False, True)
     batch, in_dim = data.shape