Ejemplo n.º 1
0
class _SparseTensorDescriptor(ctypes.Structure):
    """A C structure for an MLIR sparse tensor."""
    _fields_ = [
        # A pointer for the MLIR sparse tensor storage.
        ("storage", ctypes.POINTER(ctypes.c_ulonglong)),
        # An MLIR MemRef descriptor for the shape of the sparse tensor.
        ("shape", runtime.make_nd_memref_descriptor(1, ctypes.c_ulonglong)),
    ]
Ejemplo n.º 2
0
def build_compile_and_run_SDDMMM(attr: st.EncodingAttr, opt: str,
                                 support_lib: str, compiler):
    # Build.
    module = build_SDDMM(attr)
    func = str(module.operation.regions[0].blocks[0].operations[0].operation)
    module = ir.Module.parse(func + boilerplate(attr))

    # Compile.
    compiler(module)
    engine = execution_engine.ExecutionEngine(module,
                                              opt_level=0,
                                              shared_libs=[support_lib])

    # Set up numpy input and buffer for output.
    a = np.array([[1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1, 8.1],
                  [1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2, 8.2],
                  [1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3, 8.3],
                  [1.4, 2.4, 3.4, 4.4, 5.4, 6.4, 7.4, 8.4],
                  [1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5],
                  [1.6, 2.6, 3.6, 4.6, 5.6, 6.6, 7.6, 8.6],
                  [1.7, 2.7, 3.7, 4.7, 5.7, 6.7, 7.7, 8.7],
                  [1.8, 2.8, 3.8, 4.8, 5.8, 6.8, 7.8, 8.8]], np.float64)
    b = np.ones((8, 8), np.float64)
    c = np.zeros((8, 8), np.float64)

    mem_a = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(a)))
    mem_b = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(b)))
    mem_c = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(c)))

    # Allocate a MemRefDescriptor to receive the output tensor.
    # The buffer itself is allocated inside the MLIR code generation.
    ref_out = rt.make_nd_memref_descriptor(2, ctypes.c_double)()
    mem_out = ctypes.pointer(ctypes.pointer(ref_out))

    # Invoke the kernel and get numpy output.
    # Built-in bufferization uses in-out buffers.
    # TODO: replace with inplace comprehensive bufferization.
    engine.invoke('main', mem_out, mem_a, mem_b, mem_c)

    # Sanity check on computed result. Only a few elements
    # are sampled from the full dense matrix multiplication.
    full_matmul = np.matmul(a, b)
    expected = np.zeros((8, 8), np.float64)
    expected[0, 0] = 1.0 * full_matmul[0, 0]
    expected[0, 2] = 2.0 * full_matmul[0, 2]
    expected[4, 1] = 3.0 * full_matmul[4, 1]
    c = rt.ranked_memref_to_numpy(mem_out[0])
    if np.allclose(c, expected):
        pass
    else:
        quit(f'FAILURE')
Ejemplo n.º 3
0
def build_compile_and_run_SpMM(attr: st.EncodingAttr, support_lib: str,
                               compiler):
    # Build.
    module = build_SpMM(attr)
    func = str(module.operation.regions[0].blocks[0].operations[0].operation)
    module = ir.Module.parse(func + boilerplate(attr))

    # Compile.
    compiler(module)
    engine = execution_engine.ExecutionEngine(module,
                                              opt_level=0,
                                              shared_libs=[support_lib])

    # Set up numpy input and buffer for output.
    a = np.array(
        [[1.1, 0.0, 0.0, 1.4], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 3.3, 0.0]],
        np.float64)
    b = np.array([[1.0, 2.0], [4.0, 3.0], [5.0, 6.0], [8.0, 7.0]], np.float64)
    c = np.zeros((3, 2), np.float64)

    mem_a = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(a)))
    mem_b = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(b)))
    mem_c = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(c)))
    # Allocate a MemRefDescriptor to receive the output tensor.
    # The buffer itself is allocated inside the MLIR code generation.
    ref_out = rt.make_nd_memref_descriptor(2, ctypes.c_double)()
    mem_out = ctypes.pointer(ctypes.pointer(ref_out))

    # Invoke the kernel and get numpy output.
    # Built-in bufferization uses in-out buffers.
    # TODO: replace with inplace comprehensive bufferization.
    engine.invoke('main', mem_out, mem_a, mem_b, mem_c)

    # Sanity check on computed result.
    expected = np.matmul(a, b)
    c = rt.ranked_memref_to_numpy(mem_out[0])
    if np.allclose(c, expected):
        pass
    else:
        quit(f'FAILURE')