Add tk assertion

nod-ai · Oct 8, 2024 · fda11c9 · fda11c9
1 parent 9f34a73
commit fda11c9
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/gemmbench/gemm_utils.py b/gemmbench/gemm_utils.py
@@ -116,7 +116,10 @@ def generate_mlir(config: GemmConfig):
         return mlir_template_B
     return mlir_template
 
+
 def generate_tk_mlir(config: GemmConfig):
+    assert config.operand_element_type == 'f16', "Unsupported problem"
+    assert config.accumulator_element_type == 'f32', "Unsupported problem"
     # Input sizes
     M = tkl.sym.M
     N = tkl.sym.N