[WIP] [CodeGen] Fix fp16 -> fp32 codegen on X86-64

[ajtulloch]

LLVM fails to generate calls to AVX/AVX512 variants of vcvtph2ps by default, which is inconvenient.

This adds a pattern matches (I'll generalize to arbitrary vector length) which dispatches to the AVX/AVX512 intrinsic functions when available.

This speeds this simple script up by approximately ~5x on my Haswell Core i7.

import tvm
import numpy as np
# Build a simple op that does 8-wide vectorized Float16 to Float32 conversion.

N = tvm.var("N")
V = 8
X = tvm.placeholder((N, V), name="X", dtype="float16")

def tvm_from_fp16(x):
    return x.astype("float32")

Y = tvm.compute(X.shape, lambda *i: tvm_from_fp16(X(*i)))

s = tvm.create_schedule(Y.op)
s[Y].vectorize(s[Y].op.axis[1])

target = tvm.target.create("llvm -mcpu=core-avx2")

with target:
    func = tvm.build(s, [X, Y])
    # print(func.get_source("asm"))
    # print(func.get_source("ll"))
    te = func.time_evaluator(func.entry_name, ctx=tvm.cpu(), min_repeat_ms=1000)
    x_np = tvm.nd.array(np.random.randn(10000, 8).astype(np.float16))
    y_np = tvm.nd.array(np.random.randn(10000, 8).astype(np.float32))
    print("Before: {:.2f}us".format(te(x_np, y_np).mean * 10 ** 6))
    np.testing.assert_equal(x_np.asnumpy().astype(np.float32), y_np.asnumpy())
    
target = tvm.target.create("llvm -mcpu=core-avx2 -mattr=+avx2,+f16c")

with target:
    func = tvm.build(s, [X, Y])
    # print(func.get_source("asm"))
    # print(func.get_source("ll"))
    te = func.time_evaluator(func.entry_name, ctx=tvm.cpu(), min_repeat_ms=1000)
    x_np = tvm.nd.array(np.random.randn(10000, 8).astype(np.float16))
    y_np = tvm.nd.array(np.random.randn(10000, 8).astype(np.float32))
    print("After: {:.2f}us".format(te(x_np, y_np).mean * 10 ** 6))
    np.testing.assert_equal(x_np.asnumpy().astype(np.float32), y_np.asnumpy())


# Before: 74.09us
# After: 14.90us

[ajtulloch]

@tqchen, @yidawang, @yzhliu, this may be of interest to you folks.

One question - is it possible to add unit tests that only execute when the underlying machine supports architectural features? It would be convenient to add a unit test for this that runs on machines with AVX or AVX-512, but I don't know how to express that in the current testing infrastructure.

[hlu1]

@ajtulloch, cpu-info can detect AVX2 or AVX512: https://github.com/pytorch/cpuinfo/blob/40c5f3695b053e5c3d642d9bc34113f3baa71ef2/include/cpuinfo.h#L1009

[yidawang]

LGTM in general. Just have some questions, mostly for my own educational purpose :)

[yidawang]

how time flies

[yidawang]

AVX-512

[yidawang]

I don't get the logic here. If lanes==8, why could has_avx512f be true?

[ajtulloch]

lanes == 8 is a property of the Expr we are compiling, has_avx512f is a property of the TargetMachine we are generating code for.

[yidawang]

Does has_avx512f==true imply has_f16c==true? And, is there a case that lanes==16 but has_avx512f==false? I am not familiar with the F16C instruction set.

[yidawang]

So the mattr flag needs to be set by the users?

[ajtulloch]

I've fixed this to query the TargetMachine directly, so no changes are needed (assuming users are using the existing -mcpu=core-avx2, -mcpu=skylake-avx512, etc)

[yidawang]

How is fp16 handled without F16C?

[ajtulloch]

Without f16c (i.e. vector/scalar vcvtph2ps), we invoke compiler-rt's builtins i.e. https://github.com/llvm-mirror/compiler-rt/blob/7a739a0dfb6d408c6d587e5c7b52abd89fc3fdd3/lib/builtins/fp_extend_impl.inc#L40.

https://godbolt.org/z/23tsVK is an example.

[ajtulloch]

@hlu1

@ajtulloch, cpu-info can detect AVX2 or AVX512: https://github.com/pytorch/cpuinfo/blob/40c5f3695b053e5c3d642d9bc34113f3baa71ef2/include/cpuinfo.h#L1009

The problem with that is that this only works if we are generating code for the current host architecture, which isn't what we want. I think the right way here is to continue to pass -mcpu (or -mattr) to LLVM when we want to generate architecture-specific code, and then we can query the constructed llvm::TargetMachine instance to see if it has the supplied feature (which is how LLVM chooses how to select instructions already, so we're guaranteed to be consistent). This is how Rust, etc implements this kind of thing (rust-lang/rust#31709). We can pass -mattr=-f16c,-avx512f, etc if we want to disable this optimization.

[tqchen]

Thanks @ajtulloch , @yidawang @hlu1 this is now merged