tensorcore support

Makefile

@@ -507,6 +507,7 @@ SOURCE_FILES = \
   Error.cpp \
   Expr.cpp \
   ExtractTileOperations.cpp \
+  ExtractWMMAOperations.cpp \
   FastIntegerDivide.cpp \
   FindCalls.cpp \
   FindIntrinsics.cpp \
@@ -709,6 +710,7 @@ HEADER_FILES = \
   Extern.h \
   ExternFuncArgument.h \
   ExtractTileOperations.h \
+  ExtractWMMAOperations.h \
   FastIntegerDivide.h \
   FindCalls.h \
   FindIntrinsics.h \

apps/cuda_mat_mul/CMakeLists.txt

@@ -27,12 +27,53 @@ find_package(Halide REQUIRED)
 # Generator
 add_halide_generator(mat_mul.generator SOURCES mat_mul_generator.cpp)
 
-# Filters
-add_halide_library(mat_mul FROM mat_mul.generator FEATURES cuda cuda_capability_50 PARAMS size=1024)
+# Filters. The operand type picks the schedule: half precision gets the tensor
+# cores, which need compute capability 7.0 or above, and float gets a schedule
+# that accumulates in ordinary registers.
+add_halide_library(
+    mat_mul
+    FROM mat_mul.generator
+    FEATURES cuda cuda_capability_50
+    PARAMS size=1024 A.type=float32 B.type=float32 out.type=float32
+)
+add_halide_library(
+    mat_mul_f16
+    FROM mat_mul.generator
+    GENERATOR mat_mul
+    FEATURES cuda cuda_capability_80
+    PARAMS size=1024 A.type=float16 B.type=float16 out.type=float32
+)
+
+add_halide_library(
+    mat_mul_f16_acc16
+    FROM mat_mul.generator
+    GENERATOR mat_mul
+    FEATURES cuda cuda_capability_80
+    PARAMS size=1024 A.type=float16 B.type=float16 out.type=float16
+)
+add_halide_library(
+    mat_mul_bf16
+    FROM mat_mul.generator
+    GENERATOR mat_mul
+    FEATURES cuda cuda_capability_80
+    PARAMS size=1024 A.type=bfloat16 B.type=bfloat16 out.type=float32
+)
+add_halide_library(
+    mat_mul_u8
+    FROM mat_mul.generator
+    GENERATOR mat_mul
+    FEATURES cuda cuda_capability_80
+    PARAMS size=1024 A.type=uint8 B.type=uint8 out.type=int32
+)
 
 # Main executable
 add_executable(runner runner.cpp)
-target_link_libraries(runner PRIVATE mat_mul Halide::Tools CUDA::cudart CUDA::cublas)
+target_link_libraries(
+    runner
+    PRIVATE
+    mat_mul mat_mul_f16 mat_mul_f16_acc16 mat_mul_bf16 mat_mul_u8 Halide::Halide Halide::Tools
+    CUDA::cudart CUDA::cublas
+)
 
 # Test that the app actually works!
 add_test(NAME mat_mul COMMAND runner)

apps/cuda_mat_mul/Makefile

@@ -2,11 +2,16 @@ include ../support/Makefile.inc
 
 MATRIX_SIZE ?= 1024
 
-CUDA_SDK ?= /usr/local/cuda-10.0
+CUDA_SDK ?= /usr/local/cuda
 
 CXXFLAGS += -I $(CUDA_SDK)/include
 LDFLAGS += -L $(CUDA_SDK)/lib64 -Wl,-rpath,$(CUDA_SDK)/lib64
 
+# The float variant runs anywhere, but the half variant is scheduled onto the
+# tensor cores, which need compute capability 7.0 or above.
+FLOAT_TARGET ?= host-cuda-cuda_capability_50
+HALF_TARGET ?= host-cuda-cuda_capability_80
+
 all: $(BIN)/$(HL_TARGET)/runner
 
 $(GENERATOR_BIN)/mat_mul.generator: mat_mul_generator.cpp $(GENERATOR_DEPS)
@@ -15,11 +20,39 @@ $(GENERATOR_BIN)/mat_mul.generator: mat_mul_generator.cpp $(GENERATOR_DEPS)
 
 $(BIN)/%/mat_mul.a: $(GENERATOR_BIN)/mat_mul.generator
 	@mkdir -p $(@D)
-	$^ -g mat_mul -e $(GENERATOR_OUTPUTS) -o $(@D) target=host-cuda-cuda_capability_50 size=$(MATRIX_SIZE)
+	$^ -g mat_mul -f mat_mul -e $(GENERATOR_OUTPUTS) -o $(@D) \
+	    target=$(FLOAT_TARGET) size=$(MATRIX_SIZE) \
+	    A.type=float32 B.type=float32 out.type=float32
+
+$(BIN)/%/mat_mul_f16.a: $(GENERATOR_BIN)/mat_mul.generator
+	@mkdir -p $(@D)
+	$^ -g mat_mul -f mat_mul_f16 -e $(GENERATOR_OUTPUTS) -o $(@D) \
+	    target=$(HALF_TARGET) size=$(MATRIX_SIZE) \
+	    A.type=float16 B.type=float16 out.type=float32
+
+$(BIN)/%/mat_mul_f16_acc16.a: $(GENERATOR_BIN)/mat_mul.generator
+	@mkdir -p $(@D)
+	$^ -g mat_mul -f mat_mul_f16_acc16 -e $(GENERATOR_OUTPUTS) -o $(@D) \
+	    target=$(HALF_TARGET) size=$(MATRIX_SIZE) \
+	    A.type=float16 B.type=float16 out.type=float16
+
+$(BIN)/%/mat_mul_bf16.a: $(GENERATOR_BIN)/mat_mul.generator
+	@mkdir -p $(@D)
+	$^ -g mat_mul -f mat_mul_bf16 -e $(GENERATOR_OUTPUTS) -o $(@D) \
+	    target=$(HALF_TARGET) size=$(MATRIX_SIZE) \
+	    A.type=bfloat16 B.type=bfloat16 out.type=float32
+
+$(BIN)/%/mat_mul_u8.a: $(GENERATOR_BIN)/mat_mul.generator
+	@mkdir -p $(@D)
+	$^ -g mat_mul -f mat_mul_u8 -e $(GENERATOR_OUTPUTS) -o $(@D) \
+	    target=$(HALF_TARGET) size=$(MATRIX_SIZE) \
+	    A.type=uint8 B.type=uint8 out.type=int32
 
-$(BIN)/%/runner: runner.cpp $(BIN)/%/mat_mul.a
+$(BIN)/%/runner: runner.cpp $(BIN)/%/mat_mul.a $(BIN)/%/mat_mul_f16.a \
+                 $(BIN)/%/mat_mul_bf16.a $(BIN)/%/mat_mul_u8.a \
+                 $(BIN)/%/mat_mul_f16_acc16.a
 	@mkdir -p $(@D)
-	$(CXX) $(CXXFLAGS) -I$(BIN)/$* -Wall $^ -o $@ $(LDFLAGS) -lcudart -lcublas
+	$(CXX) $(CXXFLAGS) -I$(BIN)/$* -Wall $^ -o $@ $(LDFLAGS) $(LIBHALIDE_LDFLAGS) -lcudart -lcublas
 
 test: $(BIN)/$(HL_TARGET)/runner
 	$^ $(MATRIX_SIZE)