41 lines · python
1# RUN: %PYTHON %s | FileCheck %s2# This is just a smoke test that the dialect is functional.3from array import array4 5from mlir.ir import *6from mlir.dialects import rocdl, arith7from mlir.extras import types as T8 9 10def constructAndPrintInModule(f):11 print("\nTEST:", f.__name__)12 with Context(), Location.unknown():13 module = Module.create()14 with InsertionPoint(module.body):15 f()16 print(module)17 return f18 19 20# CHECK-LABEL: testSmoke21@constructAndPrintInModule22def testSmoke():23 v_len = 1624 f32 = F32Type.get()25 # Note: this isn't actually the right type for the intrinsic (should be f16)26 # but array doesn't support f16.27 v16f32 = T.vector(v_len, f32)28 f32_array = array("f", [0.0] * v_len)29 a_frag = arith.constant(v16f32, f32_array)30 b_frag = arith.constant(v16f32, f32_array)31 c_frag = arith.constant(v16f32, f32_array)32 33 c_frag = rocdl.wmma_f16_16x16x16_f16(v16f32, a_frag, b_frag, c_frag, opsel=False)34 # CHECK: %{{.*}} = "rocdl.wmma.f16.16x16x16.f16"35 print(c_frag)36 assert isinstance(c_frag, OpView)37 # CHECK: Value(%{{.*}} = "rocdl.wmma.f16.16x16x16.f16"38 c_frag = rocdl.wmma_f16_16x16x16_f16_(v16f32, a_frag, b_frag, c_frag, opsel=False)39 print(c_frag)40 assert isinstance(c_frag, Value)41