diff --git a/Makefile b/Makefile index 46996ac81b50..5a47aa20d8da 100644 --- a/Makefile +++ b/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 \ diff --git a/apps/cuda_mat_mul/CMakeLists.txt b/apps/cuda_mat_mul/CMakeLists.txt index 803f28c5ecdb..734274d1a8c1 100644 --- a/apps/cuda_mat_mul/CMakeLists.txt +++ b/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) diff --git a/apps/cuda_mat_mul/Makefile b/apps/cuda_mat_mul/Makefile index e0dfb78900fe..5f25604d9a02 100644 --- a/apps/cuda_mat_mul/Makefile +++ b/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) diff --git a/apps/cuda_mat_mul/mat_mul_generator.cpp b/apps/cuda_mat_mul/mat_mul_generator.cpp index 6f2cb17c8cd6..02b9ad95b45e 100644 --- a/apps/cuda_mat_mul/mat_mul_generator.cpp +++ b/apps/cuda_mat_mul/mat_mul_generator.cpp @@ -1,5 +1,7 @@ #include "Halide.h" +#include + using namespace Halide; namespace { @@ -12,61 +14,305 @@ void set_alignment_and_bounds(OutputImageParam p, int size) { .set_stride(size); } +// A square matrix multiply. The operands are untyped, so their type is a +// generator param, and it picks the schedule: the tensor cores multiply +// halves, bfloats and bytes, and anything else accumulates in ordinary +// registers. The output type is the accumulator type, so it picks between the +// accumulators an operand type can pair with. +// +// GFlop/s on an RTX 5060 Ti, against cublas and against the ceiling of the +// instructions used. Rows within a pair of types are comparable; rows in +// different pairs are not, because narrower operands or a narrower +// accumulator are less work. +// +// 1024 2048 4096 ceiling +// Halide f32 6937 10199 7306 25960 +// cublas f32 14688 16782 17545 +// +// Halide f16 -> f32 40396 47019 49314 51541 +// cublas f16 -> f32 42324 49596 51212 +// +// Halide bf16 -> f32 40347 47007 49129 51541 +// cublas bf16 -> f32 42315 49601 51179 +// +// Halide f16 -> f16 60745 87075 86973 99626 +// cublas f16 -> f16 73466 76443 90182 +// +// Halide u8 -> i32 63588 83223 90390 100650 +// cublas u8 -> i32 120109 128990 140195 +// +// The tensor core ceilings are measured, by issuing wmma instructions back to +// back out of registers. The float one is 36 SMs times the 2817 MHz this part +// averages while benchmarking times the 256 flops per SM per clock the cuda +// cores do. cublas is asked for a beta of zero, so that it writes its output +// rather than accumulating onto it, which is what the filters here do. +// +// The tensor core schedules reach 87% to 96% of their ceilings, and beat +// cublas at f16 -> f16 at 2048. Two rows fall short for reasons outside the +// schedule. +// +// At eight bits cublas is 39% past the ceiling, so it is not using wmma, which +// multiplies bytes no faster than it multiplies halves into halves. The mma +// instructions reach 188355 GOP/s at the same shape, 1.87x, for the same 8192 +// ops per instruction. Reaching that needs mma's fragment layout, which this +// schedule cannot express. +// +// The float row is at 28% because sharing a staged panel across a block needs +// the reduction chunk loop above the thread loops and the accumulator below +// them, living across it. Such an accumulator lands at block level, where a +// Register allocation is sized for the whole block tile and spills. So it +// stages per thread and shares nothing. The tensor core schedules only escape +// this because a tile allocation at block level is already per-lane. +// class MatMul : public Halide::Generator { public: GeneratorParam size{"size", 1024}; - Input> A{"A"}; - Input> B{"B"}; - Output> out{"out"}; + // How many tensor core tiles of accumulator each warp holds, and how many + // warps per block in each dimension. Zero means use the measured shapes + // below. + GeneratorParam tiles_x{"tiles_x", 0}; + GeneratorParam tiles_y{"tiles_y", 0}; + GeneratorParam warps_x{"warps_x", 0}; + GeneratorParam warps_y{"warps_y", 0}; + // How much of the reduction is staged in shared memory at a time. Zero + // means use the depth that goes with the shape below. + GeneratorParam block_r{"block_r", 0}; + // Extra elements per row of the shared panels, to spread consecutive rows + // across banks. Zero means sixteen bytes, the least that keeps each row + // aligned for both the widest asynchronous copy and the tensor core loads. + // That is a different number of elements per operand type. + GeneratorParam pad_a{"pad_a", 0}; + GeneratorParam pad_b{"pad_b", 0}; + + Input> A{"A"}; + Input> B{"B"}; + + // The output type is the accumulator type - there is no tensor core store + // from a half fragment into single precision memory. A half accumulator + // halves the registers it takes, but loses accuracy over a long reduction, + // so only ask for it when the numerics allow. + Output> out{"out"}; + + // Asking for a type the tensor cores multiply is what asks for them. + bool use_tensor_cores() const { + Type t = A.type(); + return t == Float(16) || t == BFloat(16) || t == Int(8) || t == UInt(8); + } + + // The accumulator each operand type pairs with. Halves can also + // accumulate into halves, which is what asking for a half output does. + Type natural_accumulator() const { + Type t = A.type(); + if (t == Int(8) || t == UInt(8)) { + return Int(32); + } + return Float(32); + } void generate() { - // 688 us on an RTX 2060 - // cublas is 512 us on the same card + _halide_user_assert(A.type() == B.type()) + << "The two operands must have the same type, but they are " + << A.type() << " and " << B.type() << ".\n"; + _halide_user_assert(out.type() == natural_accumulator() || + (out.type() == Float(16) && A.type() == Float(16))) + << "A " << A.type() << " matrix multiply accumulates into " + << natural_accumulator() + << (A.type() == Float(16) ? " or float16" : "") + << ", but a " << out.type() << " output was asked for.\n"; + r = RDom(0, size, "r"); - Var x("x"), y("y"), p("p"); + Type acc = out.type(); + prod(x, y) = cast(acc, 0); + // The widening happens here, at the multiply, rather than in the + // operand wrappers the schedule stages through shared memory, so that + // narrow operands are staged narrow and reach the tensor cores so. + prod(x, y) += cast(acc, A(x, r)) * cast(acc, B(r, y)); - Func prod("prod"); - RDom r(0, size); - prod(x, y) += A(x, r) * B(r, y); out(x, y) = prod(x, y); + } - Var xi, yi, xio, xii, yii, xo, yo, x_pair, xiio, ty; - RVar rxo, rxi; - - if (!using_autoscheduler()) { - out.bound(x, 0, size) - .bound(y, 0, size) - .tile(x, y, xi, yi, 64, 16) - .tile(xi, yi, xii, yii, 4, 8) - .gpu_blocks(x, y) - .gpu_threads(xi, yi) - .unroll(xii) - .unroll(yii); - prod.compute_at(out, xi) - .vectorize(x) - .unroll(y) - .update() - .reorder(x, y, r) - .vectorize(x) - .unroll(y) - .unroll(r, 8); - A.in().compute_at(prod, r).vectorize(_0).unroll(_1); - B.in().compute_at(prod, r).vectorize(_0).unroll(_1); - - set_alignment_and_bounds(A, size); - set_alignment_and_bounds(B, size); - set_alignment_and_bounds(out, size); - } else { + void schedule() { + if (using_autoscheduler()) { A.dim(0).set_estimate(0, size).dim(1).set_estimate(0, size); B.dim(0).set_estimate(0, size).dim(1).set_estimate(0, size); + out.bound(x, 0, size).bound(y, 0, size); + return; + } + + set_alignment_and_bounds(A, size); + set_alignment_and_bounds(B, size); + set_alignment_and_bounds(out, size); + + out.bound(x, 0, size).bound(y, 0, size); + + if (use_tensor_cores()) { + schedule_tensor_cores(); + } else { + schedule_cuda(); + } + } + +private: + // See the table above. + void schedule_cuda() { + Var xi, yi, xii, yii; + + out.tile(x, y, xi, yi, 64, 16) + .tile(xi, yi, xii, yii, 4, 8) + .gpu_blocks(x, y) + .gpu_threads(xi, yi) + .unroll(xii) + .unroll(yii); + + prod.compute_at(out, xi) + .vectorize(x) + .unroll(y) + .update() + .reorder(x, y, r) + .vectorize(x) + .unroll(y) + .unroll(r, 8); + + A.in().compute_at(prod, r).vectorize(_0).unroll(_1); + B.in().compute_at(prod, r).vectorize(_0).unroll(_1); + } + + // See the table above. + void schedule_tensor_cores() { + // Each operand tile loaded feeds tiles_x (or tiles_y) multiplies, so + // the tile counts are what get reuse out of the loads. + const int tile = 16; + int tx = tiles_x, ty = tiles_y, wx = warps_x, wy = warps_y; + // The staging depth goes with the shape below unless asked for. + int br = 32; + if (tx == 0 || ty == 0 || wx == 0 || wy == 0) { + // Measured on an RTX 5060 Ti by sweeping shapes at each size and + // operand type. There is no trend worth extrapolating from, so + // these are measured points. The operand type matters as much as + // the size: bytes want a tall block and a deep staged panel, where + // the 16-bit types want a wide block and a shallow one. Bfloats pick + // out the same shapes as halves at every size, so they share a row. + const bool bytes = A.type().bits() == 8; + const bool half_accumulator = out.type() == Float(16); + if (bytes) { + if ((int)size <= 1024) { + tx = 2, ty = 8, wx = 2, wy = 1, br = 64; + } else if ((int)size <= 2048) { + tx = 2, ty = 10, wx = 2, wy = 1, br = 32; + } else { + tx = 2, ty = 8, wx = 2, wy = 1, br = 64; + } + } else if (half_accumulator) { + if ((int)size <= 1024) { + tx = 4, ty = 4, wx = 2, wy = 1; + } else if ((int)size <= 2048) { + tx = 5, ty = 4, wx = 2, wy = 2; + } else { + tx = 4, ty = 4, wx = 2, wy = 2; + } + } else { + if ((int)size <= 1024) { + tx = 8, ty = 2, wx = 1, wy = 1; + } else if ((int)size <= 2048) { + tx = 5, ty = 4, wx = 1, wy = 2; + } else { + tx = 4, ty = 2, wx = 2, wy = 2; + } + } + } + if (block_r) { + br = block_r; } + const int block_x = tile * tx * wx; + const int block_y = tile * ty * wy; - // Always specify bounds for outputs, whether autoscheduled or not - out - .bound(x, 0, size) - .bound(y, 0, size); + const int pa = pad_a ? (int)pad_a : 16 / A.type().bytes(); + const int pb = pad_b ? (int)pad_b : 16 / A.type().bytes(); + + // out and prod are tiled the same way: blocks of warps, each warp + // holding several tensor core tiles. mmx, mmy and mmr are the + // dimensions of one tensor core operation. + Var xi("xi"), yi("yi"), xw("xw"), yw("yw"); + Var mmx("mmx"), mmy("mmy"); + RVar ro("ro"), ri("ri"), mmr("mmr"); + + out.tile(x, y, xi, yi, block_x, block_y) + .tile(xi, yi, xw, yw, xi, yi, tile * tx, tile * ty) + .tile(xi, yi, mmx, mmy, tile, tile) + .gpu_blocks(x, y) + .gpu_threads(xw, yw) + .unroll(xi) + .unroll(yi) + .tile_store(mmx, mmy); + + // The accumulators live in tensor core registers for the whole + // reduction, and are written out to memory once at the end. They sit + // at block level so that the reduction loop can be above the loop over + // warps, which lets every warp share one staged panel. + prod.compute_at(out, x) + .tile(x, y, xw, yw, xi, yi, tile * tx, tile * ty) + .tile(xi, yi, mmx, mmy, tile, tile) + .gpu_threads(xw, yw) + .tile_init(mmx, mmy) + .unroll(xi) + .unroll(yi); + + prod.update() + .split(r, ro, ri, br) + .split(ri, ri, mmr, tile) + .tile(x, y, xw, yw, xi, yi, tile * tx, tile * ty) + .tile(xi, yi, mmx, mmy, tile, tile) + .reorder(xi, yi, ri, xw, yw, ro) + .gpu_threads(xw, yw) + .unroll(xi) + .unroll(yi) + .unroll(ri) + .tile_matmul(mmr, mmx, mmy); + + // Stage the operand panels into shared memory once per reduction step, + // to be shared by every warp in the block. Each thread moves sixteen + // bytes at a time along the dense dimension, so that the reads from + // global memory coalesce and the writes to shared memory can be done + // as asynchronous copies. + // Each thread moves sixteen bytes, the widest asynchronous copy the + // hardware has. How many elements that is depends on the operand type. + const int vec = 16 / A.type().bytes(); + Var ko("ko"), kv("kv"), xo("xo"), xv("xv"); + Var t("t"), ti("ti"), to("to"); + + // B.in() is dense in the reduction dimension, which is its _0. + B.in() + .compute_at(prod, ro) + .store_in(MemoryType::GPUSharedAsync) + .align_storage(_0, br + pa) + .split(_0, ko, kv, vec) + .fuse(ko, _1, t) + .split(t, t, ti, 32) + .split(t, t, xw, wx) + .split(t, to, yw, wy) + .gpu_lanes(ti) + .gpu_threads(xw, yw) + .vectorize(kv); + + // A.in() is dense in x, which is its _0. + A.in() + .compute_at(prod, ro) + .store_in(MemoryType::GPUSharedAsync) + .align_storage(_0, block_x + pb) + .split(_0, xo, xv, vec) + .fuse(xo, _1, t) + .split(t, t, ti, 32) + .split(t, t, xw, wx) + .split(t, to, yw, wy) + .gpu_lanes(ti) + .gpu_threads(xw, yw) + .vectorize(xv); } + + Var x{"x"}, y{"y"}; + RDom r; + Func prod{"prod"}; }; } // namespace diff --git a/apps/cuda_mat_mul/runner.cpp b/apps/cuda_mat_mul/runner.cpp index 898496632802..c5df75a578b3 100644 --- a/apps/cuda_mat_mul/runner.cpp +++ b/apps/cuda_mat_mul/runner.cpp @@ -1,17 +1,143 @@ +// TODO: this links libHalide purely to get Halide::float16_t and +// Halide::bfloat16_t, which an AOT program should not have to do. Float16.h +// depends on nothing but HalideRuntime.h and already declares halide_type_of +// for both types, but its bodies live in Float16.cpp inside libHalide, and the +// header is not distributed. Making those definitions inline and shipping the +// header, or moving the types to the runtime, would let this go back to +// linking only the generated code. +#include "Halide.h" + #include "HalideBuffer.h" #include "HalideRuntimeCuda.h" #include "halide_benchmark.h" -#include "mat_mul.h" #include #include +#include #include +#include "mat_mul.h" +#include "mat_mul_bf16.h" +#include "mat_mul_f16.h" +#include "mat_mul_f16_acc16.h" +#include "mat_mul_u8.h" + +using Halide::bfloat16_t; +using Halide::float16_t; using Halide::Runtime::Buffer; -using Halide::Tools::benchmark; + +namespace { + +// The same matrix multiply is compiled from one generator at each pair of +// operand and accumulator types, and each is compared against cublas doing the +// same thing. Both are handed the same device buffers - the Halide side fills +// them and copies them down, and cublas is given the pointers out of them - so +// the two see identical data, and both answers are checked. +// +// The operands are small integers so that every dot product is exact in every +// accumulator here, which lets the results be checked for equality rather than +// to a tolerance. That makes them unusually compressible, so the numbers were +// checked against dense random operands too: no configuration moved by more +// than a couple of percent, which is what you would expect of a multiply that +// is issue-bound rather than waiting on memory. + +// Time one call. Both implementations queue work asynchronously, and cublas +// does a heuristic lookup on the host for every call, so synchronizing per +// launch would measure that host work rather than letting it overlap with the +// GPU. Batch the launches instead, and sync once. `sync` has to match the +// launcher: Halide runs on its own CUDA context, so cudaDeviceSynchronize does +// not wait for it. +template +double bench(F &&launch, S &&sync) { + const int batch = 5; + return Halide::Tools::benchmark(5, 1, + [&]() { + for (int i = 0; i < batch; i++) { + launch(); + } + sync(); + }) / + batch; +} + +double gflops(int size, double seconds) { + return 2.0 * size * size * size / seconds * 1e-9; +} + +template +void fill_ints(Buffer &b, int modulus) { + b.for_each_value([&](T &v) { v = (T)(rand() % modulus); }); + b.set_host_dirty(); +} + +// Spot check on strides coprime with the tile sizes, so the samples land at +// varying offsets within a tile. +template +bool check(Buffer &Ab, Buffer &Bb, Buffer &Cb, + int size, const char *name) { + for (int y = 0; y < size; y += 97) { + for (int x = 0; x < size; x += 89) { + double correct = 0; + for (int k = 0; k < size; k++) { + correct += (double)Ab(x, k) * (double)Bb(k, y); + } + if ((double)Cb(x, y) != correct) { + printf("%s: bad result at %d %d: %f != %f\n", + name, x, y, (double)Cb(x, y), correct); + return false; + } + } + } + return true; +} + +// One row of the table: run the Halide filter, check it, time it, then time +// cublas on the same device buffers. +template +bool row(const char *name, int size, int modulus, Filter filter, Cublas cublas) { + Buffer Ab(size, size), Bb(size, size); + Buffer Cb(size, size); + fill_ints(Ab, modulus); + fill_ints(Bb, modulus); + + if (filter(Ab.raw_buffer(), Bb.raw_buffer(), Cb.raw_buffer()) != 0) { + printf("%s: filter returned an error\n", name); + return false; + } + Cb.copy_to_host(); + if (!check(Ab, Bb, Cb, size, name)) { + return false; + } + double t = bench([&]() { filter(Ab.raw_buffer(), Bb.raw_buffer(), Cb.raw_buffer()); }, + [&]() { Cb.device_sync(); }); + printf(" Halide %-12s %9.0f GFlop/s\n", name, gflops(size, t)); + + // Running the filter left everything on the device, so cublas gets the + // same buffers, output included. Halide's dense-first-dimension layout is + // what cublas calls column major with a leading dimension of the size, so + // no transposing is needed, and the answer is checked the same way to + // confirm that. + void *Ad = (void *)halide_cuda_get_device_ptr(nullptr, Ab.raw_buffer()); + void *Bd = (void *)halide_cuda_get_device_ptr(nullptr, Bb.raw_buffer()); + void *Cd = (void *)halide_cuda_get_device_ptr(nullptr, Cb.raw_buffer()); + cublas(Ad, Bd, Cd, size); + cudaDeviceSynchronize(); + Cb.set_device_dirty(); + Cb.copy_to_host(); + if (!check(Ab, Bb, Cb, size, name)) { + printf(" (that was cublas, not Halide)\n"); + return false; + } + t = bench([&]() { cublas(Ad, Bd, Cd, size); }, + []() { cudaDeviceSynchronize(); }); + printf(" cublas %-12s %9.0f GFlop/s\n", name, gflops(size, t)); + return true; +} + +cublasHandle_t handle; + +} // namespace int main(int argc, char **argv) { - // Our Generator is compiled using cuda_capability_50; if the system running this - // test doesn't have at least that, quietly skip the test. const auto *interface = halide_cuda_device_interface(); assert(interface->compute_capability != nullptr); int major, minor; @@ -28,66 +154,77 @@ int main(int argc, char **argv) { size = atoi(argv[1]); } - // Check correctness using small-integer matrices - if (1) { - Buffer A(size, size), B(size, size), C(size, size); - A.for_each_value([](float &v) { v = (rand() & 3) - 1; }); - B.for_each_value([](float &v) { v = (rand() & 3) - 1; }); - A.set_host_dirty(); - B.set_host_dirty(); - mat_mul(A, B, C); - C.copy_to_host(); - for (int y = 0; y < size; y++) { - for (int x = 0; x < size; x++) { - float correct = 0.f; - for (int k = 0; k < size; k++) { - correct += A(x, k) * B(k, y); - } - float actual = C(x, y); - if (correct != actual) { - printf("%d %d: %f vs %f\n", x, y, correct, actual); - return -1; - } - } - } - } + cublasCreate(&handle); + static float alpha = 1.0f, beta = 0.0f; + int failures = 0; - // Benchmark it - { - Buffer A(size, size), B(size, size), C(size, size); - double t = Halide::Tools::benchmark(5, 5, [&]() { - mat_mul(A, B, C); - C.device_sync(); - }); - printf("Halide time: %f\n", t); - } + // A gemm at one pair of types, given how cublas should read the buffers. + auto gemm_ex = [](cudaDataType at, cudaDataType ct, cublasComputeType_t comp, + const void *al, const void *be) { + return [=](void *A, void *B, void *C, int n) { + cublasGemmEx(handle, CUBLAS_OP_N, CUBLAS_OP_N, n, n, n, al, + A, at, n, B, at, n, be, C, ct, n, comp, + CUBLAS_GEMM_DEFAULT); + }; + }; - // Benchmark cublas -#ifdef _MSC_VER - // https://github.com/halide/Halide/issues/5053 - printf("Skipping cublas on Windows; see https://github.com/halide/Halide/issues/5053\n"); -#else - { - float *A, *B, *C; - cudaMalloc((void **)&A, size * size * 4); - cudaMalloc((void **)&B, size * size * 4); - cudaMalloc((void **)&C, size * size * 4); - cublasHandle_t handle; - cublasCreate(&handle); - float alpha = 1.0f, beta = 1.0f; - double t = Halide::Tools::benchmark(5, 5, [&]() { - cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, - size, size, size, &alpha, A, size, B, size, &beta, C, size); - cudaDeviceSynchronize(); + failures += !row( + "f32 -> f32", size, 4, + [](halide_buffer_t *A, halide_buffer_t *B, halide_buffer_t *C) { + return mat_mul(A, B, C); + }, + [](void *A, void *B, void *C, int n) { + cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, n, n, n, &alpha, + (const float *)A, n, (const float *)B, n, &beta, + (float *)C, n); }); - cudaFree(A); - cudaFree(B); - cudaFree(C); - cublasDestroy(handle); - printf("cublas time: %f\n", t); + + if (ver < 70) { + printf("[SKIP] Tensor cores require compute capability 7.0 or above; " + "this system has %d.%d.\n", + major, minor); + } else { + failures += !row( + "f16 -> f32", size, 4, + [](halide_buffer_t *A, halide_buffer_t *B, halide_buffer_t *C) { + return mat_mul_f16(A, B, C); + }, + gemm_ex(CUDA_R_16F, CUDA_R_32F, CUBLAS_COMPUTE_32F, &alpha, &beta)); + + failures += !row( + "bf16 -> f32", size, 4, + [](halide_buffer_t *A, halide_buffer_t *B, halide_buffer_t *C) { + return mat_mul_bf16(A, B, C); + }, + gemm_ex(CUDA_R_16BF, CUDA_R_32F, CUBLAS_COMPUTE_32F, &alpha, &beta)); + + // Zeros and ones, so that the dot products stay under 2048, the + // largest integer half precision represents exactly. + static __half halpha = __float2half(1.f), hbeta = __float2half(0.f); + failures += !row( + "f16 -> f16", size, 2, + [](halide_buffer_t *A, halide_buffer_t *B, halide_buffer_t *C) { + return mat_mul_f16_acc16(A, B, C); + }, + gemm_ex(CUDA_R_16F, CUDA_R_16F, CUBLAS_COMPUTE_16F, &halpha, &hbeta)); + + // cublas takes signed bytes where the Halide variant takes unsigned + // ones. The hardware runs both at the same rate, and these values are + // small enough to mean the same thing either way. + static int32_t ialpha = 1, ibeta = 0; + failures += !row( + "u8 -> i32", size, 4, + [](halide_buffer_t *A, halide_buffer_t *B, halide_buffer_t *C) { + return mat_mul_u8(A, B, C); + }, + gemm_ex(CUDA_R_8I, CUDA_R_32I, CUBLAS_COMPUTE_32I, &ialpha, &ibeta)); } -#endif + cublasDestroy(handle); + if (failures) { + printf("%d configuration(s) failed\n", failures); + return 1; + } printf("Success!\n"); return 0; } diff --git a/python_bindings/src/halide/halide_/PyEnums.cpp b/python_bindings/src/halide/halide_/PyEnums.cpp index 93978cc792a5..b42cfa2bada2 100644 --- a/python_bindings/src/halide/halide_/PyEnums.cpp +++ b/python_bindings/src/halide/halide_/PyEnums.cpp @@ -49,7 +49,11 @@ void define_enums(py::module &m) { .value("GPUShared", MemoryType::GPUShared) .value("GPUTexture", MemoryType::GPUTexture) .value("LockedCache", MemoryType::LockedCache) - .value("VTCM", MemoryType::VTCM); + .value("VTCM", MemoryType::VTCM) + .value("Tile", MemoryType::Tile) + .value("GPUSharedAsync", MemoryType::GPUSharedAsync) + // Deprecated alias for Tile. + .value("AMXTile", MemoryType::Tile); py::enum_(m, "NameMangling") .value("Default", NameMangling::Default) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index e5bfba34892a..38002130f8b2 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -117,6 +117,7 @@ target_sources( Extern.h ExternFuncArgument.h ExtractTileOperations.h + ExtractWMMAOperations.h FastIntegerDivide.h FindCalls.h FindIntrinsics.h @@ -299,6 +300,7 @@ target_sources( Error.cpp Expr.cpp ExtractTileOperations.cpp + ExtractWMMAOperations.cpp FastIntegerDivide.cpp FindCalls.cpp FindIntrinsics.cpp diff --git a/src/CanonicalizeGPUVars.cpp b/src/CanonicalizeGPUVars.cpp index 609323f8e5dd..4e02304a37ac 100644 --- a/src/CanonicalizeGPUVars.cpp +++ b/src/CanonicalizeGPUVars.cpp @@ -6,6 +6,7 @@ #include "IR.h" #include "IRMutator.h" #include "Substitute.h" +#include "Util.h" namespace Halide { namespace Internal { @@ -36,14 +37,21 @@ class CountGPUBlocksThreads : public IRVisitor { // Counters that track the number of blocks, threads, and lanes loops that // we're inside of, respectively. Lanes loops also count as threads loops. - int nb = 0, nt = 0, nl = 0; + int nb = 0, nt = 0, nl = 0, nto = 0; + + // Whether we're already inside a lane dimension. Every lane dimension maps + // to the innermost thread dimension, so one nested inside another is the + // same dimension, not a new one. + bool in_lanes = false; void visit(const For *op) override { // Figure out how much to increment each counter by based on the loop // type. int db = op->for_type == ForType::GPUBlock; - int dl = op->for_type == ForType::GPULane; + int dl = (op->for_type == ForType::GPULane) && !in_lanes; int dt = op->for_type == ForType::GPUThread; + ScopedValue old_in_lanes(in_lanes, + in_lanes || op->for_type == ForType::GPULane); // The threads counter includes lanes loops dt += dl; @@ -52,11 +60,13 @@ class CountGPUBlocksThreads : public IRVisitor { nb += db; nl += dl; nt += dt; + nto += op->for_type == ForType::GPUThread; // Update the maximum counter values seen. nblocks = std::max(nb, nblocks); nthreads = std::max(nt, nthreads); nlanes = std::max(nl, nlanes); + nthreads_excluding_lanes = std::max(nto, nthreads_excluding_lanes); // Visit the body IRVisitor::visit(op); @@ -65,6 +75,23 @@ class CountGPUBlocksThreads : public IRVisitor { nb -= db; nl -= dl; nt -= dt; + nto -= op->for_type == ForType::GPUThread; + } + + void visit(const Realize *op) override { + // extract_wmma_operations will wrap the statements that touch this + // allocation in loops over the lanes of a warp, so count it as a lane + // dimension. + const bool wmma = op->memory_type == MemoryType::Tile; + int dl = wmma && !in_lanes; + ScopedValue old_in_lanes(in_lanes, in_lanes || wmma); + nl += dl; + nt += dl; + nlanes = std::max(nl, nlanes); + nthreads = std::max(nt, nthreads); + IRVisitor::visit(op); + nl -= dl; + nt -= dl; } public: @@ -73,13 +100,30 @@ class CountGPUBlocksThreads : public IRVisitor { int nblocks = 0; int nthreads = 0; int nlanes = 0; + // Threads not counting lanes, which is what a thread loop's depth is when + // the lane dimension comes from an enclosing tensor core allocation rather + // than from anything in this loop's body. + int nthreads_excluding_lanes = 0; }; class CanonicalizeGPUVars : public IRMutator { map gpu_vars; + // Whether we're inside a tensor core accumulator allocation. + // extract_wmma_operations will introduce a loop over the lanes of a warp + // somewhere inside it, which takes the innermost thread dimension, so the + // thread loops in here are one dimension further out than they look. + bool in_wmma_alloc = false; + using IRMutator::visit; + Stmt visit(const Realize *op) override { + ScopedValue old(in_wmma_alloc, + in_wmma_alloc || + op->memory_type == MemoryType::Tile); + return IRMutator::visit(op); + } + std::string find_replacement(const std::string &suffix, const std::string &name) { vector v = split_string(name, suffix); internal_assert(v.size() == 2); @@ -107,7 +151,9 @@ class CanonicalizeGPUVars : public IRMutator { name += gpu_block_name(counter.nblocks); debug(5) << "Replacing " << op->name << " with GPU block name " << name << "\n"; } else if (op->for_type == ForType::GPUThread) { - name += gpu_thread_name(counter.nthreads); + name += gpu_thread_name(in_wmma_alloc ? + counter.nthreads_excluding_lanes + 1 : + counter.nthreads); debug(5) << "Replacing " << op->name << " with GPU thread name " << name << "\n"; } else if (op->for_type == ForType::GPULane) { name += gpu_thread_name(0); diff --git a/src/CodeGen_D3D12Compute_Dev.cpp b/src/CodeGen_D3D12Compute_Dev.cpp index be20e97e6113..f58b43452271 100644 --- a/src/CodeGen_D3D12Compute_Dev.cpp +++ b/src/CodeGen_D3D12Compute_Dev.cpp @@ -1406,7 +1406,7 @@ void CodeGen_D3D12Compute_Dev::CodeGen_D3D12Compute_C::visit(const Select *op) { } bool is_shared_allocation(const Allocate *op) { - return op->memory_type == MemoryType::GPUShared; + return is_gpu_shared(op->memory_type); } void CodeGen_D3D12Compute_Dev::CodeGen_D3D12Compute_C::visit(const Allocate *op) { diff --git a/src/CodeGen_Metal_Dev.cpp b/src/CodeGen_Metal_Dev.cpp index 61fbb8602c10..08145925aa10 100644 --- a/src/CodeGen_Metal_Dev.cpp +++ b/src/CodeGen_Metal_Dev.cpp @@ -539,7 +539,7 @@ void CodeGen_Metal_Dev::CodeGen_Metal_C::visit(const Select *op) { void CodeGen_Metal_Dev::CodeGen_Metal_C::visit(const Allocate *op) { - if (op->memory_type == MemoryType::GPUShared) { + if (is_gpu_shared(op->memory_type)) { // Already handled op->body.accept(this); } else { @@ -776,7 +776,7 @@ void CodeGen_Metal_Dev::CodeGen_Metal_C::add_kernel(const Stmt &s, const Allocate *shared_alloc = nullptr; shared_name = "__shared"; visit_with(s, [&](auto *self, const Allocate *op) { - if (op->memory_type == MemoryType::GPUShared) { + if (is_gpu_shared(op->memory_type)) { internal_assert(shared_alloc == nullptr) << "Found multiple shared allocations in metal kernel\n"; shared_alloc = op; diff --git a/src/CodeGen_OpenCL_Dev.cpp b/src/CodeGen_OpenCL_Dev.cpp index 66acd48c08f6..88e858355f1e 100644 --- a/src/CodeGen_OpenCL_Dev.cpp +++ b/src/CodeGen_OpenCL_Dev.cpp @@ -780,7 +780,7 @@ void CodeGen_OpenCL_Dev::CodeGen_OpenCL_C::visit(const Allocate *op) { user_assert(!op->new_expr.defined()) << "Allocate node inside OpenCL kernel has custom new expression.\n" << "(Memoization is not supported inside GPU kernels at present.)\n"; - if (op->memory_type == MemoryType::GPUShared) { + if (is_gpu_shared(op->memory_type)) { // Already handled op->body.accept(this); } else { @@ -1063,7 +1063,7 @@ void CodeGen_OpenCL_Dev::CodeGen_OpenCL_C::add_kernel(Stmt s, const Allocate *shared_alloc = nullptr; shared_name = "__shared"; visit_with(s, [&](auto *self, const Allocate *op) { - if (op->memory_type == MemoryType::GPUShared) { + if (is_gpu_shared(op->memory_type)) { internal_assert(shared_alloc == nullptr) << "Found multiple shared allocations in metal kernel\n"; shared_alloc = op; diff --git a/src/CodeGen_PTX_Dev.cpp b/src/CodeGen_PTX_Dev.cpp index 8dcfc0016aa0..afed95155d23 100644 --- a/src/CodeGen_PTX_Dev.cpp +++ b/src/CodeGen_PTX_Dev.cpp @@ -1,4 +1,5 @@ #include "CodeGen_PTX_Dev.h" + #include "CSE.h" #include "CanonicalizeGPUVars.h" #include "CodeGen_GPU_Dev.h" @@ -7,6 +8,7 @@ #include "ConciseCasts.h" #include "Debug.h" #include "ExprUsesVar.h" +#include "ExtractWMMAOperations.h" #include "IREquality.h" #include "IRMatch.h" #include "IRMutator.h" @@ -14,11 +16,15 @@ #include "IRPrinter.h" #include "LLVM_Headers.h" #include "LLVM_Runtime_Linker.h" +#include "ModulusRemainder.h" +#include "MultiRamp.h" #include "Simplify.h" #include "Solve.h" +#include "Substitute.h" #include "Target.h" #include +#include namespace Halide { namespace Internal { @@ -80,6 +86,7 @@ class CodeGen_PTX_Dev : public CodeGen_LLVM, public CodeGen_GPU_Dev { void visit(const Load *) override; void visit(const Store *) override; void visit(const Atomic *) override; + void visit(const ProducerConsumer *) override; void codegen_vector_reduce(const VectorReduce *op, const Expr &init) override; // @} @@ -101,6 +108,62 @@ class CodeGen_PTX_Dev : public CodeGen_LLVM, public CodeGen_GPU_Dev { * intrinsic functions to call to get them. */ std::string simt_intrinsic(const std::string &name); + /** The memory type of each allocation made inside the kernel, so that + * copies into shared memory can be recognized. */ + Scope alloc_memory_type; + + /** Whether we're inside a producer node. */ + bool in_producer = false; + + /** The groups of asynchronous copies committed so far, oldest first, and + * the group of any copies issued since the last commit. Waits are FIFO - + * the hardware can only wait for all but the newest N groups - so a + * group's position here is what determines the N we emit for it. */ + std::vector committed_groups; + int uncommitted_group = -1; + + /** Try to emit a store into shared memory as an asynchronous copy, which + * moves the data straight from global memory without routing it through + * registers. Returns false if this store isn't one we can do that for. */ + bool codegen_async_copy(const Store *op, const char **reason); + + /** Close the current group of asynchronous copies, if there is one. */ + void commit_copies(); + + /** Emit a wait that leaves at most n groups of copies outstanding. */ + void emit_copy_wait(int n); + + /** Wait for the asynchronous copies in the given group to have landed. */ + void await_copies(int group); + + /** Wait for every asynchronous copy issued so far to have landed. */ + void await_all_copies(); + + /** Emit calls to the nvvm warp-level matrix multiply-accumulate + * intrinsics that drive the tensor cores. */ + // @{ + void codegen_wmma(const Call *op); + llvm::Value *codegen_wmma_raw(const Call *op); + void codegen_wmma_store(const Store *op); + void split_fragment(const Expr &e, std::vector &args); + // @} + + /** A fragment is a fixed number of 32-bit registers per lane. Keeping it in + * that form all the way to and from its allocation matters, because NVPTX + * holds a wide float vector in pairs of registers, and packing and + * unpacking one around every tensor core instruction costs more + * instructions than the instructions themselves. + */ + // @{ + bool is_fragment_alloc(const std::string &name); + llvm::Type *fragment_reg_type(Type t); + llvm::Value *fragment_reg_ptr(const std::string &name, Type t, int i); + void codegen_fragment_store(const Store *op); + llvm::Value *call_wmma_intrinsic(const std::string &name, + const std::vector &args, + const std::vector &overloads); + // @} + bool supports_atomic_add(const Type &t) const override; }; @@ -126,6 +189,34 @@ Type CodeGen_PTX_Dev::upgrade_type_for_storage(const Type &t) const { return CodeGen_LLVM::upgrade_type_for_storage(t); } +namespace { + +// The size of the thread block a kernel will be launched with. ptxas allocates +// registers on the assumption that a block may be as large as the hardware +// allows unless we tell it otherwise, which costs occupancy. +class BlockSize : public IRVisitor { + using IRVisitor::visit; + + void visit(const For *op) override { + for (int i = 0; i < 3; i++) { + if (ends_with(op->name, gpu_thread_name(i))) { + if (auto e = as_const_int(simplify(op->extent()))) { + extent[i] = std::max(extent[i], (int)*e); + } else { + known = false; + } + } + } + IRVisitor::visit(op); + } + +public: + int extent[3] = {1, 1, 1}; + bool known = true; +}; + +} // namespace + void CodeGen_PTX_Dev::add_kernel(Stmt stmt, const std::string &name, const std::vector &args) { @@ -204,6 +295,26 @@ void CodeGen_PTX_Dev::add_kernel(Stmt stmt, module->getOrInsertNamedMetadata("nvvm.annotations")->addOperand(md_node); + // Tell ptxas how large the thread block is. Without this it has to assume + // the largest block the hardware supports, and allocates registers for that + // rather than for the block we actually launch. + BlockSize block_size; + stmt.accept(&block_size); + if (block_size.known) { + const char *annotation[] = {"maxntidx", "maxntidy", "maxntidz"}; + for (int i = 0; i < 3; i++) { + llvm::Metadata *args[] = { + llvm::ValueAsMetadata::get(function), + MDString::get(*context, annotation[i]), + llvm::ValueAsMetadata::get(ConstantInt::get(i32_t, block_size.extent[i]))}; + module->getOrInsertNamedMetadata("nvvm.annotations") + ->addOperand(MDNode::get(*context, args)); + } + debug(2) << "Kernel " << name << " has block size " + << block_size.extent[0] << "x" << block_size.extent[1] + << "x" << block_size.extent[2] << "\n"; + } + // Now verify the function is ok verifyFunction(*function); @@ -264,6 +375,15 @@ void CodeGen_PTX_Dev::init_module() { } void CodeGen_PTX_Dev::visit(const Call *op) { + if (op->is_intrinsic(Call::cuda_await_copies)) { + internal_assert(op->args.size() == 1); + auto group = as_const_int(op->args[0]); + internal_assert(group) << "cuda_await_copies group is not a constant integer\n"; + await_copies((int)*group); + value = ConstantInt::get(i32_t, 0); + return; + } + if (op->is_intrinsic(Call::gpu_thread_barrier)) { // Even though we always insert a __syncthreads equivalent // (which has both a device and shared memory fence) @@ -271,6 +391,10 @@ void CodeGen_PTX_Dev::visit(const Call *op) { // arguments internal_assert(op->args.size() == 1) << "gpu_thread_barrier() intrinsic must specify memory fence type.\n"; + // A barrier tells other threads the shared memory this thread wrote is + // ready, so any asynchronous copies must have landed by now. + await_all_copies(); + auto fence_type_ptr = as_const_int(op->args[0]); internal_assert(fence_type_ptr) << "gpu_thread_barrier() parameter is not a constant integer.\n"; @@ -288,6 +412,16 @@ void CodeGen_PTX_Dev::visit(const Call *op) { return; } + if (is_wmma_intrinsic(op)) { + codegen_wmma(op); + return; + } + + internal_assert(!op->is_intrinsic(Call::wmma_fragment_to_matrix_d) && + !op->is_intrinsic(Call::wmma_lane_owns)) + << "A tensor core accumulator store was broken apart during lowering. " + << op->name << " only has meaning as part of one.\n"; + // TODO: It would be better if CodeGen_LLVM could handle overloaded intrin calls by default. value = call_overloaded_intrin(op->type, op->name, op->args); if (!value) { @@ -295,6 +429,281 @@ void CodeGen_PTX_Dev::visit(const Call *op) { } } +namespace { + +WMMAMatrixLayout matrix_in_memory(const string &name, const MultiRamp &mr, int rows, int cols, const Expr &access = Expr()) { + WMMAMatrixLayout result; + user_assert(wmma_matrix_layout(mr, rows, cols, &result)) + << "The memory a tensor core instruction moves a " << rows << "x" << cols + << " matrix of " << name << " to or from is not a dense tile by the time it " + << "reaches the backend. One cause is a shared memory allocation made inside " + << "the loop over GPU threads, which gets striped across them; compute it at " + << "a loop outside the threads instead. The addresses accessed are:\n" + << mr.to_expr() << "\nThe access is:\n" + << access << "\n"; + return result; +} + +} // namespace + +// The nvvm intrinsics take and return a fragment as 32-bit registers. Halide +// represents half precision as llvm's half, which is what the intrinsics for +// it take directly, but bfloat and the eight-bit integers have no llvm vector +// type in the signature - those intrinsics take the lanes packed into an i32. +bool wmma_reg_is_packed_i32(Type t) { + return t.bits() < 32 && t.element_of() != Float(16); +} + +void CodeGen_PTX_Dev::split_fragment(const Expr &e, vector &args) { + // One llvm value per 32-bit register. + const int num_regs = e.type().bits() * e.type().lanes() / 32; + if (const Load *load = e.as()) { + if (is_fragment_alloc(load->name)) { + llvm::Type *reg_type = fragment_reg_type(e.type()); + for (int i = 0; i < num_regs; i++) { + args.push_back(builder->CreateAlignedLoad( + reg_type, fragment_reg_ptr(load->name, e.type(), i), llvm::Align(4))); + } + return; + } + } + Value *v = codegen(e); + const int lanes_per_reg = 32 / e.type().bits(); + for (int i = 0; i < e.type().lanes() / lanes_per_reg; i++) { + Value *reg = lanes_per_reg == 1 ? + builder->CreateExtractElement(v, i) : + slice_vector(v, i * lanes_per_reg, lanes_per_reg); + if (wmma_reg_is_packed_i32(e.type())) { + reg = builder->CreateBitCast(reg, i32_t); + } + args.push_back(reg); + } +} + +Value *CodeGen_PTX_Dev::call_wmma_intrinsic(const std::string &name, + const vector &args, + const vector &overloads) { + llvm::Intrinsic::ID id = llvm::Intrinsic::lookupIntrinsicID(name); + internal_assert(id != llvm::Intrinsic::not_intrinsic) + << "Could not find the nvvm intrinsic " << name << "\n"; + llvm::Function *fn = llvm::Intrinsic::getOrInsertDeclaration(module.get(), id, overloads); + return builder->CreateCall(fn, args); +} + +void CodeGen_PTX_Dev::codegen_wmma(const Call *op) { + Value *result = codegen_wmma_raw(op); + + // Reassemble the returned struct into a Halide vector. + llvm::Type *result_type = llvm_type_of(op->type); + const int num_regs = op->type.bits() * op->type.lanes() / 32; + // A fragment that is a single register comes back as that register rather + // than as a struct holding one of them. + auto get_reg = [&](int i) { + return result->getType()->isStructTy() ? + builder->CreateExtractValue(result, i) : + result; + }; + if (op->type.bits() == 32) { + value = UndefValue::get(result_type); + for (int i = 0; i < num_regs; i++) { + value = builder->CreateInsertElement(value, get_reg(i), i); + } + } else { + vector regs; + regs.reserve(num_regs); + for (int i = 0; i < num_regs; i++) { + Value *reg = get_reg(i); + if (wmma_reg_is_packed_i32(op->type)) { + reg = builder->CreateBitCast( + reg, get_vector_type(llvm_type_of(op->type.element_of()), + 32 / op->type.bits())); + } + regs.push_back(reg); + } + value = concat_vectors(regs); + } + internal_assert(value->getType() == result_type) + << "Unexpected result type from a tensor core instruction\n"; +} + +bool CodeGen_PTX_Dev::is_fragment_alloc(const std::string &name) { + const MemoryType *t = alloc_memory_type.find(name); + return t && *t == MemoryType::Tile; +} + +llvm::Type *CodeGen_PTX_Dev::fragment_reg_type(Type t) { + return t.bits() == 32 ? llvm_type_of(t.element_of()) : + get_vector_type(llvm_type_of(t.element_of()), 32 / t.bits()); +} + +llvm::Value *CodeGen_PTX_Dev::fragment_reg_ptr(const std::string &name, Type t, int i) { + return codegen_buffer_pointer(name, t.element_of(), Expr(i * (32 / t.bits()))); +} + +void CodeGen_PTX_Dev::codegen_fragment_store(const Store *op) { + const Type t = op->value.type(); + const int num_regs = t.bits() * t.lanes() / 32; + llvm::Type *reg_type = fragment_reg_type(t); + + const Call *call = op->value.as(); + if (call && is_wmma_intrinsic(call)) { + // Take the registers straight out of the struct the instruction + // returns, without ever making a vector of them. + Value *result = codegen_wmma_raw(call); + for (int i = 0; i < num_regs; i++) { + builder->CreateAlignedStore(builder->CreateExtractValue(result, i), + fragment_reg_ptr(op->name, t, i), llvm::Align(4)); + } + return; + } + + // Anything else (a zero-initialization, say) does become a vector, but it + // is still written a register at a time so that the allocation only ever + // sees register-sized accesses. + Value *v = codegen(op->value); + const int lanes_per_reg = 32 / t.bits(); + for (int i = 0; i < num_regs; i++) { + Value *reg = lanes_per_reg == 1 ? + builder->CreateExtractElement(v, i) : + slice_vector(v, i * lanes_per_reg, lanes_per_reg); + internal_assert(reg->getType() == reg_type); + builder->CreateAlignedStore(reg, fragment_reg_ptr(op->name, t, i), llvm::Align(4)); + } +} + +// The element type the wmma intrinsic names use for a Halide type. The +// hardware takes 16-bit floats or 8-bit integers as multiplicands, and +// accumulates the first into 16 or 32-bit floats and the second into 32-bit +// integers. +std::string wmma_type_suffix(Type t) { + if (t == Float(16)) { + return "f16"; + } else if (t == BFloat(16)) { + return "bf16"; + } else if (t == Float(32)) { + return "f32"; + } else if (t == Int(8)) { + return "s8"; + } else if (t == UInt(8)) { + return "u8"; + } else if (t == Int(32)) { + return "s32"; + } + user_error << "There is no tensor core instruction for " << t << ".\n"; + return ""; +} + +llvm::Value *CodeGen_PTX_Dev::codegen_wmma_raw(const Call *op) { + // The nvvm wmma intrinsics take and return fragments as a flat list of + // 32-bit registers, packaged up as a literal struct. We represent them in + // Halide IR as vectors, so most of the work here is repacking. + auto get_int_arg = [&](int i) { + auto v = as_const_int(op->args[i]); + internal_assert(v) << "Expected a constant integer argument to " << op->name << "\n"; + return (int)*v; + }; + + const int M = get_int_arg(0), N = get_int_arg(1), K = get_int_arg(2); + const char *layouts[] = {"row", "col"}; + + std::ostringstream name; + name << "llvm.nvvm.wmma.m" << M << "n" << N << "k" << K << "."; + + vector args; + vector overloads; + + if (op->is_intrinsic(Call::wmma_mma)) { + // Half precision multiplicands name the instruction after the d and c + // operands, which for us are always the same type. Everything else + // names it after the multiplicands, which are in args 5 and 6. + const Type operand_type = op->args[5].type().element_of(); + std::string signature; + if (operand_type == Float(16)) { + const std::string acc = wmma_type_suffix(op->type.element_of()); + signature = acc + "." + acc; + } else { + signature = wmma_type_suffix(operand_type); + } + name << "mma." << layouts[get_int_arg(3)] << "." << layouts[get_int_arg(4)] + << "." << signature; + split_fragment(op->args[5], args); + split_fragment(op->args[6], args); + split_fragment(op->args[7], args); + } else { + // The a operand is M x K, the b operand is K x N, and the accumulator + // is M x N. + const bool is_a = op->is_intrinsic(Call::wmma_matrix_to_fragment_a); + const bool is_b = op->is_intrinsic(Call::wmma_matrix_to_fragment_b); + // The simplifier is free to have rewritten the load of the matrix into + // a dense load followed by a transpose, which is what happens to a + // column-major matrix. is_load_of_multiramp undoes that. + const Expr &arg = op->args[wmma_matrix_arg(op)]; + MultiRamp mr; + const Load *matrix = is_load_of_multiramp(arg, Scope::empty_scope(), &mr); + user_assert(matrix && matrix->type.element_of() == arg.type().element_of()) + << "The matrix a tensor core instruction takes a fragment out of is not a " + << "load with an affine index by the time it reaches the backend.\n"; + WMMAMatrixLayout mem = matrix_in_memory(matrix->name, mr, + is_b ? K : M, is_a ? K : N, arg); + const std::string type_suffix = wmma_type_suffix(op->type.element_of()); + name << "load." << (is_a ? "a" : is_b ? "b" : + "c") + << "." << (mem.row_major ? "row" : "col") << ".stride." << type_suffix; + + Value *ptr = codegen_buffer_pointer(matrix->name, matrix->type.element_of(), mem.base); + overloads.push_back(ptr->getType()); + args.push_back(ptr); + args.push_back(codegen(cast(Int(32), mem.stride))); + } + + return call_wmma_intrinsic(name.str(), args, overloads); +} + +void CodeGen_PTX_Dev::codegen_wmma_store(const Store *op) { + // Each lane writes the entries of the matrix that it holds, which the + // predicate describes and which one wmma store instruction does for the + // whole warp. + // A store to a column-major matrix arrives as a dense store of the + // transpose of it, because the simplifier rewrites stores to make them + // dense. Undoing that gives back the store as the extraction pass wrote it, + // and the column-major layout falls out of the index as usual. + Expr index; + const Call *inflate = peel_store_permutations(op, &index).as(); + internal_assert(inflate && inflate->args.size() == 5); + Expr predicate = op->predicate; + while (const Shuffle *shuffle = predicate.as()) { + predicate = shuffle->vectors[0]; + } + internal_assert(predicate.as() && + predicate.as()->is_intrinsic(Call::wmma_lane_owns)) + << "A store of a tensor core accumulator lost its predicate\n"; + + auto get_int_arg = [&](int i) { + auto v = as_const_int(inflate->args[i]); + internal_assert(v); + return (int)*v; + }; + const int M = get_int_arg(0), N = get_int_arg(1), K = get_int_arg(2); + const Expr &fragment = inflate->args[3]; + + MultiRamp mr; + internal_assert(is_multiramp(index, Scope::empty_scope(), &mr)); + WMMAMatrixLayout mem = matrix_in_memory(op->name, mr, M, N); + + std::ostringstream name; + name << "llvm.nvvm.wmma.m" << M << "n" << N << "k" << K << ".store.d." + << (mem.row_major ? "row" : "col") << ".stride." + << wmma_type_suffix(fragment.type().element_of()); + + Value *ptr = codegen_buffer_pointer(op->name, op->value.type().element_of(), mem.base); + vector args{ptr}; + vector overloads{ptr->getType()}; + split_fragment(fragment, args); + args.push_back(codegen(cast(Int(32), mem.stride))); + + call_wmma_intrinsic(name.str(), args, overloads); +} + string CodeGen_PTX_Dev::simt_intrinsic(const string &name) { if (ends_with(name, gpu_thread_name(0))) { return "llvm.nvvm.read.ptx.sreg.tid.x"; @@ -328,7 +737,8 @@ void CodeGen_PTX_Dev::visit(const For *loop) { void CodeGen_PTX_Dev::visit(const Allocate *alloc) { user_assert(!alloc->new_expr.defined()) << "Allocate node inside PTX kernel has custom new expression.\n" << "(Memoization is not supported inside GPU kernels at present.)\n"; - if (alloc->memory_type == MemoryType::GPUShared) { + ScopedBinding bind(alloc_memory_type, alloc->name, alloc->memory_type); + if (is_gpu_shared(alloc->memory_type)) { // PTX uses zero in address space 3 as the base address for shared memory Value *shared_base = Constant::getNullValue(PointerType::get(*context, 3)); sym_push(alloc->name, shared_base); @@ -369,7 +779,6 @@ void CodeGen_PTX_Dev::visit(const AssertStmt *op) { } void CodeGen_PTX_Dev::visit(const Load *op) { - // Do aligned 4-wide 32-bit loads as a single i128 load. const Ramp *r = op->index.as(); // TODO: lanes >= 4, not lanes == 4 @@ -388,13 +797,282 @@ void CodeGen_PTX_Dev::visit(const Load *op) { CodeGen_LLVM::visit(op); } +// The name of the Func a marked store belonged to, for error messages. The +// store itself is named after the packed allocation it ended up in. +std::string async_copy_func_name(const Call *marker) { + internal_assert(marker->args.size() == 3); + const StringImm *name = marker->args[2].as(); + internal_assert(name) << "cuda_bypass_registers name is not a string\n"; + return name->value; +} + +// A copy from global memory into shared memory can be done by the hardware +// without going through registers, which saves the load, the store, and the +// registers in between. The copy is asynchronous, so it has to be waited for +// before the data is used; that happens at the end of the producer. +bool CodeGen_PTX_Dev::codegen_async_copy(const Store *op, const char **reason) { + if (target.get_cuda_capability_lower_bound() < 80) { + *reason = "asynchronous copies require CUDA compute capability 8.0 or above"; + return false; + } + if (emit_atomic_stores) { + *reason = "the store is atomic"; + return false; + } + // Asynchronous copies need something to wait for them, which only happens + // inside a producer. + if (!in_producer) { + *reason = "the store is not inside a produce node"; + return false; + } + + // The value must be one the schedule asked to have moved without passing + // through registers, and it must be a plain load from something we didn't + // allocate in here, which is to say global memory. An ordinary store that + // happens to match the pattern is left synchronous. + // CSE and LICM lift common subexpressions of a stored value into Lets + // around it, which would hide the marker. Substituting them back in leaves + // an expression that stands alone, which is what the analysis below and + // codegen_buffer_pointer both need. It has to be held in a local, because + // everything below points into it. + const Expr stored = substitute_in_all_lets(op->value); + const Call *marker = stored.as(); + if (!(marker && marker->is_intrinsic(Call::cuda_bypass_registers))) { + *reason = "the store was not marked as an asynchronous copy"; + return false; + } + internal_assert(marker->args.size() == 3); + const Expr &copied = marker->args[0]; + auto group = as_const_int(marker->args[1]); + internal_assert(group) << "cuda_bypass_registers group is not a constant integer\n"; + + const Load *src = copied.as(); + if (!src) { + // A load that isn't dense is broken up into a shuffle of dense loads + // well before we get here, so say what that means for the copy rather + // than describing it as not being a load. + Expr value = copied; + if (const Shuffle *s = value.as(); + s && !s->vectors.empty()) { + *reason = "the source is not read densely. Each copy moves one run of " + "bytes, so the Func must read its source with a stride of one"; + return false; + } + *reason = "the value stored is not a load from a buffer outside the kernel. " + "An asynchronous copy moves bytes untouched, so the Func must be a " + "plain copy - no cast, no arithmetic, and no boundary condition"; + return false; + } + if (alloc_memory_type.contains(src->name)) { + *reason = "the value stored is loaded from another allocation inside the " + "kernel. An asynchronous copy reads from global memory"; + return false; + } + if (!is_const_one(op->predicate) || !is_const_one(src->predicate)) { + *reason = "the load or the store is predicated"; + return false; + } + + // The hardware copies 4, 8 or 16 bytes at a time, from and to consecutive + // addresses. + const Type t = copied.type(); + const int bytes = t.bytes() * t.lanes(); + if (!(bytes == 4 || bytes == 8 || bytes == 16)) { + *reason = "each thread must copy 4, 8 or 16 bytes at a time. Vectorize the " + "copy along its dense dimension by that many bytes' worth"; + return false; + } + Expr dst_base = op->index, src_base = src->index; + if (t.lanes() > 1) { + // Shared allocations are given an offset into one big block after the + // last simplification pass, so the indices need simplifying here. + // strided_ramp_base returns undefined unless the stride is exactly + // one, so this checks the density as well as finding the address. + dst_base = strided_ramp_base(simplify(op->index)); + src_base = strided_ramp_base(simplify(src->index)); + if (!dst_base.defined() || !src_base.defined()) { + *reason = "the source and the destination are not both indexed densely"; + return false; + } + } + // The hardware needs both addresses aligned to the width of the copy. A + // stride that isn't a multiple of it - which is what an odd align_storage + // produces - shows up here as an unprovable alignment. + if (t.lanes() > 1) { + // Use the alignment lowering worked out, which knows what the loop + // variables in the index are multiples of. + // Only the destination is checked. Its alignment is what align_storage + // controls, so it is the one a schedule can get wrong. A source in a + // buffer whose strides are not known until runtime has no provable + // alignment either way, and rejecting those would fail schedules that + // are fine. + auto aligned = [&](const ModulusRemainder &a, const Expr &base) { + auto ok = [&](const ModulusRemainder &m) { + return m.modulus % t.lanes() == 0 && m.remainder % t.lanes() == 0; + }; + return ok(a) || ok(modulus_remainder(base)); + }; + if (!aligned(op->alignment, dst_base)) { + *reason = "the destination is not known to be aligned to the width of " + "the copy. Any align_storage on this Func has to be a multiple " + "of the number of elements each thread copies"; + return false; + } + } + + Value *dst = codegen_buffer_pointer(op->name, t.element_of(), dst_base); + Value *src_ptr = codegen_buffer_pointer(src->name, t.element_of(), src_base); + + // Shared allocations are already in the shared address space. The source is + // a kernel argument, so it's global, but it comes in as a generic pointer. + llvm::Type *shared_ptr_t = PointerType::get(*context, 3); + llvm::Type *global_ptr_t = PointerType::get(*context, 1); + if (dst->getType() != shared_ptr_t) { + *reason = "the destination did not end up in the shared address space"; + return false; + } + src_ptr = builder->CreateAddrSpaceCast(src_ptr, global_ptr_t); + + std::ostringstream name; + name << "llvm.nvvm.cp.async.ca.shared.global." << bytes; + llvm::Intrinsic::ID id = llvm::Intrinsic::lookupIntrinsicID(name.str()); + internal_assert(id != llvm::Intrinsic::not_intrinsic) + << "Could not find the nvvm intrinsic " << name.str() << "\n"; + llvm::Function *fn = llvm::Intrinsic::getOrInsertDeclaration(module.get(), id); + // Copies are committed in groups, so close the previous group before + // starting one for a different batch. + if (uncommitted_group != -1 && uncommitted_group != (int)*group) { + commit_copies(); + } + builder->CreateCall(fn, {dst, src_ptr}); + uncommitted_group = (int)*group; + return true; +} + +void CodeGen_PTX_Dev::visit(const ProducerConsumer *op) { + if (!op->is_producer) { + CodeGen_LLVM::visit(op); + return; + } + + ScopedValue old_in(in_producer, true); + codegen(op->body); +} + +void CodeGen_PTX_Dev::commit_copies() { + if (uncommitted_group == -1) { + return; + } + llvm::Intrinsic::ID id = + llvm::Intrinsic::lookupIntrinsicID("llvm.nvvm.cp.async.commit.group"); + internal_assert(id != llvm::Intrinsic::not_intrinsic); + builder->CreateCall(llvm::Intrinsic::getOrInsertDeclaration(module.get(), id), {}); + committed_groups.push_back(uncommitted_group); + uncommitted_group = -1; +} + +void CodeGen_PTX_Dev::emit_copy_wait(int n) { + llvm::Intrinsic::ID id = + llvm::Intrinsic::lookupIntrinsicID("llvm.nvvm.cp.async.wait.group"); + internal_assert(id != llvm::Intrinsic::not_intrinsic); + llvm::Function *fn = llvm::Intrinsic::getOrInsertDeclaration(module.get(), id); + vector args; + if (fn->getFunctionType()->getNumParams() == 1) { + args.push_back(ConstantInt::get(i32_t, n)); + } + builder->CreateCall(fn, args); +} + +void CodeGen_PTX_Dev::await_copies(int group) { + commit_copies(); + // The wait is FIFO, so waiting for this group means letting everything + // committed after it stay outstanding. Searching from the newest end finds + // the most recent batch with this group, which is the one just issued. + for (size_t i = committed_groups.size(); i > 0; i--) { + if (committed_groups[i - 1] != group) { + continue; + } + emit_copy_wait((int)(committed_groups.size() - i)); + committed_groups.erase(committed_groups.begin(), + committed_groups.begin() + i); + return; + } + // Nothing from that group is outstanding, so there is nothing to wait for. +} + +void CodeGen_PTX_Dev::await_all_copies() { + commit_copies(); + if (committed_groups.empty()) { + return; + } + emit_copy_wait(0); + committed_groups.clear(); +} + void CodeGen_PTX_Dev::visit(const Store *op) { + if (is_wmma_matrix_store(op)) { + codegen_wmma_store(op); + return; + } + // Issue atomic store if we are inside an Atomic node. if (emit_atomic_stores) { user_assert(is_const_one(op->predicate)) << "Atomic update does not support predicated store.\n"; user_assert(op->value.type().bits() >= 32) << "CUDA: 8-bit or 16-bit atomics are not supported.\n"; } + if (is_fragment_alloc(op->name)) { + codegen_fragment_store(op); + return; + } + + { + const char *reason = ""; + if (codegen_async_copy(op, &reason)) { + return; + } + // Asking for that memory type is a promise that the stores to it are + // copies the hardware can make asynchronously. If one isn't, say so + // rather than quietly emitting a load and a store instead. + const Expr stored = substitute_in_all_lets(op->value); + const Call *marker = stored.as(); + if (marker && marker->is_intrinsic(Call::cuda_bypass_registers)) { + user_error + << async_copy_func_name(marker) << " is scheduled in GPUSharedAsync memory, but this " + << "store to it cannot be done with an asynchronous copy, because " + << reason << ".\n\n" + << "An asynchronous copy moves bytes from global memory into shared " + << "memory without routing them through registers. It requires that " + << "the Func is a plain copy of a buffer or another Func - no cast, " + << "arithmetic, or boundary condition, because the bytes move " + << "untouched - and that each thread stores a dense vector of 4, 8 or " + << "16 bytes, aligned to its own size at both ends. It needs CUDA " + << "compute capability 8.0 or above.\n\n" + << "The alignment of the destination is set by align_storage, which " + << "fixes the stride of the staged Func. Padding it to avoid bank " + << "conflicts is usually a good idea, but the padded stride has to " + << "stay a multiple of the vector width or the rows stop being " + << "aligned enough to copy into.\n\n" + << "The usual way to get a Func that is a plain copy is Func::in, " + << "which makes a wrapper that does nothing but hold a staged copy of " + << "something. Vectorizing its dense dimension by a whole number of " + << "bytes gives each thread one copy to issue, and its other " + << "dimensions are spread over the threads of the block as usual:\n\n" + << " A.in()\n" + << " .compute_at(consumer, r)\n" + << " .store_in(MemoryType::GPUSharedAsync)\n" + << " .tile(x, y, xi, yi, 256, 8) // 256 = 8 elements x 32 threads\n" + << " .vectorize(xi, 8) // 8 halves is 16 bytes\n" + << " .gpu_threads(xi, yi);\n\n" + << "The tile has to divide into the thread counts the block already " + << "has: its width over the vector width is the number of threads in " + << "x, and its height the number in y. Anything left over is covered " + << "by the serial loops the tile leaves outside.\n\n" + << "The store that could not be made asynchronous was:\n" + << Stmt(op); + } + } + // Do aligned 4-wide 32-bit stores as a single i128 store. const Ramp *r = op->index.as(); // TODO: lanes >= 4, not lanes == 4 diff --git a/src/CodeGen_Vulkan_Dev.cpp b/src/CodeGen_Vulkan_Dev.cpp index 2918eba7b85b..b7d8dd844667 100644 --- a/src/CodeGen_Vulkan_Dev.cpp +++ b/src/CodeGen_Vulkan_Dev.cpp @@ -1877,7 +1877,7 @@ void CodeGen_Vulkan_Dev::SPIRV_Emitter::visit(const Allocate *op) { uint32_t array_size = 0; SpvStorageClass storage_class = SpvStorageClassGeneric; - if (op->memory_type == MemoryType::GPUShared) { + if (is_gpu_shared(op->memory_type)) { // Allocation of shared memory must be declared at global scope storage_class = SpvStorageClassWorkgroup; // shared across workgroup diff --git a/src/CodeGen_WebGPU_Dev.cpp b/src/CodeGen_WebGPU_Dev.cpp index 39a88337af1d..23db34fa02ad 100644 --- a/src/CodeGen_WebGPU_Dev.cpp +++ b/src/CodeGen_WebGPU_Dev.cpp @@ -410,7 +410,7 @@ void CodeGen_WebGPU_Dev::CodeGen_WGSL::add_kernel( } void CodeGen_WebGPU_Dev::CodeGen_WGSL::visit(const Allocate *op) { - if (op->memory_type == MemoryType::GPUShared) { + if (is_gpu_shared(op->memory_type)) { internal_assert(!workgroup_allocations.count(op->name)); workgroup_allocations.insert({op->name, op}); op->body.accept(this); diff --git a/src/CodeGen_X86.cpp b/src/CodeGen_X86.cpp index f88b67a68c91..d8c6351b89c7 100644 --- a/src/CodeGen_X86.cpp +++ b/src/CodeGen_X86.cpp @@ -1648,9 +1648,9 @@ void CodeGen_X86::visit(const Allocate *op) { void CodeGen_X86::visit(const Load *op) { if (const auto *mt = mem_type.find(op->name)) { - if (*mt == MemoryType::AMXTile) { + if (*mt == MemoryType::Tile) { const Ramp *ramp = op->index.as(); - internal_assert(ramp) << "Expected AMXTile to have index ramp\n"; + internal_assert(ramp) << "Expected a tile to have index ramp\n"; Value *ptr = codegen_buffer_pointer(op->name, op->type, ramp->base); LoadInst *load = builder->CreateAlignedLoad(llvm_type_of(upgrade_type_for_storage(op->type)), ptr, llvm::Align(op->type.bytes())); add_tbaa_metadata(load, op->name, op->index); @@ -1663,11 +1663,11 @@ void CodeGen_X86::visit(const Load *op) { void CodeGen_X86::visit(const Store *op) { if (const auto *mt = mem_type.find(op->name)) { - if (*mt == MemoryType::AMXTile) { + if (*mt == MemoryType::Tile) { Value *val = codegen(op->value); Halide::Type value_type = op->value.type(); const Ramp *ramp = op->index.as(); - internal_assert(ramp) << "Expected AMXTile to have index ramp\n"; + internal_assert(ramp) << "Expected a tile to have index ramp\n"; Value *ptr = codegen_buffer_pointer(op->name, value_type, ramp->base); StoreInst *store = builder->CreateAlignedStore(val, ptr, llvm::Align(value_type.bytes())); add_tbaa_metadata(store, op->name, op->index); diff --git a/src/Deserialization.cpp b/src/Deserialization.cpp index e31c27eea492..54bdfa5f631e 100644 --- a/src/Deserialization.cpp +++ b/src/Deserialization.cpp @@ -180,14 +180,16 @@ MemoryType Deserializer::deserialize_memory_type(Serialize::MemoryType memory_ty return MemoryType::Register; case Serialize::MemoryType::GPUShared: return MemoryType::GPUShared; + case Serialize::MemoryType::GPUSharedAsync: + return MemoryType::GPUSharedAsync; case Serialize::MemoryType::GPUTexture: return MemoryType::GPUTexture; case Serialize::MemoryType::LockedCache: return MemoryType::LockedCache; case Serialize::MemoryType::VTCM: return MemoryType::VTCM; - case Serialize::MemoryType::AMXTile: - return MemoryType::AMXTile; + case Serialize::MemoryType::Tile: + return MemoryType::Tile; default: user_error << "unknown memory type " << (int)memory_type << "\n"; return MemoryType::Auto; diff --git a/src/Expr.h b/src/Expr.h index 0f71b860f4b9..665148c247aa 100644 --- a/src/Expr.h +++ b/src/Expr.h @@ -402,11 +402,37 @@ enum class MemoryType { * on Hexagon */ VTCM, - /** AMX Tile register for X86. Any data that would be used in an AMX matrix - * multiplication must first be loaded into an AMX tile register. */ - AMXTile, + /** Storage for a matrix tile, in whatever form the target's matrix unit + * keeps one: an AMX tile register on x86, or an NVIDIA tensor core + * fragment striped across the registers of the 32 lanes of a warp. Either + * way the layout is not architecturally specified, so the only legal + * accesses are the whole-tile ones a matrix unit provides - filling a + * tile, loading one from memory, storing one back, and multiplying a pair + * of them into a third. Which of those a given tile takes part in, and so + * which role it plays in the multiply, follows from how it is used. + * + * Schedule these with tile_init, tile_load, tile_store and tile_matmul, + * which set this memory type where it is needed. */ + Tile, + + /** GPU shared memory, written by an asynchronous copy instruction that + * moves data straight from global memory without routing it through + * registers. The only stores allowed to such an allocation are unpredicated + * copies of a whole number of 4, 8 or 16 byte chunks from a densely indexed + * global buffer, because that is all the hardware can do. On GPU APIs with + * no such instruction this is ordinary shared memory. */ + GPUSharedAsync, + + /** Deprecated alias for Tile, which covers the tile storage of every + * target's matrix unit rather than just x86's. */ + AMXTile = Tile, }; +/** Whether a MemoryType places an allocation in GPU shared memory. */ +inline bool is_gpu_shared(MemoryType t) { + return t == MemoryType::GPUShared || t == MemoryType::GPUSharedAsync; +} + /** Whether a MemoryType is backed by tile-shaped storage with native 2D * loads and stores. Generic optimizations that rewrite the structure of * loads or stores (deinterleaving, strided-load staging, etc.) should @@ -414,7 +440,7 @@ enum class MemoryType { * dedicated lowering pass that requires the original 2D-shaped loads * and stores to remain intact. */ inline bool is_tile_memory_type(MemoryType t) { - return t == MemoryType::AMXTile; + return t == MemoryType::Tile; } namespace Internal { diff --git a/src/ExtractTileOperations.cpp b/src/ExtractTileOperations.cpp index eb1dfc7660df..ef5124d77a79 100644 --- a/src/ExtractTileOperations.cpp +++ b/src/ExtractTileOperations.cpp @@ -121,11 +121,6 @@ Matmul convert_to_matmul(const Store *op, const string &new_name) { return fail("the vector reduction is not of a widening multiply"); } - if (lhs.type().bits() != 8 || - rhs.type().bits() != 8) { - return fail("the vector reduction operand or result types are not supported"); - } - } else { // Lower a widening_mul intrinsic, as they can be used but aren't lifted to for bf16. Expr reduce_value = simplify(lower_intrinsics(reduce->value)); @@ -137,65 +132,40 @@ Matmul convert_to_matmul(const Store *op, const string &new_name) { rhs = mul->b; } - // There may be a broadcast next (it can get hoisted outside of other ops) - auto debroadcast = [](Expr &e) -> int { - if (const Broadcast *b = e.as()) { - int lanes = b->lanes; - e = b->value; - return lanes; - } else { - return 1; - } - }; - int lhs_broadcast = debroadcast(lhs); - int rhs_broadcast = debroadcast(rhs); - - // Unpack the casts, if it was a direct multiply. This should only happen - // for floats (the integer branch above already extracted the cast inputs - // from the widening_mul intrinsic). - if (reduce->type.is_float()) { - const auto *lhs_cast = lhs.as(); - const auto *rhs_cast = rhs.as(); - if (!lhs_cast || !rhs_cast) { - return fail("the vector reduction is not of a widening multiply"); - } - lhs = lhs_cast->value; - rhs = rhs_cast->value; - if (!lhs.type().is_bfloat() || - rhs.type().element_of() != lhs.type().element_of()) { - return fail("the vector reduction operand or result types are not supported"); - } - } - - // Underneath all of this must be a load + // Underneath all of this must be a load, though it may be wrapped in a + // broadcast over the dimension it doesn't depend on, in the widening cast + // (for floats - the integer branch above already extracted the cast inputs + // from the widening_mul intrinsic), and in a lane permutation. The cast is + // checked against the load's type below. // TODO: What if we want to multiply by the same matrix multiple times? It might be a let binding. - const auto *lhs_load = lhs.as(); - const auto *rhs_load = rhs.as(); + MultiRamp lhs_mr, rhs_mr; + Scope empty_scope; + const auto *lhs_load = is_load_of_multiramp(lhs, empty_scope, &lhs_mr); + const auto *rhs_load = is_load_of_multiramp(rhs, empty_scope, &rhs_mr); if (!lhs_load || !rhs_load) { - return fail("the matrix multiply operands are not loads"); + return fail("the matrix multiply operands are not loads with affine indices"); } // The loads must be unpredicated if (!is_const_one(lhs_load->predicate) || !is_const_one(rhs_load->predicate)) { return fail("the matrix multiply operands are predicated loads"); } - // Now we analyze the load indices as multiramps - MultiRamp lhs_mr, rhs_mr; - Scope empty_scope; - if (!is_multiramp(lhs_load->index, empty_scope, &lhs_mr) || - !is_multiramp(rhs_load->index, empty_scope, &rhs_mr)) { - return fail("the matrix multiply loads indices are not affine"); - } - - // Add back on any broadcasts as a stride-0 outer dim. - auto add_broadcast = [](MultiRamp &mr, int extent) { - if (extent > 1) { - mr.strides.push_back(make_zero(mr.base.type())); - mr.lanes.push_back(extent); + if (reduce->type.is_int_or_uint()) { + // The tile registers are typed by what was loaded, and the signedness + // of those types picks which of the four integer tdpb instructions + // gets used, so a cast between the load and the multiply would change + // the meaning of the multiply without changing the instruction. + if (lhs.type().element_of() != lhs_load->type.element_of() || + rhs.type().element_of() != rhs_load->type.element_of()) { + return fail("the matrix multiply operands are cast after being loaded"); } - }; - add_broadcast(lhs_mr, lhs_broadcast); - add_broadcast(rhs_mr, rhs_broadcast); + if (lhs_load->type.bits() != 8 || rhs_load->type.bits() != 8) { + return fail("the vector reduction operand or result types are not supported"); + } + } else if (!lhs_load->type.is_bfloat() || + rhs_load->type.element_of() != lhs_load->type.element_of()) { + return fail("the vector reduction operand or result types are not supported"); + } // In a matrix multiply with row-major inputs and outputs, the algorithm // looks like: @@ -234,7 +204,6 @@ Matmul convert_to_matmul(const Store *op, const string &new_name) { // Now deduce LHS and RHS. First some helpers. auto swap_sides = [&]() { - std::swap(lhs, rhs); std::swap(lhs_mr, rhs_mr); std::swap(lhs_load, rhs_load); }; @@ -309,8 +278,8 @@ Matmul convert_to_matmul(const Store *op, const string &new_name) { { Type t = op->value.type(); bool result_ok = t.bytes() * I * J <= 1024; - bool lhs_ok = lhs.type().bytes() * I * K <= 1024; - bool rhs_ok = rhs.type().bytes() * K * J <= 1024; + bool lhs_ok = lhs_load->type.bytes() * I * K <= 1024; + bool rhs_ok = rhs_load->type.bytes() * K * J <= 1024; if (!result_ok || !lhs_ok || !rhs_ok) { return fail("one more more matrices are too large to fit in AMX registers (more than 1024 bytes)"); } @@ -424,86 +393,18 @@ class ExtractTileOperations : public IRMutator { int found_J = -1; int found_K = -1; - // An AMXTile allocation may represent multiple AMX accumulator + // A tile allocation may represent multiple AMX accumulator // registers as 2D sub-tiles. This map tracks those. std::vector amx_subtiles; - // Returns a unique subtile index for a load or store index, or -1 if it - // overlaps with an existing subtile, or is otherwise poorly behaved. - int get_subtile(const Expr &index) { - MultiRamp mr; - if (!is_multiramp(index, Scope::empty_scope(), &mr)) { - user_error << "Access to AMX tile not affine: " << index << "\n"; - } - if (!can_prove(mr.alias_free())) { - // What are you doing? - user_error << "Access to AMX tile may have duplicated lanes: " << index << "\n"; - } - if (amx_subtiles.empty()) { - amx_subtiles.push_back(std::move(mr)); - return 0; - } - - // All strides and lanes must match across all subtiles, or we give up. - const MultiRamp &first = amx_subtiles[0]; - if (mr.dimensions() != first.dimensions()) { - user_error - << "Access to AMX tile does not have the same shape as other accesses to the same memory."; - return -1; - } - for (int i = 0; i < first.dimensions(); i++) { - if (!can_prove(mr.strides[i] == first.strides[i]) || - mr.lanes[i] != first.lanes[i]) { - user_error - << "Access to AMX tile has different size and strides to other " - << "accesses to the same memory. All accesses must have the same " - << "subtile size and strides: " << index; - } - } - - // Now check for disjointedness - // Add a synthetic dimension, the purpose of which will become clear. - mr.strides.emplace_back(); - mr.lanes.push_back(2); - for (int i = 0; i < (int)amx_subtiles.size(); i++) { - auto &other = amx_subtiles[i]; - // One of two things must be true: - // 1) All of the lanes of mr equal the corresponding lane of - // other. We've already checked the strides and lanes, so it's just - // a matter of checking the base. - if (can_prove(mr.base == other.base)) { - return i; - } - - // 2) None of the lanes or mr equal any of the lanes of other. To do - // this we'll construct a combined mr that can be either 'mr' or - // 'other', and ask if it's alias-free. This is what the synthetic - // dimension was for. - mr.strides.back() = mr.base - other.base; - if (!can_prove(mr.alias_free())) { - user_error - << "Failed to prove access to AMX does not partially overlap " - << "another distinct access: " << index; - return -1; - } - } - - // Didn't already exist and didn't alias with anything. - mr.strides.pop_back(); - mr.lanes.pop_back(); - amx_subtiles.push_back(std::move(mr)); - return (int)amx_subtiles.size() - 1; - } - // Returns an index expression for a given load or store index. user_asserts if impossible std::string get_subtile_name(const Expr &index) { - int idx = get_subtile(index); - internal_assert(idx >= 0); // errors handled already + int idx = get_subtile(index, "AMX tile", &amx_subtiles); return amx_name + std::to_string(idx); } Stmt visit(const Allocate *op) override { - if (op->memory_type == MemoryType::AMXTile) { + if (op->memory_type == MemoryType::Tile) { user_assert( (op->type.is_int() && op->type.bits() == 32) || (op->type.is_float() && op->type.bits() == 32)) @@ -521,7 +422,7 @@ class ExtractTileOperations : public IRMutator { pass = 0; body = mutate(body); user_assert(found_I >= 0 && found_J >= 0 && found_K >= 0) - << op->name << " is stored in AMXTile memory, but no matrix multiply " + << op->name << " is stored in Tile memory, but no matrix multiply " << "operation was found that stores to it, so the shape of the tile " << "was unable to be determined.\n"; pass = 1; @@ -529,7 +430,7 @@ class ExtractTileOperations : public IRMutator { for (int i = 0; i < (int)amx_subtiles.size(); i++) { body = Allocate::make(amx_name + std::to_string(i), op->type.element_of(), - MemoryType::AMXTile, {256}, const_true(), body); + MemoryType::Tile, {256}, const_true(), body); } return body; } diff --git a/src/ExtractTileOperations.h b/src/ExtractTileOperations.h index 918e3b1b9940..38da92329fca 100644 --- a/src/ExtractTileOperations.h +++ b/src/ExtractTileOperations.h @@ -11,7 +11,7 @@ namespace Halide { namespace Internal { -/** Rewrite any AMX tile operations that have been stored in the AMXTile memory +/** Rewrite any AMX tile operations that have been stored in the Tile memory * type as intrinsic calls, to be used in the X86 backend. */ Stmt extract_tile_operations(const Stmt &s); diff --git a/src/ExtractWMMAOperations.cpp b/src/ExtractWMMAOperations.cpp new file mode 100644 index 000000000000..7d15427b2ff1 --- /dev/null +++ b/src/ExtractWMMAOperations.cpp @@ -0,0 +1,861 @@ +#include "ExtractWMMAOperations.h" + +#include "CanonicalizeGPUVars.h" +#include "FindIntrinsics.h" +#include "IREquality.h" +#include "IRMutator.h" +#include "IROperator.h" +#include "MultiRamp.h" +#include "Simplify.h" +#include "Substitute.h" +#include "Util.h" + +#include + +/** \file Support extraction of NVIDIA tensor core (wmma) instructions. + * + * The wmma instructions are warp-level: the 32 lanes of a warp cooperate to + * hold a tile of a matrix, and each instruction is executed by the warp as a + * whole. The layout of the tile across the lanes' registers is not specified + * by the architecture, so the only way to get data into or out of a tile is + * with the wmma load and store instructions, which move a tile between the + * registers of a warp and a 2D array in memory. + * + * A tile allocation holds one of the three matrices of a multiply, and + * which one follows from how it is used. An allocation accumulated into by a + * matrix multiply is the accumulator; one read as an operand of a multiply is + * that operand. The role determines which accesses are legal: + * + * - An accumulator may be zero-initialized, filled from a matrix in memory, + * accumulated into by a matrix multiply, and copied back out to memory. + * - An operand may be filled from a matrix in memory and read by a matrix + * multiply. + * + * Anything else is an error. Every one of those operations is wrapped in a loop + * over the 32 lanes of a warp, because nothing in the schedule says the tile is + * spread over a warp - that's a consequence of asking for tensor core storage. + * + * Operands don't have to be staged in fragments. If a multiply reads its + * operand straight out of shared or global memory, the load into registers is + * synthesized at the multiply instead. Staging one explicitly is how a schedule + * says to load it once and reuse it across several multiplies. + */ + +namespace Halide { +namespace Internal { + +using std::string; +using std::vector; + +namespace { + +// The tile shapes the hardware supports, in the order we try them. +struct Shape { + int M, N, K; +}; +const Shape supported_shapes[] = {{16, 16, 16}, {32, 8, 16}, {8, 32, 16}}; + +// A matrix operand is stored either with its rows contiguous or its columns +// contiguous. +enum class Layout { + Row, + Col, +}; + +// Which of the three matrices of the multiply a fragment holds. +enum class Role { + Unknown, + A, + B, + Accumulator, +}; + +Call::IntrinsicOp intrinsic_for_role(Role role) { + switch (role) { + case Role::A: + return Call::wmma_matrix_to_fragment_a; + case Role::B: + return Call::wmma_matrix_to_fragment_b; + default: + return Call::wmma_matrix_to_fragment_c; + } +} + +const char *role_name(Role role) { + switch (role) { + case Role::A: + return "the first operand"; + case Role::B: + return "the second operand"; + case Role::Accumulator: + return "the accumulator"; + default: + return "no part"; + } +} + +// Every warp is 32 lanes, and an accumulator tile holds 256 elements whatever +// its shape, so each lane holds 8 of them whatever their type. +constexpr int warp_lanes = 32; +constexpr int accumulator_elements = 8; +// Half precision operand fragments are eight 32-bit registers per lane +// whatever the shape, which is sixteen elements. For the shapes that hold +// fewer than that the hardware replicates elements across lanes. +constexpr int half_operand_elements = 16; + +// The multiplicand and accumulator type combinations the tensor cores have an +// instruction for. Half precision multiplicands accumulate into halves or +// floats, brain floats accumulate into floats, and eight-bit integers +// accumulate into 32-bit ones. The multiplicands must match each other: at +// this shape there is no instruction for mixed signedness. +bool wmma_types_supported(Type operand, Type accumulator) { + if (operand == Float(16)) { + return accumulator == Float(16) || accumulator == Float(32); + } else if (operand == BFloat(16)) { + return accumulator == Float(32); + } else if (operand == Int(8) || operand == UInt(8)) { + return accumulator == Int(32); + } + return false; +} + +// The shape of the matrix each fragment is taken out of. +void fragment_matrix_shape(Role role, const Shape &shape, int *rows, int *cols) { + *rows = role == Role::B ? shape.K : shape.M; + *cols = role == Role::A ? shape.K : shape.N; +} + +// How many elements of a fragment each lane holds. Every type but half +// precision holds exactly its share of the matrix, so this follows from the +// shape; half precision is the fixed size above however small the shape is. +int elements_per_lane(Role role, const Shape &shape, Type t) { + if (role == Role::Accumulator) { + return accumulator_elements; + } + if (t == Float(16)) { + return half_operand_elements; + } + int rows, cols; + fragment_matrix_shape(role, shape, &rows, &cols); + return rows * cols / warp_lanes; +} + +// One operand of a matrix multiply, described in the canonical [K, N, M] +// (innermost first) coordinate system. +struct Operand { + const Load *load = nullptr; + MultiRamp mr; + vector strides; +}; + +// Try to interpret an operand as the left-hand side of the multiply: an M x K +// matrix that doesn't depend on the N coordinate. +bool is_lhs(const Operand &op, Layout *layout, Expr *stride) { + if (!is_const_zero(op.strides[1])) { + return false; + } + if (is_const_one(op.strides[0])) { + *layout = Layout::Row; + *stride = op.strides[2]; + } else if (is_const_one(op.strides[2])) { + *layout = Layout::Col; + *stride = op.strides[0]; + } else { + return false; + } + return true; +} + +// Try to interpret an operand as the right-hand side of the multiply: a K x N +// matrix that doesn't depend on the M coordinate. +bool is_rhs(const Operand &op, Layout *layout, Expr *stride) { + if (!is_const_zero(op.strides[2])) { + return false; + } + if (is_const_one(op.strides[1])) { + *layout = Layout::Row; + *stride = op.strides[0]; + } else if (is_const_one(op.strides[0])) { + *layout = Layout::Col; + *stride = op.strides[1]; + } else { + return false; + } + return true; +} + +// The wmma instructions are warp-level, so every statement that touches an +// accumulator has to be executed by all 32 lanes of a warp. Nothing in the +// schedule says so - it's a consequence of asking for tensor core storage - so +// the loop over lanes is introduced here. +// The lane of the warp a tensor core operation runs in. Every wmma intrinsic +// takes it, because their values depend on it and nothing else in their +// arguments does. +Expr make_lane(const string &name) { + return Variable::make(Int(32), name); +} + +Stmt in_lane_loop(const Expr &lane, Stmt s) { + const Variable *v = lane.as(); + internal_assert(v) << "the lane of a tensor core operation is not a variable\n"; + return For::make(v->name, 0, warp_lanes - 1, ForType::GPULane, Partition::Never, + DeviceAPI::CUDA, std::move(s)); +} + +// The matrix a wmma load or store moves, described as a Load of every element +// of it, with the lanes in row-major order regardless of how it is laid out in +// memory. A row-major matrix has consecutive columns one element apart and its +// rows stride apart; a column-major one is the other way around, which changes +// the strides of the index but not the order of the lanes. So the argument is +// the matrix value, and the memory layout is recoverable from the index. +Expr make_matrix_index(const Expr &base, int rows, int cols, Layout layout, + const Expr &stride) { + Expr one = make_one(base.type()); + Expr col_stride = layout == Layout::Row ? one : stride; + Expr row_stride = layout == Layout::Row ? stride : one; + return Ramp::make(Ramp::make(base, col_stride, cols), + Broadcast::make(row_stride, cols), rows); +} + +Expr make_matrix_address(const string &name, Type element_type, const Expr &base, + int rows, int cols, Layout layout, const Expr &stride, + const Buffer<> &image, const Parameter ¶m) { + Expr index = make_matrix_index(base, rows, cols, layout, stride); + const int lanes = rows * cols; + return Load::make(element_type.with_lanes(lanes), name, index, image, param, + const_true(lanes), ModulusRemainder(), false); +} + +Expr make_matrix_to_fragment(Role role, const Shape &shape, Layout layout, + const Load *load, const Expr &base, const Expr &stride, + const Expr &lane) { + int rows, cols; + fragment_matrix_shape(role, shape, &rows, &cols); + Expr address = make_matrix_address(load->name, load->type.element_of(), base, + rows, cols, layout, stride, load->image, load->param); + Type type = load->type.element_of().with_lanes( + elements_per_lane(role, shape, load->type.element_of())); + return Call::make(type, intrinsic_for_role(role), + {shape.M, shape.N, shape.K, std::move(address), lane}, + Call::Intrinsic); +} + +// A store to a fragment recognized as the accumulation of a matrix multiply, +// broken down into the pieces the multiply is built from. +struct MatmulInfo { + Shape shape; + Operand lhs, rhs; + Layout lhs_layout = Layout::Row, rhs_layout = Layout::Row; + Expr lda, ldb; + Type accumulator_type; + vector> peeled_lets; +}; + +MatmulInfo analyze_matmul(const Store *op) { + // We expect the pattern: + // + // out[idx] = reduce_add(widen(lhs[multiramp]) * widen(rhs[multiramp])) + out[idx] + // + // Though either operand may have been hoisted out to a broadcast or had a + // lane permutation left on it by vectorization. + + auto fail = [&](const char *reason) -> MatmulInfo { + user_error << "Matrix multiply not recognized. Store to a Tile " + << "allocation must be a zero-initialization, a fill from a matrix " + << "in memory, or a sum of a vector reduce op and a load from the " + << "same allocation. In the following store, " << reason << ".\n" + << Stmt(op); + return MatmulInfo{}; + }; + + MatmulInfo info; + + // Peel lets + Expr value = op->value; + while (const Let *let = value.as()) { + info.peeled_lets.emplace_back(let->name, let->value); + value = let->body; + } + + // The RHS must be an add + const auto *add = value.as(); + if (!add) { + return fail("the right-hand-side is not an add"); + } + + // The add must be between a vector reduce and a load. The simplifier will + // have placed the vector reduce to the left, due to canonicalization of + // commutative ops. + const auto *reduce = add->a.as(); + if (!reduce || reduce->op != VectorReduce::Add) { + return fail("the right-hand-side is not a vector reduction plus a load"); + } + + // The load must be to the same addresses as the store (i.e. this is a +=) + const auto *load = add->b.as(); + if (!load || load->name != op->name || !equal(load->index, op->index)) { + return fail("the right-hand-side load is not from the same address as the store"); + } + + // There must be no predicate on the load or store + if (!is_const_one(load->predicate) || !is_const_one(op->predicate)) { + return fail("the load or store is predicated"); + } + + info.accumulator_type = reduce->type.element_of(); + if (!(info.accumulator_type == Float(32) || + info.accumulator_type == Float(16) || + info.accumulator_type == Int(32))) { + return fail("the accumulator type is not float32, float16, or int32"); + } + + // The vector reduce must be of a widening multiply. FindIntrinsics does + // not lift float widening muls, so we just expect a multiply of two casts. + Expr reduce_value = simplify(lower_intrinsics(reduce->value)); + const auto *mul = reduce_value.as(); + if (!mul) { + return fail("the vector reduction is not of a multiply"); + } + // Under the casts, broadcasts and lane permutations that vectorization may + // have left on each operand there must be a load. + Scope empty_scope; + info.lhs.load = is_load_of_multiramp(mul->a, empty_scope, &info.lhs.mr); + info.rhs.load = is_load_of_multiramp(mul->b, empty_scope, &info.rhs.mr); + if (!info.lhs.load || !info.rhs.load) { + return fail("the matrix multiply operands are not loads with affine indices"); + } + if (!is_const_one(info.lhs.load->predicate) || !is_const_one(info.rhs.load->predicate)) { + return fail("the matrix multiply operands are predicated loads"); + } + const Type operand_type = info.lhs.load->type.element_of(); + if (info.rhs.load->type.element_of() != operand_type) { + return fail("the matrix multiply operands do not have the same type"); + } + if (!wmma_types_supported(operand_type, info.accumulator_type)) { + return fail("there is no tensor core instruction that multiplies these " + "operands into an accumulator of this type"); + } + + // In a matrix multiply with row-major inputs and outputs, the algorithm + // looks like: + // + // C(j, i) += A(k, i) * B(j, k) + // + // (Recall that for matrices where the rows are stored densely in memory, + // Halide is indexed col-major.) The canonical loop nest order, from + // innermost out, is k, j, i, which is the coordinate system we compare the + // operands' access patterns in. In that coordinate system, i indexes the M + // rows of the output, j indexes its N columns, and k indexes the reduction. + int MN = reduce->type.lanes(); + int K = reduce->value.type().lanes() / MN; + + // Deduce which operand is which and what tile shape this is by trying each + // supported shape and seeing which one the access patterns fit. + const Shape *shape = nullptr; + for (const Shape &candidate : supported_shapes) { + if (candidate.M * candidate.N != MN || candidate.K != K) { + continue; + } + vector canonical_shape{candidate.K, candidate.N, candidate.M}; + if (!info.lhs.mr.strides_for_shape(canonical_shape, &info.lhs.strides) || + !info.rhs.mr.strides_for_shape(canonical_shape, &info.rhs.strides)) { + continue; + } + if (is_lhs(info.lhs, &info.lhs_layout, &info.lda) && + is_rhs(info.rhs, &info.rhs_layout, &info.ldb)) { + shape = &candidate; + break; + } + if (is_lhs(info.rhs, &info.lhs_layout, &info.lda) && + is_rhs(info.lhs, &info.rhs_layout, &info.ldb)) { + std::swap(info.lhs, info.rhs); + shape = &candidate; + break; + } + } + + if (!shape) { + return fail("the operands' access patterns do not describe a matrix " + "multiply of a tile shape the tensor cores support (16x16x16, " + "32x8x16, or 8x32x16)"); + } + info.shape = *shape; + return info; +} + +// Whether a store to a fragment fills it, as opposed to a matrix multiply +// accumulating into it. +bool is_fill(const Store *op) { + if (is_const_zero(op->value)) { + return true; + } + MultiRamp mr; + const Load *load = is_load_of_multiramp(op->value, Scope::empty_scope(), &mr); + return load && load->name != op->name; +} + +Stmt convert_to_tile_store(const Store *op, const Expr &store_index, + const string &new_name, const Shape &shape) { + auto fail = [&](const char *reason) { + user_error << "Store of a tensor core accumulator to memory not supported. " + << reason << ".\n" + << Stmt(op); + return Stmt{}; + }; + + if (!is_const_one(op->predicate)) { + return fail("The store has a predicate"); + } + MultiRamp mr; + if (!is_multiramp(store_index, Scope::empty_scope(), &mr)) { + return fail("The store index is not affine"); + } + WMMAMatrixLayout mem; + if (!wmma_matrix_layout(mr, shape.M, shape.N, &mem)) { + return fail("The store is not to a dense tile of the deduced matrix shape"); + } + Layout layout = mem.row_major ? Layout::Row : Layout::Col; + + // Each lane writes the entries of the matrix that it holds. The index + // enumerates the matrix in row-major order, which is the lane order of the + // accumulator. + Expr index = make_matrix_index(mem.base, shape.M, shape.N, layout, mem.stride); + Type element_type = op->value.type().element_of(); + Expr frag = Load::make(element_type.with_lanes(accumulator_elements), new_name, + Ramp::make(0, 1, accumulator_elements)); + const int lanes = shape.M * shape.N; + Expr lane = make_lane(unique_name("wmma_lane") + gpu_thread_name(0)); + Expr matrix = Call::make(element_type.with_lanes(lanes), Call::wmma_fragment_to_matrix_d, + {shape.M, shape.N, shape.K, std::move(frag), lane}, + Call::Intrinsic); + Expr owned = Call::make(UInt(1, lanes), Call::wmma_lane_owns, + {shape.M, shape.N, shape.K, lane}, Call::Intrinsic); + return in_lane_loop(lane, + Store::make(op->name, std::move(matrix), std::move(index), + op->param, std::move(owned), ModulusRemainder(), + op->is_streaming)); +} + +// Everything we learn about one tile allocation from the way it is +// used. The role and the shape come from the matrix multiplies it takes part +// in, so neither is known until those have been found. +struct Fragment { + string name; + // The prefix of the names of the per-fragment allocations this one becomes. + string fragment_name; + Type element_type; + Role role = Role::Unknown; + Shape shape{}; + bool found_shape = false; + // An allocation may hold several fragments as disjoint sub-tiles, each of + // which becomes its own allocation. + vector subtiles; + + Type value_type() const { + return element_type.with_lanes( + elements_per_lane(role, shape, element_type)); + } +}; + +class ExtractWMMAOperations : public IRMutator { + using IRMutator::visit; + + // Everything we've learned about each fragment allocation, and the names of + // the ones we're currently inside of. The records outlive the first pass so + // that the second one can use them. + std::map fragments; + vector in_scope; + + // In the first pass we recognize the matrix multiplies, which is what tells + // us what role each fragment plays and what shape it is. In the second we + // rewrite everything, which needs to know both. + int pass = 0; + + // The loops over GPU blocks, threads, and lanes we're inside of. + vector gpu_loop_vars; + + Fragment *find_fragment(const string &name) { + for (const string &n : in_scope) { + if (n == name) { + return &fragments[name]; + } + } + return nullptr; + } + + // A fragment allocation may sit outside the loops over GPU threads, in + // which case each thread gets its own copy of it and only ever touches its + // own slice. Any dependence of the index on the thread is selecting between + // those copies, not between subtiles within one, so drop it. + Expr index_within_thread(const Expr &index) { + Expr idx = index; + for (const string &v : gpu_loop_vars) { + idx = substitute(v, 0, idx); + } + return simplify(idx); + } + + string subtile_name(Fragment *f, const Expr &index) { + int idx = get_subtile(index_within_thread(index), + "tensor core fragment", &f->subtiles); + internal_assert(idx >= 0); // errors handled already + return f->fragment_name + std::to_string(idx); + } + + void set_role(Fragment *f, Role role) { + user_assert(f->role == Role::Unknown || f->role == role) + << "The tensor core fragment " << f->name << " is used as both " + << role_name(f->role) << " and " << role_name(role) << " of a matrix " + << "multiply. Those are held in registers in different layouts, so a " + << "fragment can only play one of those roles. Stage it through memory " + << "in between.\n"; + f->role = role; + } + + void set_shape(Fragment *f, const Shape &shape) { + user_assert(!f->found_shape || + (shape.M == f->shape.M && shape.N == f->shape.N && + shape.K == f->shape.K)) + << "Found inconsistent tile shapes for the tensor core fragment " + << f->name << " across the matrix multiplies that use it: " + << f->shape.M << "x" << f->shape.N << "x" << f->shape.K << " vs " + << shape.M << "x" << shape.N << "x" << shape.K << "."; + f->shape = shape; + f->found_shape = true; + } + + // Which subtile of a fragment an access refers to, as an index over just + // the tile. An operand is read by the multiply at the shape of the whole + // reduction, with one axis broadcast, so its index has to be projected back + // down to the tile before it can be compared against the fill that wrote + // it. + string operand_subtile_name(Fragment *f, const Expr &base, Role role, + const Shape &shape, Layout layout, const Expr &stride) { + int rows, cols; + fragment_matrix_shape(role, shape, &rows, &cols); + return subtile_name(f, make_matrix_index(base, rows, cols, layout, stride)); + } + + // Note what a matrix multiply tells us about an operand staged in a + // fragment. An operand read straight out of memory tells us nothing, + // because its load gets synthesized at the multiply. + void record_operand(const Operand &operand, Role role, const Shape &shape, + Layout layout, const Expr &stride) { + if (Fragment *f = find_fragment(operand.load->name)) { + set_role(f, role); + set_shape(f, shape); + operand_subtile_name(f, operand.mr.base, role, shape, layout, stride); + } + } + + // The value a matrix multiply uses for one of its operands: the fragment it + // was staged in, or a load synthesized here if it wasn't staged. + Expr operand_value(const Operand &operand, Role role, const Shape &shape, + Layout layout, const Expr &stride, const Expr &lane) { + if (Fragment *f = find_fragment(operand.load->name)) { + const int lanes = f->value_type().lanes(); + const string name = + operand_subtile_name(f, operand.mr.base, role, shape, layout, stride); + return Load::make(f->value_type(), name, Ramp::make(0, 1, lanes)); + } + return make_matrix_to_fragment(role, shape, layout, operand.load, + operand.mr.base, stride, lane); + } + + Stmt convert_to_fill(const Store *op, Fragment *f) { + int rows, cols; + fragment_matrix_shape(f->role, f->shape, &rows, &cols); + MultiRamp dest_mr; + WMMAMatrixLayout dest; + user_assert(is_multiramp(op->index, Scope::empty_scope(), &dest_mr) && + wmma_matrix_layout(dest_mr, rows, cols, &dest)) + << "A tensor core fragment must be filled a whole tile at a time, but " + << "this fill is not to a dense " << rows << "x" << cols << " tile of " + << f->name << ".\n" + << Stmt(op); + const string name = subtile_name( + f, make_matrix_index(dest.base, rows, cols, + dest.row_major ? Layout::Row : Layout::Col, dest.stride)); + const int lanes = f->value_type().lanes(); + Expr lane = make_lane(unique_name("wmma_lane") + gpu_thread_name(0)); + Expr value; + if (is_const_zero(op->value)) { + // Zeroing a fragment is layout-independent, so it doesn't need an + // instruction - the registers just get set to zero. + value = make_zero(f->value_type()); + } else { + auto fail = [&](const char *reason) { + user_error << "Fill of a tensor core fragment not supported. " + << reason << ".\n" + << Stmt(op); + return Expr{}; + }; + + MultiRamp mr; + const Load *matrix = + is_load_of_multiramp(op->value, Scope::empty_scope(), &mr); + internal_assert(matrix); // is_fill checked this + int rows, cols; + fragment_matrix_shape(f->role, f->shape, &rows, &cols); + WMMAMatrixLayout mem; + if (find_fragment(matrix->name)) { + value = fail("A fragment can only be filled from a matrix in memory, " + "not from another fragment, because the layout in " + "registers is not known"); + } else if (matrix->type.element_of() != f->element_type) { + value = fail("A fragment can only be filled from a matrix of the same " + "type, because the hardware does not convert on the way in"); + } else if (!is_const_one(matrix->predicate)) { + value = fail("The load is predicated"); + } else if (!wmma_matrix_layout(mr, rows, cols, &mem)) { + value = fail("The matrix loaded from is not a dense tile of the right " + "shape"); + } else { + value = make_matrix_to_fragment( + f->role, f->shape, mem.row_major ? Layout::Row : Layout::Col, + matrix, mem.base, mem.stride, lane); + } + } + return in_lane_loop( + lane, Store::make(name, std::move(value), Ramp::make(0, 1, lanes))); + } + + Stmt convert_to_matmul(const Store *op, Fragment *f, const MatmulInfo &info) { + Expr lane = make_lane(unique_name("wmma_lane") + gpu_thread_name(0)); + Expr a = operand_value(info.lhs, Role::A, info.shape, info.lhs_layout, info.lda, lane); + Expr b = operand_value(info.rhs, Role::B, info.shape, info.rhs_layout, info.ldb, lane); + + Type acc_type = info.accumulator_type.with_lanes(accumulator_elements); + Expr frag_idx = Ramp::make(0, 1, accumulator_elements); + const string name = subtile_name(f, op->index); + Expr c = Load::make(acc_type, name, frag_idx); + + Expr mma = Call::make(acc_type, Call::wmma_mma, + {info.shape.M, info.shape.N, info.shape.K, + (int)info.lhs_layout, (int)info.rhs_layout, + std::move(a), std::move(b), std::move(c), lane}, + Call::Intrinsic); + + Stmt store = in_lane_loop( + lane, Store::make(name, std::move(mma), frag_idx)); + for (const auto &[let_name, v] : reverse_view(info.peeled_lets)) { + store = LetStmt::make(let_name, v, store); + } + return store; + } + + Stmt visit(const For *op) override { + if (!is_gpu(op->for_type)) { + return IRMutator::visit(op); + } + gpu_loop_vars.push_back(op->name); + Stmt s = IRMutator::visit(op); + gpu_loop_vars.pop_back(); + return s; + } + + Stmt visit(const Allocate *op) override { + if (op->memory_type != MemoryType::Tile) { + return IRMutator::visit(op); + } + + user_assert(op->type == Float(32) || op->type == Float(16) || + op->type == Int(32)) + << "Tensor core fragments must hold 32-bit or 16-bit floats, or " + << "32-bit integers, but " << op->name << " holds " << op->type << ".\n"; + + Fragment &f = fragments[op->name]; + if (pass == 0) { + f.name = op->name; + f.fragment_name = op->name + ".wmma."; + f.element_type = op->type; + } + + in_scope.push_back(op->name); + Stmt body = mutate(op->body); + in_scope.pop_back(); + + if (pass == 0) { + user_assert(f.role != Role::Unknown) + << op->name << " is stored in Tile memory, but no matrix " + << "multiply was found that accumulates into it or reads it as an " + << "operand, so we can't tell what layout it should have.\n"; + return op; + } + + // Each fragment is one tile's worth of storage per lane. The allocations + // stay outside the loops over lanes, because a loop over lanes is a loop + // over threads, and register allocations outside a thread loop already + // get replicated per thread. + for (int i = 0; i < (int)f.subtiles.size(); i++) { + body = Allocate::make(f.fragment_name + std::to_string(i), f.element_type, + MemoryType::Tile, {f.value_type().lanes()}, + const_true(), body); + } + return body; + } + + Stmt visit(const Atomic *op) override { + if (find_fragment(op->producer_name)) { + // A tensor core fragment is per-thread register storage, so there's + // nothing for another thread to race with. The atomic is there + // because the fragment is scheduled outside the loops over threads, + // which makes it look shared. + user_assert(op->mutex_name.empty()) + << "Accumulating into a tensor core fragment should not need a " + << "mutex.\n"; + return mutate(op->body); + } + return IRMutator::visit(op); + } + + Stmt visit(const Free *op) override { + Fragment *f = find_fragment(op->name); + if (!f || pass == 0) { + return op; + } + Stmt s; + for (int i = 0; i < (int)f->subtiles.size(); i++) { + Stmt free = Free::make(f->fragment_name + std::to_string(i)); + s = s.defined() ? Block::make(std::move(s), std::move(free)) : std::move(free); + } + return s; + } + + Stmt visit(const ProducerConsumer *op) override { + Fragment *f = find_fragment(op->name); + if (!f) { + return IRMutator::visit(op); + } + return ProducerConsumer::make(f->fragment_name, op->is_producer, mutate(op->body)); + } + + Expr visit(const Load *op) override { + user_assert(!find_fragment(op->name)) + << "The tensor core fragment " << op->name + << " is used outside a tensor core instruction.\n"; + return IRMutator::visit(op); + } + + Stmt visit(const Store *op) override { + Fragment *f = find_fragment(op->name); + + if (!f) { + // A store to memory of a load from a fragment copies a tile out. + Expr store_index; + const Load *load = peel_store_permutations(op, &store_index).as(); + Fragment *src = load ? find_fragment(load->name) : nullptr; + if (src) { + if (pass == 0) { + subtile_name(src, load->index); + return op; + } + user_assert(src->role == Role::Accumulator) + << "The tensor core fragment " << src->name << " is copied out to " + << "memory, but it holds " << role_name(src->role) << " of a matrix " + << "multiply. Only an accumulator can be copied out.\n"; + return convert_to_tile_store(op, store_index, + subtile_name(src, load->index), src->shape); + } + // Not a copy of a tile out to memory. Recurse, so that any use of a + // fragment buried in here gets reported as an error. + return IRMutator::visit(op); + } + + if (is_fill(op)) { + // A fill needs to know the role and the shape, so it waits for the + // second pass. It tells us neither, so there's nothing to record in + // the first. + return pass == 0 ? Stmt(op) : convert_to_fill(op, f); + } + + MatmulInfo info = analyze_matmul(op); + if (pass == 0) { + set_role(f, Role::Accumulator); + set_shape(f, info.shape); + subtile_name(f, op->index); + record_operand(info.lhs, Role::A, info.shape, info.lhs_layout, info.lda); + record_operand(info.rhs, Role::B, info.shape, info.rhs_layout, info.ldb); + return op; + } + return convert_to_matmul(op, f, info); + } + +public: + void next_pass() { + pass = 1; + } +}; + +} // namespace + +Stmt extract_wmma_operations(const Stmt &s) { + ExtractWMMAOperations mutator; + // The first pass only looks. What it learns about a fragment from one use + // of it is needed at all the others, including the ones it has already + // walked past. + mutator(s); + mutator.next_pass(); + return mutator(s); +} + +bool is_wmma_intrinsic(const Call *op) { + return (op->is_intrinsic(Call::wmma_matrix_to_fragment_a) || + op->is_intrinsic(Call::wmma_matrix_to_fragment_b) || + op->is_intrinsic(Call::wmma_matrix_to_fragment_c) || + op->is_intrinsic(Call::wmma_mma)); +} + +Expr peel_store_permutations(const Store *op, Expr *index) { + Expr value = op->value; + *index = op->index; + while (const Shuffle *shuffle = value.as()) { + if (!shuffle->is_transpose()) { + break; + } + // Transposing both sides of the store cancels out. + const int rows = value.type().lanes() / shuffle->transpose_factor(); + *index = Shuffle::make_transpose(*index, rows); + value = shuffle->vectors[0]; + } + return value; +} + +bool wmma_matrix_layout(const MultiRamp &mr, int rows, int cols, + WMMAMatrixLayout *result) { + std::vector strides; + if (mr.total_lanes() != rows * cols || + !mr.strides_for_shape({cols, rows}, &strides)) { + return false; + } + // A row-major matrix has its columns one element apart. + if (is_const_one(strides[0])) { + result->row_major = true; + result->stride = strides[1]; + } else if (is_const_one(strides[1])) { + result->row_major = false; + result->stride = strides[0]; + } else { + return false; + } + result->base = mr.base; + return true; +} + +bool is_wmma_matrix_store(const Store *op) { + Expr index; + const Call *value = peel_store_permutations(op, &index).as(); + return value && value->is_intrinsic(Call::wmma_fragment_to_matrix_d); +} + +int wmma_matrix_arg(const Call *op) { + if (op->is_intrinsic(Call::wmma_matrix_to_fragment_a) || + op->is_intrinsic(Call::wmma_matrix_to_fragment_b) || + op->is_intrinsic(Call::wmma_matrix_to_fragment_c)) { + return 3; + } + return -1; +} + +} // namespace Internal +} // namespace Halide diff --git a/src/ExtractWMMAOperations.h b/src/ExtractWMMAOperations.h new file mode 100644 index 000000000000..78ffb7767c4a --- /dev/null +++ b/src/ExtractWMMAOperations.h @@ -0,0 +1,94 @@ +#ifndef HALIDE_EXTRACT_WMMA_OPERATIONS_H +#define HALIDE_EXTRACT_WMMA_OPERATIONS_H + +/** \file + * Defines the lowering pass that injects calls to the warp-level matrix + * multiply-accumulate intrinsics that drive NVIDIA tensor cores. + */ + +#include "Expr.h" +#include "MultiRamp.h" + +namespace Halide { +namespace Internal { + +struct Call; +struct Store; + +/** Rewrite matrix multiplies that accumulate into Tile memory as + * calls to the wmma intrinsics understood by the PTX backend, and wrap them in + * a loop over the 32 lanes of a warp. */ +Stmt extract_wmma_operations(const Stmt &s); + +/** Whether a Call is one of the wmma intrinsics produced by the pass above. */ +bool is_wmma_intrinsic(const Call *op); + +/** The index of the argument of a wmma intrinsic that gives the matrix it + * moves to or from memory, or -1 if it doesn't have one (which is the case for + * wmma_mma, whose operands are fragments rather than matrices). + * + * That argument is a Load of the whole matrix, with its lanes in row-major + * order regardless of how the matrix is laid out in memory. The layout in + * memory, and the distance between its rows or columns, are recoverable from + * the index of that Load: whichever dimension has stride one is the dense one. + * + * Being a Load rather than a bare handle to the allocation means that the + * passes which track uses of an allocation (dead allocation removal, shared + * memory liveness and packing, closure construction) see the access, and see + * its true footprint. It is never actually loaded - codegen takes its address + * - so passes that rewrite the *structure* of a load must leave it alone, or + * the tile shape will no longer be recoverable. */ +int wmma_matrix_arg(const Call *op); + +/** Whether a Store is the copy of a tensor core accumulator out to memory + * produced by the pass above. Such a store writes the whole matrix, with each + * lane of the warp writing the entries it holds: + * + * out[matrix] = wmma_fragment_to_matrix(M, N, K, fragment, lane) + * with predicate wmma_lane_owns(M, N, K, lane) + * + * wmma_fragment_to_matrix_d permutes this lane's fragment up into a whole + * matrix, leaving the entries the lane doesn't hold undefined, and + * wmma_lane_owns says whether this lane holds each entry of the matrix in its + * fragment. Both are opaque, because the mapping from matrix entry to lane + * isn't architecturally specified. The layout of the matrix in memory is + * recoverable from the index, as it is for a load. + * + * Passes that rewrite the structure of a store must leave these alone. */ +bool is_wmma_matrix_store(const Store *op); + +/** Peel the lane permutations the simplifier may have applied to the value of a + * store, moving the inverse of each onto the index, and return what's left of + * the value. Transposing both sides of a store cancels out, so this recovers + * the store as it was written. + * + * The simplifier rewrites a store whose index isn't dense into a dense store of + * a transposed value, so a wmma store to a column-major matrix arrives at the + * backend as a dense store of the transpose of that matrix. That's the same + * thing as a column-major store, which is one of the instructions, so the + * backend just undoes the rewrite rather than preventing it. */ +Expr peel_store_permutations(const Store *op, Expr *index); + +/** How a matrix a wmma instruction moves is laid out in memory. */ +struct WMMAMatrixLayout { + /** The address of its first entry. */ + Expr base; + /** The distance between its rows if it's row-major, or between its + * columns if it's column-major, in elements. */ + Expr stride; + bool row_major; +}; + +/** Read how a matrix is laid out in memory off the access pattern of the load + * or store of it: whichever dimension has stride one is the dense one, and the + * other stride is the distance between its rows or columns. The extraction pass + * uses this to build the access pattern, and the backend to read it back, so + * they can't disagree about the convention. Returns false if the access isn't a + * dense tile of the given shape. */ +bool wmma_matrix_layout(const MultiRamp &mr, int rows, int cols, + WMMAMatrixLayout *result); + +} // namespace Internal +} // namespace Halide + +#endif diff --git a/src/FlattenNestedRamps.cpp b/src/FlattenNestedRamps.cpp index 2cb552ec58d7..2e9e28a8eac8 100644 --- a/src/FlattenNestedRamps.cpp +++ b/src/FlattenNestedRamps.cpp @@ -2,6 +2,7 @@ #include "Bounds.h" #include "CSE.h" #include "Deinterleave.h" +#include "ExtractWMMAOperations.h" #include "IRMutator.h" #include "IROperator.h" #include "MultiRamp.h" @@ -63,6 +64,20 @@ class FlattenRamps : public IRMutator { return Shuffle::make_slice(v, n * inner_lanes, 1, inner_lanes); } + Expr visit(const Call *op) override { + if (wmma_matrix_arg(op) >= 0) { + // This pass turns a load of a nested ramp into a concatenation of + // dense loads, which is several load nodes rather than one, so + // unlike a reshaping shuffle it can't be undone by the backend. The + // matrix a wmma instruction takes a fragment out of is never + // actually loaded - only its address is used - so there's nothing + // to gain by rewriting it anyway. The other arguments are scalar + // constants, so there's nothing to do to them either. + return op; + } + return IRMutator::visit(op); + } + Expr visit(const Load *op) override { // Convert a load of a bounded span of indices into a shuffle // of a dense or strided load if possible. @@ -178,6 +193,14 @@ class FlattenRamps : public IRMutator { } Stmt visit(const Store *op) override { + if (is_wmma_matrix_store(op)) { + // One wmma store instruction writes the whole matrix for the whole + // warp, so it must not be broken into a store per row. Unlike the + // load side, this can't be undone afterwards, because it would have + // become several statements. + return op; + } + // If the index is a multiramp, unroll into a sequence of per-inner-ramp // stores, for the same reason as the Load visitor above. if (op->index.type().is_vector()) { diff --git a/src/Func.cpp b/src/Func.cpp index 6382b0c8ec26..fabfa33c4904 100644 --- a/src/Func.cpp +++ b/src/Func.cpp @@ -1712,6 +1712,49 @@ Stage &Stage::vectorize(const VarOrRVar &var) { return *this; } +namespace { + +// Reorder the given dimensions to be the innermost loops of a stage, in the +// order given, leaving the relative order of the others unchanged. Naming +// every dimension makes reorder produce this exact total order, rather than +// just permuting the named ones among the positions they already hold. The +// outermost sentinel is left off so that it stays put. +Stage &reorder_innermost(Stage &stage, const std::vector &vars) { + std::vector order = vars; + for (const VarOrRVar &v : stage.split_vars()) { + if (v.name() == Var::outermost().name()) { + continue; + } + bool named = std::any_of(vars.begin(), vars.end(), + [&](const VarOrRVar &w) { return w.name() == v.name(); }); + if (!named) { + order.push_back(v); + } + } + return stage.reorder(order); +} + +} // namespace + +Stage &Stage::tile_init(const VarOrRVar &x, const VarOrRVar &y) { + return reorder_innermost(*this, {x, y}).vectorize(x).vectorize(y); +} + +Stage &Stage::tile_load(const VarOrRVar &x, const VarOrRVar &y) { + return reorder_innermost(*this, {x, y}).vectorize(x).vectorize(y); +} + +Stage &Stage::tile_store(const VarOrRVar &x, const VarOrRVar &y) { + return reorder_innermost(*this, {x, y}).vectorize(x).vectorize(y); +} + +Stage &Stage::tile_matmul(const VarOrRVar &r, const VarOrRVar &x, const VarOrRVar &y) { + // Vectorizing a reduction dimension needs the stage marked atomic, which + // is asking to reassociate the sum rather than asking for atomic memory + // operations. + return reorder_innermost(atomic(), {r, x, y}).vectorize(r).vectorize(x).vectorize(y); +} + Stage &Stage::unroll(const VarOrRVar &var) { set_dim_type(var, ForType::Unrolled); return *this; @@ -2574,6 +2617,35 @@ Func &Func::parallel(const VarOrRVar &var) { return *this; } +Func &Func::tile_init(const VarOrRVar &x, const VarOrRVar &y) { + invalidate_cache(); + store_in(MemoryType::Tile); + Stage(func, func.definition(), 0).tile_init(x, y); + return *this; +} + +Func &Func::tile_load(const VarOrRVar &x, const VarOrRVar &y) { + invalidate_cache(); + store_in(MemoryType::Tile); + Stage(func, func.definition(), 0).tile_load(x, y); + return *this; +} + +// A tile_store reads a tile and writes somewhere more general, so this Func is +// the destination rather than the tile. +Func &Func::tile_store(const VarOrRVar &x, const VarOrRVar &y) { + invalidate_cache(); + Stage(func, func.definition(), 0).tile_store(x, y); + return *this; +} + +Func &Func::tile_matmul(const VarOrRVar &r, const VarOrRVar &x, const VarOrRVar &y) { + invalidate_cache(); + store_in(MemoryType::Tile); + Stage(func, func.definition(), 0).tile_matmul(r, x, y); + return *this; +} + Func &Func::vectorize(const VarOrRVar &var) { invalidate_cache(); Stage(func, func.definition(), 0).vectorize(var); diff --git a/src/Func.h b/src/Func.h index 46516e3e572e..703e22df9247 100644 --- a/src/Func.h +++ b/src/Func.h @@ -499,6 +499,47 @@ class Stage { } // @} + /** Schedule this stage as one operation of the target's matrix unit, + * which computes on whole tiles held in MemoryType::Tile. There is one + * method per thing such a unit can do: + * + * - tile_init fills a tile with a value uniform across it, usually zero. + * - tile_load copies a tile in from more general memory. + * - tile_store copies a tile back out to more general memory. On targets + * whose tiles live in general purpose registers this may also fold in + * elementwise work whose other operands are uniform across the tile; on + * targets whose tiles are a separate register file, such as x86, it + * cannot, because no elementwise instruction can read a tile. + * - tile_matmul multiplies a pair of tiles into this one, taking the + * reduction dimension first. Its operands come either from memory, in + * which case they are staged into tiles implicitly, or from Funcs held + * in MemoryType::Tile. + * + * The dimensions passed are the ones that make up a tile, innermost + * first. They are reordered into that order and vectorized. That is all + * these do: they are sugar over reorder and vectorize, plus atomic for the + * multiply, which is what permits a reduction dimension to be vectorized. + * Asking for it is asking to reassociate the sum, which a matrix unit does + * anyway, rather than asking for atomic memory operations. + * + * Which dimensions make up a tile, and their order, is fixed by the + * instruction rather than free to schedule, so these move them to be the + * innermost loops in the order given. The surrounding loops keep their + * relative order, and can be reordered separately either side of these + * calls. There is no need to name the tile dimensions in such a reorder. + * + * On CUDA these are warp-level, with an inner loop over the 32 lanes of a + * warp that the schedule does not name. That loop becomes the innermost + * GPU thread dimension, so schedule these outside the innermost thread + * loops of any other Func computed alongside them, and give those Funcs a + * 32-wide gpu_lanes loop of their own to match. */ + // @{ + Stage &tile_init(const VarOrRVar &x, const VarOrRVar &y); + Stage &tile_load(const VarOrRVar &x, const VarOrRVar &y); + Stage &tile_store(const VarOrRVar &x, const VarOrRVar &y); + Stage &tile_matmul(const VarOrRVar &r, const VarOrRVar &x, const VarOrRVar &y); + // @} + /** Get the Vars and RVars of this definition, from innermost out, with * splits applied. This represents all the potentially-valid compute_at * sites for this Stage. The RVars returned will be symbolic and not tied to @@ -1521,6 +1562,17 @@ class Func { * innermost one. */ Func &vectorize(const VarOrRVar &var); + /** Schedule this Func as one operation of the target's matrix unit. See + * Stage::tile_init. Everything but tile_store also sets this Func's + * memory type to MemoryType::Tile, since those are the ones that produce + * a tile; tile_store reads one and writes somewhere more general. */ + // @{ + Func &tile_init(const VarOrRVar &x, const VarOrRVar &y); + Func &tile_load(const VarOrRVar &x, const VarOrRVar &y); + Func &tile_store(const VarOrRVar &x, const VarOrRVar &y); + Func &tile_matmul(const VarOrRVar &r, const VarOrRVar &x, const VarOrRVar &y); + // @} + /** Mark a dimension to be completely unrolled. The dimension * should have constant extent - e.g. because it is the inner * dimension following a split by a constant factor. For most uses diff --git a/src/FuseGPUThreadLoops.cpp b/src/FuseGPUThreadLoops.cpp index f1a02e976545..3d1ce92b971d 100644 --- a/src/FuseGPUThreadLoops.cpp +++ b/src/FuseGPUThreadLoops.cpp @@ -29,6 +29,97 @@ using std::vector; namespace { +// Being filled by the copy engine is a property of the stores to an +// allocation, not of the memory it lives in, so MemoryType::GPUSharedAsync is +// only the schedule's way of saying it. Rewrite such allocations to ordinary +// shared memory, wrapping the value of each store to them in a +// cuda_bypass_registers intrinsic and awaiting the copies where the data is +// consumed. Everything below then sees one kind of shared memory, which is +// what lets all of a kernel's shared allocations be packed together. +class MarkAsyncCopies : public IRMutator { + using IRMutator::visit; + + // The allocations being filled by the copy engine, and the group each + // one's copies belong to. Copies in a group are awaited together. + map groups; + int next_group = 0; + + // Only CUDA has a copy engine to drive. Elsewhere the allocation still + // becomes ordinary shared memory, but the stores to it stay ordinary too. + DeviceAPI device_api = DeviceAPI::None; + + Stmt visit(const For *op) override { + ScopedValue d(device_api, op->device_api == DeviceAPI::None ? + device_api : + op->device_api); + return IRMutator::visit(op); + } + + Stmt visit(const Allocate *op) override { + if (op->memory_type != MemoryType::GPUSharedAsync) { + return IRMutator::visit(op); + } + + // One group per allocation, so that consuming one Func doesn't wait + // for the copies into another. + auto [it, inserted] = groups.emplace(op->name, next_group++); + internal_assert(inserted) + << "Two asynchronously copied allocations are both named " << op->name << "\n"; + + Stmt body = mutate(op->body); + groups.erase(it); + + return Allocate::make(op->name, op->type, MemoryType::GPUShared, + op->extents, mutate(op->condition), body, + op->new_expr, op->free_function, op->padding); + } + + Stmt visit(const Store *op) override { + auto it = groups.find(op->name); + if (it == groups.end() || device_api != DeviceAPI::CUDA) { + return IRMutator::visit(op); + } + // The Func name is carried along because the allocations get packed + // together below, after which the store no longer knows which Func it + // belongs to, and that is the name an error has to name. + Expr value = Call::make(op->value.type(), Call::cuda_bypass_registers, + {mutate(op->value), it->second, StringImm::make(op->name)}, + Call::Intrinsic); + return Store::make(op->name, value, mutate(op->index), op->param, + mutate(op->predicate), op->alignment, op->is_streaming); + } + + Stmt visit(const ProducerConsumer *op) override { + Stmt body = mutate(op->body); + auto it = groups.find(op->name); + if (op->is_producer && it != groups.end() && device_api == DeviceAPI::CUDA) { + // At the end of the producer, which is before the barrier that + // publishes the data to the rest of the block, and before any of + // it is read. + Expr wait = Call::make(Int(32), Call::cuda_await_copies, + {it->second}, Call::Intrinsic); + body = Block::make(body, Evaluate::make(wait)); + } + return op->with(body); + } + +public: + using IRMutator::mutate; +}; + +// The copy engine moves up to 16 bytes at a time, and needs its destination +// aligned to the width of the copy. Allocations are packed one after another, +// so a group has to end on a 16-byte boundary for whatever follows it to be +// copyable into. Round a group's size up accordingly, given the size in bytes +// of the units it is measured in. +Expr round_up_group_size(const Expr &size, int unit_bytes) { + const int alignment = 16; + if (unit_bytes >= alignment) { + return size; + } + return align_up(size, alignment / unit_bytes); +} + class ExtractBlockSize : public IRVisitor { protected: Expr block_extent[3], block_count[3]; @@ -227,15 +318,17 @@ class ReplaceForWithIf : public IRMutator { // An allocation inside the thread loops stays in register/local memory // (handled by ExtractRegisterAllocations) rather than being pulled out to the // block level (handled by ExtractSharedAndHeapAllocations) if it has a fixed -// size or an explicit register/stack memory type. +// size or an explicit register/stack/fragment memory type. A tensor core +// fragment is per-lane, so it belongs in registers wherever it was declared. bool allocation_goes_to_registers(const Allocate *op, bool in_threads) { bool fixed_size_thread_allocation = (op->constant_allocation_size() != 0) && in_threads; return (fixed_size_thread_allocation && op->memory_type != MemoryType::Heap && - op->memory_type != MemoryType::GPUShared && + !is_gpu_shared(op->memory_type) && op->memory_type != MemoryType::GPUTexture) || op->memory_type == MemoryType::Register || - op->memory_type == MemoryType::Stack; + op->memory_type == MemoryType::Stack || + op->memory_type == MemoryType::Tile; } // Rename an allocation and all of its loads, stores, and frees. Relies on the @@ -476,7 +569,7 @@ class ExtractSharedAndHeapAllocations : public IRMutator { } user_assert(op->memory_type == MemoryType::Auto || - op->memory_type == MemoryType::GPUShared || + is_gpu_shared(op->memory_type) || op->memory_type == MemoryType::GPUTexture || op->memory_type == MemoryType::Heap) << "Allocation " << op->name << " must live in shared or heap memory, " @@ -878,7 +971,7 @@ class ExtractSharedAndHeapAllocations : public IRMutator { int ratio = alloc.widest_type.bytes() / alloc_type.bytes(); internal_assert(ratio != 0) << "alloc_type should have been at most as wide as the widest type in group\n"; - total_size += alloc.max_size * ratio; + total_size += round_up_group_size(alloc.max_size * ratio, alloc_type.bytes()); } // Upgrade the alloc type to the widest type found, and @@ -953,7 +1046,8 @@ class ExtractSharedAndHeapAllocations : public IRMutator { offset = Variable::make(Int(32), name + "." + std::to_string(i - 1) + ".offset"); int ratio = (widest_type.bytes() / cluster[i - 1].widest_type.bytes()); internal_assert(ratio != 0); - offset += simplify((cluster[i - 1].max_size + ratio - 1) / ratio); + offset += simplify(round_up_group_size( + (cluster[i - 1].max_size + ratio - 1) / ratio, widest_type.bytes())); } else { if (memory_type == MemoryType::Heap) { // One slice of a larger global allocation @@ -1160,7 +1254,8 @@ class ExtractRegisterAllocations : public IRMutator { user_assert(op->memory_type == MemoryType::Stack || op->memory_type == MemoryType::Register || op->memory_type == MemoryType::Heap || - op->memory_type == MemoryType::Auto) + op->memory_type == MemoryType::Auto || + op->memory_type == MemoryType::Tile) << "Allocation " << op->name << " is scheduled inside a loop over GPU threads, so " << "it must live in stack memory, heap memory, or registers. " << "Shared allocations at this loop level are not yet supported.\n"; @@ -1362,6 +1457,7 @@ class InjectThreadBarriers : public IRMutator { debug(4) << "Encountered store to " << op->name << "\n"; auto mem_type = memory_type_for_name(op->name); switch (mem_type) { + case MemoryType::GPUSharedAsync: case MemoryType::GPUShared: debug(4) << " memory type is shared\n"; shared_stores.insert(op->name); @@ -1376,7 +1472,7 @@ class InjectThreadBarriers : public IRMutator { case MemoryType::Register: case MemoryType::LockedCache: case MemoryType::VTCM: - case MemoryType::AMXTile: + case MemoryType::Tile: break; } @@ -1387,6 +1483,7 @@ class InjectThreadBarriers : public IRMutator { debug(4) << "Encountered load from " << op->name << "\n"; auto mem_type = memory_type_for_name(op->name); switch (mem_type) { + case MemoryType::GPUSharedAsync: case MemoryType::GPUShared: debug(4) << " memory type is shared\n"; shared_loads.insert(op->name); @@ -1401,7 +1498,7 @@ class InjectThreadBarriers : public IRMutator { case MemoryType::Register: case MemoryType::LockedCache: case MemoryType::VTCM: - case MemoryType::AMXTile: + case MemoryType::Tile: break; } @@ -1689,6 +1786,9 @@ Stmt fuse_gpu_thread_loops(Stmt s) { // into the innermost GPU block. FuseGPUThreadLoops would then // merge the predicate into the merged GPU thread. s = NormalizeIfStatements()(s); + // Must run before the allocations are packed together, because packing + // them relies on there being only one kind of shared memory. + s = MarkAsyncCopies().mutate(s); s = FuseGPUThreadLoops()(s); s = ZeroGPULoopMins()(s); return s; diff --git a/src/IR.cpp b/src/IR.cpp index 72058b62700d..61f59172e143 100644 --- a/src/IR.cpp +++ b/src/IR.cpp @@ -817,6 +817,8 @@ constexpr const char *intrinsic_op_names[] = { "concat_bits", "count_leading_zeros", "count_trailing_zeros", + "cuda_await_copies", + "cuda_bypass_registers", "debug_to_file", "declare_box_touched", "div_round_to_zero", @@ -902,6 +904,12 @@ constexpr const char *intrinsic_op_names[] = { "widening_shift_left", "widening_shift_right", "widening_sub", + "wmma_fragment_to_matrix_d", + "wmma_lane_owns", + "wmma_matrix_to_fragment_a", + "wmma_matrix_to_fragment_b", + "wmma_matrix_to_fragment_c", + "wmma_mma", // keep-sorted end }; diff --git a/src/IR.h b/src/IR.h index f260890a48e1..73d676d3a929 100644 --- a/src/IR.h +++ b/src/IR.h @@ -724,6 +724,20 @@ struct Call : public ExprNode { concat_bits, count_leading_zeros, count_trailing_zeros, + // cuda_await_copies(group) waits for every asynchronous copy in the + // given group to have landed in shared memory. Appears in an Evaluate + // node at the point where the copied data is first read. + cuda_await_copies, + // cuda_bypass_registers(value, group, func_name) marks a value that must be moved + // to its destination without being materialized in registers, which the + // CUDA backend does with the copy engine. It is only valid as the whole + // value of a Store, and promises that the store is a copy the copy + // engine can make: a dense, aligned, unpredicated run of 4, 8 or 16 + // bytes read from outside the kernel. `group` names a batch of such + // copies that are waited for together by cuda_await_copies, and + // `func_name` is the Func the store belonged to before the shared + // allocations were packed together, for error messages. + cuda_bypass_registers, debug_to_file, // Declares that a box region of an allocation has been touched (used by bounds inference) declare_box_touched, @@ -878,6 +892,45 @@ struct Call : public ExprNode { widening_shift_left, widening_shift_right, widening_sub, + // Every intrinsic below takes the lane of the warp it runs in, because + // its value depends on which lane that is and nothing else in its + // arguments does. Without it a call whose other arguments are all + // constants, such as wmma_lane_owns, could be hoisted out of the loop + // over lanes. They are Intrinsic rather than PureIntrinsic only to stop + // lets being hoisted out from under them, which makes instruction + // selection in the backend harder. + // + // Permute this lane's tensor core accumulator (d) fragment up into a + // whole matrix, leaving the entries this lane doesn't hold undefined. + // wmma_fragment_to_matrix_d(M, N, K, fragment, lane) + wmma_fragment_to_matrix_d, + // Whether this lane holds each entry of an M x N tensor core + // accumulator in its fragment. Used as the predicate of the store that + // copies an accumulator out to memory. + // wmma_lane_owns(M, N, K, lane) + wmma_lane_owns, + // Take this lane's share of a tensor core fragment out of a matrix. + // The a operand is M x K, the b operand is K x N, and the accumulator + // (c) is M x N. The mapping from matrix entry to lane is opaque, + // because it isn't architecturally specified. + // + // The matrix argument is a Load of the whole of it, with its lanes in + // row-major order; the layout in memory and the distance between its + // rows or columns are recoverable from the index. The hardware can only + // do this as part of a memory read, which is why the argument has to be + // a Load rather than an arbitrary matrix value. + // wmma_matrix_to_fragment_a(M, N, K, matrix, lane) + // @{ + wmma_matrix_to_fragment_a, + wmma_matrix_to_fragment_b, + wmma_matrix_to_fragment_c, + // @} + // A tensor core matrix multiply-accumulate. The layouts say how the a + // and b fragments were taken out of their matrices, which changes how + // the hardware arranges them in registers. The accumulator's fragment + // layout doesn't depend on how it was loaded, so it isn't a parameter. + // wmma_mma(M, N, K, a_layout, b_layout, a, b, c, lane) + wmma_mma, // keep-sorted end IntrinsicOpCount // Sentinel: keep last. }; diff --git a/src/IRPrinter.cpp b/src/IRPrinter.cpp index 2e4f94d7288d..f4e2b9125c02 100644 --- a/src/IRPrinter.cpp +++ b/src/IRPrinter.cpp @@ -157,6 +157,9 @@ std::ostream &operator<<(std::ostream &out, const MemoryType &t) { case MemoryType::Register: out << "Register"; break; + case MemoryType::GPUSharedAsync: + out << "GPUSharedAsync"; + break; case MemoryType::GPUShared: out << "GPUShared"; break; @@ -169,8 +172,8 @@ std::ostream &operator<<(std::ostream &out, const MemoryType &t) { case MemoryType::VTCM: out << "VTCM"; break; - case MemoryType::AMXTile: - out << "AMXTile"; + case MemoryType::Tile: + out << "Tile"; break; } return out; diff --git a/src/Lower.cpp b/src/Lower.cpp index 409ae847baa3..816a08912bdf 100644 --- a/src/Lower.cpp +++ b/src/Lower.cpp @@ -26,6 +26,7 @@ #include "Deinterleave.h" #include "EarlyFree.h" #include "ExtractTileOperations.h" +#include "ExtractWMMAOperations.h" #include "FindCalls.h" #include "FindIntrinsics.h" #include "FlattenNestedRamps.h" @@ -352,6 +353,12 @@ void lower_impl(const vector &output_funcs, s = simplify(s); log("Lowering after vectorizing:", s); + if (t.has_feature(Target::CUDA)) { + debug(1) << "Extracting tensor core operations...\n"; + s = extract_wmma_operations(s); + log("Lowering after extracting tensor core operations:", s); + } + if (t.has_gpu_feature() || t.has_feature(Target::Vulkan)) { debug(1) << "Injecting per-block gpu synchronization...\n"; diff --git a/src/LowerWarpShuffles.cpp b/src/LowerWarpShuffles.cpp index 211aa7b5c422..79d4278f27bf 100644 --- a/src/LowerWarpShuffles.cpp +++ b/src/LowerWarpShuffles.cpp @@ -653,11 +653,14 @@ class LowerWarpShuffles : public IRMutator { Stmt visit(const Allocate *op) override { if (this_lane.defined() || - op->memory_type == MemoryType::GPUShared || - op->memory_type == MemoryType::Heap) { - // Not a warp-level allocation. Warp-level storage is per-lane - // register storage; shared and heap (global) memory are never - // striped across lanes. + is_gpu_shared(op->memory_type) || + op->memory_type == MemoryType::Heap || + op->memory_type == MemoryType::Tile) { + // Not an allocation for us to stripe. Warp-level storage is + // per-lane register storage; shared and heap (global) memory are + // never striped across lanes, and tensor core accumulators are + // already striped across lanes in a layout that only the wmma + // instructions understand. return IRMutator::visit(op); } else { // Pick up this allocation and deposit it inside the loop over lanes at reduced size. diff --git a/src/MultiRamp.cpp b/src/MultiRamp.cpp index 3bd52e14d99e..d7277511e81a 100644 --- a/src/MultiRamp.cpp +++ b/src/MultiRamp.cpp @@ -34,6 +34,40 @@ void collapse_adjacent_dims(MultiRamp *m) { } } +// Walk two shapes innermost-out, emitting the coarsest shape that groups into +// both of them. Returns false if there is no such shape, which happens when a +// pair of lane counts is coprime. If from_a is non-null it receives, for each +// refined dim, the index of the dim of `a` it came from. +bool common_refinement(const std::vector &a, const std::vector &b, + std::vector *refined, std::vector *from_a) { + refined->clear(); + if (from_a) { + from_a->clear(); + } + size_t i = 0, j = 0; + int x = a.empty() ? 1 : a[0]; + int y = b.empty() ? 1 : b[0]; + while (i < a.size() && j < b.size()) { + int g = gcd(x, y); + if (g == 1) { + return false; + } + refined->push_back(g); + if (from_a) { + from_a->push_back((int)i); + } + x /= g; + y /= g; + if (x == 1 && ++i < a.size()) { + x = a[i]; + } + if (y == 1 && ++j < b.size()) { + y = b[j]; + } + } + return true; +} + } // namespace MultiRamp::MultiRamp(Expr base, std::vector strides, std::vector lanes) @@ -62,68 +96,32 @@ void MultiRamp::mul(const Expr &e) { // common refinement (the sum is not a multiramp). Adding multiramps with // different total lane counts is a caller error and triggers an assertion. bool MultiRamp::add(const MultiRamp &other) { - // We walk through both ramps' dimensions innermost-to-outermost, consuming - // gcd(a_lanes, b_lanes) of lanes at a time. When a dimension is only - // partially consumed, the remaining part of that dimension corresponds to - // an "outer" sub-dim in the refined shape and its stride must be scaled - // by the factor just consumed. + // Refine the two shapes until they agree, express both in that shape, and + // add the strides elementwise. internal_assert(total_lanes() == other.total_lanes()) << "MultiRamp::add: total lane counts must match (" << total_lanes() << " vs " << other.total_lanes() << ")"; + Expr new_base = simplify(base + other.base); if (lanes.empty()) { // Both are 0-dim scalars. - base = simplify(base + other.base); + base = new_base; return true; } - MultiRamp result; - result.base = simplify(base + other.base); - size_t ai = 0, bi = 0; - int a_lanes = lanes[0], b_lanes = other.lanes[0]; - Expr a_stride = strides[0], b_stride = other.strides[0]; - while (true) { - int next_lanes = gcd(a_lanes, b_lanes); - if (next_lanes == 1) { - // The two next lanes are coprime, e.g: - // [0, 1, 2, 100, 101, 102] + [0, 1, 100, 101, 200, 201] - // which has no common refinement. - return false; - } - result.strides.emplace_back(simplify(a_stride + b_stride)); - result.lanes.push_back(next_lanes); - a_lanes /= next_lanes; - b_lanes /= next_lanes; - bool a_done = false, b_done = false; - if (a_lanes == 1) { - ai++; - if (ai >= lanes.size()) { - a_done = true; - } else { - a_lanes = lanes[ai]; - a_stride = strides[ai]; - } - } else { - // Remaining portion of current A-dim has a scaled stride. - a_stride = simplify(a_stride * next_lanes); - } - if (b_lanes == 1) { - bi++; - if (bi >= other.lanes.size()) { - b_done = true; - } else { - b_lanes = other.lanes[bi]; - b_stride = other.strides[bi]; - } - } else { - b_stride = simplify(b_stride * next_lanes); - } - if (a_done && b_done) { - collapse_adjacent_dims(&result); - *this = std::move(result); - return true; - } - // The up-front lane-count check ensures both sides always exhaust - // together, so neither side should be done here. + std::vector refined; + if (!common_refinement(lanes, other.lanes, &refined, nullptr)) { + // e.g. [0, 1, 2, 100, 101, 102] + [0, 1, 100, 101, 200, 201] + return false; + } + std::vector a_strides, b_strides; + bool ok = (strides_for_shape(refined, &a_strides) && + other.strides_for_shape(refined, &b_strides)); + internal_assert(ok) << "a common refinement should group into both shapes\n"; + std::vector new_strides(refined.size()); + for (size_t i = 0; i < refined.size(); i++) { + new_strides[i] = simplify(a_strides[i] + b_strides[i]); } + *this = MultiRamp(new_base, new_strides, refined); + return true; } bool MultiRamp::strides_for_shape(const std::vector &target_lanes, @@ -410,6 +408,52 @@ std::optional unbroadcast(const Expr &e) { } } +// Recognize a list of constant lane indices as a MultiRamp of constants, which +// is what it means for a shuffle to be a reshaping of its input rather than an +// arbitrary gather. +bool multiramp_of_constants(const std::vector &idx, Type t, MultiRamp *result) { + const int n = (int)idx.size(); + internal_assert(n > 0); + if (n == 1) { + *result = MultiRamp(make_const(t, idx[0]), {}, {}); + return true; + } + + // The innermost dim is the longest prefix that's an arithmetic progression. + const int stride = idx[1] - idx[0]; + int extent = 1; + while (extent < n && idx[extent] == idx[0] + extent * stride) { + extent++; + } + if (n % extent) { + return false; + } + + // Every block of that length has to be the same progression. + std::vector starts; + starts.reserve(n / extent); + for (int b = 0; b < n; b += extent) { + for (int j = 0; j < extent; j++) { + if (idx[b + j] != idx[b] + j * stride) { + return false; + } + } + starts.push_back(idx[b]); + } + + MultiRamp outer; + if (!multiramp_of_constants(starts, t, &outer)) { + return false; + } + + std::vector strides{make_const(t, stride)}; + strides.insert(strides.end(), outer.strides.begin(), outer.strides.end()); + std::vector lanes{extent}; + lanes.insert(lanes.end(), outer.lanes.begin(), outer.lanes.end()); + *result = MultiRamp(outer.base, strides, lanes); + return true; +} + // Internal is_multiramp. May leave *result in a partial state on failure; // the public is_multiramp below protects callers by only committing on // success. Recursive calls go through the public wrapper, so each branch @@ -429,6 +473,19 @@ bool is_multiramp_impl(const Expr &e, const Scope &scope, MultiRamp *resul result->strides.push_back(make_zero(elem_t)); result->lanes.push_back(b->lanes); return true; + } else if (const Shuffle *s = e.as(); s && s->vectors.size() == 1) { + // A shuffle of a single vector is a reshaping of it, rather than a + // gather, when its lane indices are themselves a multiramp that + // permutes them. This is the shape flatten_nested_ramps leaves a + // strided load in, and what a transpose of a tile looks like. + MultiRamp inner, perm; + if (is_multiramp(s->vectors[0], scope, &inner) && + multiramp_of_constants(s->indices, inner.base.type(), &perm) && + inner.shuffle(perm)) { + *result = inner; + return true; + } + return false; } else if (const Ramp *r = e.as()) { if (auto stride = unbroadcast(r->stride)) { if (is_multiramp(r->base, scope, result)) { @@ -478,6 +535,7 @@ bool is_multiramp_impl(const Expr &e, const Scope &scope, MultiRamp *resul } } } + return false; } } // namespace @@ -493,6 +551,119 @@ bool is_multiramp(const Expr &e, const Scope &scope, MultiRamp *result) { return false; } +namespace { + +// Strip the cast, broadcasts and lane permutations off a load, moving the +// ones that rearrange lanes onto a copy of the load's index, where +// is_multiramp can make sense of them. A broadcast of a load is a load of a +// broadcast of the index, and likewise for a lane permutation. +// +// At most one cast is peeled, so that comparing the element type of the +// original Expr against the type of the Load tells the caller whether the +// values were cast, and to what. +const Load *peel_load(const Expr &e, Expr *index, bool cast_allowed) { + if (const Cast *cast = e.as()) { + if (!cast_allowed) { + return nullptr; + } + return peel_load(cast->value, index, false); + } else if (const Broadcast *broadcast = e.as()) { + const Load *load = peel_load(broadcast->value, index, cast_allowed); + if (load) { + *index = Broadcast::make(*index, broadcast->lanes); + } + return load; + } else if (const Shuffle *shuffle = e.as()) { + if (shuffle->vectors.size() != 1) { + return nullptr; + } + const Load *load = peel_load(shuffle->vectors[0], index, cast_allowed); + if (load) { + // Shuffling the values loaded is the same as shuffling the + // addresses loaded from. + *index = Shuffle::make({*index}, shuffle->indices); + } + return load; + } else if (const Load *load = e.as()) { + *index = load->index; + return load; + } + return nullptr; +} + +} // namespace + +int get_subtile(const Expr &index, const std::string &description, + std::vector *subtiles) { + MultiRamp mr; + if (!is_multiramp(index, Scope::empty_scope(), &mr)) { + user_error << "Access to " << description << " not affine: " << index << "\n"; + } + if (!can_prove(mr.alias_free())) { + user_error << "Access to " << description << " may have duplicated lanes: " + << index << "\n"; + } + if (subtiles->empty()) { + subtiles->push_back(std::move(mr)); + return 0; + } + + // All strides and lanes must match across all subtiles, or we give up. + const MultiRamp &first = (*subtiles)[0]; + if (mr.dimensions() != first.dimensions()) { + user_error << "Access to " << description << " does not have the same shape as " + << "other accesses to the same memory."; + } + for (int i = 0; i < first.dimensions(); i++) { + if (!can_prove(mr.strides[i] == first.strides[i]) || + mr.lanes[i] != first.lanes[i]) { + user_error << "Access to " << description << " has different size and strides " + << "to other accesses to the same memory. All accesses must have " + << "the same subtile size and strides: " << index; + } + } + + // Now check for disjointedness. Add a synthetic dimension, the purpose of + // which will become clear. + mr.strides.emplace_back(); + mr.lanes.push_back(2); + for (int i = 0; i < (int)subtiles->size(); i++) { + const MultiRamp &other = (*subtiles)[i]; + // One of two things must be true: + // 1) All of the lanes of mr equal the corresponding lane of other. + // We've already checked the strides and lanes, so it's just a matter of + // checking the base. + if (can_prove(mr.base == other.base)) { + return i; + } + + // 2) None of the lanes of mr equal any of the lanes of other. To do + // this we construct a combined mr that can be either 'mr' or 'other', + // and ask if it's alias-free. This is what the synthetic dimension was + // for. + mr.strides.back() = mr.base - other.base; + if (!can_prove(mr.alias_free())) { + user_error << "Failed to prove access to " << description << " does not " + << "partially overlap another distinct access: " << index; + } + } + + // Didn't already exist and didn't alias with anything. + mr.strides.pop_back(); + mr.lanes.pop_back(); + subtiles->push_back(std::move(mr)); + return (int)subtiles->size() - 1; +} + +const Load *is_load_of_multiramp(const Expr &e, const Scope &scope, MultiRamp *result) { + Expr index; + const Load *load = peel_load(e, &index, true); + if (load && is_multiramp(index, scope, result)) { + return load; + } + return nullptr; +} + Expr MultiRamp::operator==(const MultiRamp &other) const { // Construct the difference, and check if all strides are zero. MultiRamp diff = other; @@ -577,6 +748,71 @@ std::vector MultiRamp::alias_free_slice() { return peeled; } +bool MultiRamp::shuffle(const MultiRamp &mask) { + if (mask.total_lanes() != total_lanes()) { + return false; + } + // A permutation of [0, n) starts at zero and steps by constants. + auto base = as_const_int(simplify(mask.base)); + if (!base || *base != 0) { + return false; + } + std::vector mask_strides; + for (const Expr &e : mask.strides) { + auto s = as_const_int(simplify(e)); + if (!s) { + return false; + } + mask_strides.push_back(*s); + } + + // Sorting the mask's dims by stride has to give the prefix products of + // their lane counts. Anything else walks the lane index in a way that + // isn't a reshaping of it. + std::vector order(mask_strides.size()); + std::iota(order.begin(), order.end(), 0); + std::sort(order.begin(), order.end(), + [&](int a, int b) { return mask_strides[a] < mask_strides[b]; }); + std::vector mask_shape; + int64_t expected = 1; + for (int d : order) { + if (mask_strides[d] != expected) { + return false; + } + mask_shape.push_back(mask.lanes[d]); + expected *= mask.lanes[d]; + } + + // Refine the mask's shape and ours against each other, tracking which of + // the mask's dims each refined dim came from. + std::vector refined, from_mask_dim; + if (!common_refinement(mask_shape, lanes, &refined, &from_mask_dim)) { + return false; + } + // A common refinement groups into our dims, so this cannot fail. + std::vector refined_strides; + bool ok = strides_for_shape(refined, &refined_strides); + internal_assert(ok) << "refined shape does not group into the multiramp's dims\n"; + + // Emit the refined dims grouped and ordered the way the mask reads them. + std::vector sorted_pos(order.size()); + for (int k = 0; k < (int)order.size(); k++) { + sorted_pos[order[k]] = k; + } + std::vector new_strides; + std::vector new_lanes; + for (int d = 0; d < mask.dimensions(); d++) { + for (size_t k = 0; k < refined.size(); k++) { + if (from_mask_dim[k] == sorted_pos[d]) { + new_strides.push_back(refined_strides[k]); + new_lanes.push_back(refined[k]); + } + } + } + *this = MultiRamp(this->base, new_strides, new_lanes); + return true; +} + int MultiRamp::rotate_stride_one_innermost() { int k = -1; for (int i = 0; i < dimensions(); i++) { diff --git a/src/MultiRamp.h b/src/MultiRamp.h index 583d257d93f4..a720b04fcdd8 100644 --- a/src/MultiRamp.h +++ b/src/MultiRamp.h @@ -14,6 +14,7 @@ namespace Internal { class IRMutator; class IRVisitor; +struct Load; /** A multi-dimensional ramp. I.e. a ramp of ramps of ramps of ramps... * @@ -151,6 +152,19 @@ struct MultiRamp { * a vector in the old lane order from one in the new order. */ int rotate_stride_one_innermost(); + /** Permute the lanes by a mask. `mask` gives, for each lane of the + * result, which lane of *this to take, so its lanes must be a permutation + * of [0, total_lanes()). That is the case exactly when its base is zero + * and its strides are the prefix products of its lane counts in some + * order, which makes it a reshaping of the lane index rather than a + * gather - a transpose being the two-dimensional example. + * + * The mask's dims and ours are refined against each other until both are + * groupings of a common shape, and the result is that shape reordered the + * way the mask asks for. Returns false, leaving *this unchanged, if the + * mask isn't a permutation or the two shapes have no common refinement. */ + bool shuffle(const MultiRamp &mask); + /** The dimensionality. May be lower than you expected, because this * gets flattened when possible by the operations above. */ int dimensions() const; @@ -210,10 +224,37 @@ struct MultiRamp { const std::vector &pos) const; }; +/** Locate the subtile of a tile-memory allocation that an access refers to. + * Such an allocation may hold several tile registers as disjoint sub-tiles, and + * each one becomes its own allocation, so the accesses have to be partitioned + * between them. + * + * `subtiles` accumulates the distinct subtiles found so far. Returns the index + * within it of the one `index` refers to, appending a new one if this is the + * first access to it. Every subtile must have the same shape, and any two must + * be either identical or disjoint - a partial overlap can't be expressed as + * separate tile registers. Anything else is a user error, reported in terms of + * `description`, which should name the kind of memory (e.g. "AMX tile"). */ +int get_subtile(const Expr &index, const std::string &description, + std::vector *subtiles); + /** Check if a vector Expr is a multiramp, and assign to result if so. * Returns false and leaves *result untouched if not. */ bool is_multiramp(const Expr &e, const Scope &scope, MultiRamp *result); +/** Check if an Expr is a load of a multiramp, possibly wrapped in a cast, in + * broadcasts over dimensions the load doesn't depend on, and in the lane + * permutations the simplifier introduces when it rewrites a strided load as a + * dense load followed by a transpose. Returns the Load node underneath, with + * *result describing the addresses it loads in the lane order of `e`, or null + * if `e` isn't of that form (in which case *result is untouched). + * + * At most one cast is seen through, so `e.type().element_of()` and the + * returned Load's type together tell you whether the loaded values were cast, + * and to what. Callers that care about the type the values actually have - + * including its signedness - must use the former, not the Load's type. */ +const Load *is_load_of_multiramp(const Expr &e, const Scope &scope, MultiRamp *result); + } // namespace Internal } // namespace Halide diff --git a/src/OffloadGPULoops.cpp b/src/OffloadGPULoops.cpp index 6551c10068d2..e67f6fad3f4c 100644 --- a/src/OffloadGPULoops.cpp +++ b/src/OffloadGPULoops.cpp @@ -79,7 +79,7 @@ class ExtractBounds : public IRVisitor { user_assert(!allocate->new_expr.defined()) << "Allocate node inside GPU kernel has custom new expression.\n" << "(Memoization is not supported inside GPU kernels at present.)\n"; - if (allocate->memory_type == MemoryType::GPUShared) { + if (is_gpu_shared(allocate->memory_type)) { internal_assert(allocate->extents.size() == 1); shared_mem_size += allocate->extents[0] * allocate->type.bytes(); found_shared = true; diff --git a/src/Serialization.cpp b/src/Serialization.cpp index 2ca84e441c77..9c050e4a95b3 100644 --- a/src/Serialization.cpp +++ b/src/Serialization.cpp @@ -150,14 +150,16 @@ Serialize::MemoryType Serializer::serialize_memory_type(const MemoryType &memory return Serialize::MemoryType::Register; case MemoryType::GPUShared: return Serialize::MemoryType::GPUShared; + case MemoryType::GPUSharedAsync: + return Serialize::MemoryType::GPUSharedAsync; case MemoryType::GPUTexture: return Serialize::MemoryType::GPUTexture; case MemoryType::LockedCache: return Serialize::MemoryType::LockedCache; case MemoryType::VTCM: return Serialize::MemoryType::VTCM; - case MemoryType::AMXTile: - return Serialize::MemoryType::AMXTile; + case MemoryType::Tile: + return Serialize::MemoryType::Tile; default: user_error << "Unsupported memory type\n"; return Serialize::MemoryType::Auto; diff --git a/src/Simplify_Call.cpp b/src/Simplify_Call.cpp index d63caa2e6c39..97951f1d8148 100644 --- a/src/Simplify_Call.cpp +++ b/src/Simplify_Call.cpp @@ -59,6 +59,38 @@ Expr Simplify::visit(const Call *op, ExprInfo *info) { info->cast_to(op->type); } + // A pure elementwise call of broadcasts is a broadcast of the call. Lifting + // the broadcast out lets the surrounding expression see that the value is + // the same in every lane. The clauses below that special-case particular + // functions have already mutated their arguments, so they can use this + // directly; anything that falls through to the general case at the bottom + // gets the same treatment there. + auto lift_broadcast_out = [&](const std::vector &args) -> Expr { + if (!op->type.is_vector() || !op->is_pure()) { + return Expr(); + } + const int lanes = op->type.lanes(); + std::vector scalar_args; + scalar_args.reserve(args.size()); + for (const Expr &arg : args) { + if (arg.type().is_scalar()) { + scalar_args.push_back(arg); + } else if (const Broadcast *b = arg.as(); + b && b->lanes == lanes && b->value.type().is_scalar()) { + scalar_args.push_back(b->value); + } else { + return Expr(); + } + } + if (scalar_args.empty()) { + return Expr(); + } + Expr scalar = Call::make(op->type.element_of(), op->name, scalar_args, + op->call_type, op->func, op->value_index, + op->image, op->param); + return mutate(Broadcast::make(std::move(scalar), lanes), info); + }; + if (op->is_intrinsic(Call::unreachable)) { in_unreachable = true; return op; @@ -778,6 +810,8 @@ Expr Simplify::visit(const Call *op, ExprInfo *info) { if (auto f = as_const_float(arg)) { auto fn = it->second; return make_const(arg.type(), fn(*f), info); + } else if (Expr e = lift_broadcast_out({arg}); e.defined()) { + return e; } else if (arg.same_as(op->args[0])) { return op; } else { @@ -810,6 +844,8 @@ Expr Simplify::visit(const Call *op, ExprInfo *info) { if (auto f = as_const_float(arg)) { auto fn = it->second; return make_const(arg.type(), fn(*f), info); + } else if (Expr e = lift_broadcast_out({arg}); e.defined()) { + return e; } else if (call && (call->call_type == Call::PureExtern || call->call_type == Call::PureIntrinsic) && (it = pure_externs_truncation.find(call->name)) != pure_externs_truncation.end()) { // For any combination of these integer-valued functions, we can @@ -863,6 +899,11 @@ Expr Simplify::visit(const Call *op, ExprInfo *info) { // No else: we want to fall thru from the PureExtern clause. { auto [new_args, changed] = mutate_with_changes(op->args); + + if (Expr e = lift_broadcast_out(new_args); e.defined()) { + return e; + } + if (!changed) { return op; } else { diff --git a/src/Simplify_Stmts.cpp b/src/Simplify_Stmts.cpp index 381ebb5a6177..c656f8e714ba 100644 --- a/src/Simplify_Stmts.cpp +++ b/src/Simplify_Stmts.cpp @@ -320,6 +320,7 @@ Stmt Simplify::visit(const Provide *op) { } Stmt Simplify::visit(const Store *op) { + found_buffer_reference(op->name); Expr predicate = mutate(op->predicate, nullptr); diff --git a/src/Util.h b/src/Util.h index 01d89eb3b6ff..6c04daed65f4 100644 --- a/src/Util.h +++ b/src/Util.h @@ -605,8 +605,10 @@ inline bool is_power_of_two(int64_t x) { return (x & (x - 1)) == 0; } +/** Round x up to the next multiple of n. Works for integral types and for + * Expr. */ template -inline T align_up(T x, int n) { +inline T align_up(const T &x, int n) { return (x + n - 1) / n * n; } diff --git a/src/halide_ir.fbs b/src/halide_ir.fbs index 60fdafdf6b5d..4715678f9c1f 100644 --- a/src/halide_ir.fbs +++ b/src/halide_ir.fbs @@ -116,7 +116,8 @@ enum MemoryType: byte { GPUTexture, LockedCache, VTCM, - AMXTile, + Tile, + GPUSharedAsync, } table Range { diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index ce92c73f4516..126d4a22b1d4 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -5,6 +5,7 @@ tests( SOURCES # keep-sorted start case=no align_bounds.cpp + amx_matmul.cpp argmax.cpp arm_cpu_detect.cpp async_device_copy.cpp @@ -137,6 +138,8 @@ tests( gpu_allocation_cache.cpp gpu_arg_types.cpp gpu_assertion_in_kernel.cpp + gpu_async_copy.cpp + gpu_async_copy_errors.cpp gpu_bounds_inference_failure.cpp gpu_condition_lifting.cpp gpu_cpu_simultaneous_read.cpp @@ -334,7 +337,6 @@ tests( sve_codegen_reinterpret.cpp target.cpp target_query.cpp - tiled_matmul.cpp tiled_matmul_errors.cpp tracing.cpp tracing_bounds.cpp @@ -379,6 +381,7 @@ tests( vectorized_load_from_vectorized_allocation.cpp vectorized_reduction_bug.cpp widening_reduction.cpp + wmma_matmul.cpp x86_cpu_detect.cpp # keep-sorted end ) diff --git a/test/correctness/tiled_matmul.cpp b/test/correctness/amx_matmul.cpp similarity index 91% rename from test/correctness/tiled_matmul.cpp rename to test/correctness/amx_matmul.cpp index ea44064902db..0161968dda78 100644 --- a/test/correctness/tiled_matmul.cpp +++ b/test/correctness/amx_matmul.cpp @@ -113,40 +113,37 @@ bool matmul(int col, int row, int acc, int tile_x, int tile_y, int tile_r, bool // loads. But if you go too big you'll run out of tile registers and // compilation will fail (The LLVM AMX register allocator will spill, but it // seems to be fussy about it). Doing this also tests the case of more than - // one matrix multiply operation applied to a single AMXTile allocation. + // one matrix multiply operation applied to a single tile allocation. int outer_tile_x = col > tile_x ? 2 : 1, outer_tile_y = row > tile_y ? 2 : 1; mm.compute_at(mm.in(), x) - .store_in(MemoryType::AMXTile) .update() .tile(x, y, rxi, ryi, tile_x, tile_y, TailStrategy::GuardWithIf) .split(r, rro, rri, tile_r) - .reorder(rri, rxi, ryi, rro, x, y) - .atomic() - .vectorize(rri) - .vectorize(rxi) - .vectorize(ryi) + .reorder(rro, x, y) + .tile_matmul(rri, rxi, ryi) .tile(x, y, xi, yi, outer_tile_x, outer_tile_y) - .reorder(rri, rxi, ryi, xi, yi, rro, x, y) + .reorder(xi, yi, rro, x, y) .unroll(xi) .unroll(yi); Var ixi("ixi"), iyi("iyi"); mm.compute_at(mm.in(), x) .tile(x, y, ixi, iyi, tile_x, tile_y) - .vectorize(ixi) - .vectorize(iyi) + .tile_init(ixi, iyi) .unroll(x) .unroll(y); - // schedule the consumer - Var mmxi("mmxi"), mmyi("mmyi"); + // Schedule the consumer, which is what copies each tile back out. The + // outer tiling matches the group of tiles the accumulator holds, and the + // inner one is a single tile. + Var mmxo("mmxo"), mmyo("mmyo"), mmxi("mmxi"), mmyi("mmyi"); mm.in() - .tile(x, y, mmxi, mmyi, tile_x * outer_tile_x, tile_y * outer_tile_y) - .vectorize(mmxi, tile_x) - .vectorize(mmyi, tile_y) - .unroll(mmxi) - .unroll(mmyi); + .tile(x, y, mmxo, mmyo, tile_x * outer_tile_x, tile_y * outer_tile_y) + .tile(mmxo, mmyo, mmxi, mmyi, tile_x, tile_y) + .tile_store(mmxi, mmyi) + .unroll(mmxo) + .unroll(mmyo); Func result = mm.in(); @@ -160,7 +157,7 @@ bool matmul(int col, int row, int acc, int tile_x, int tile_y, int tile_r, bool result.realize(out); } else { // Just compile it to see if anything crashes - result.compile_to_assembly(Internal::get_test_tmp_dir() + "tiled_matmul.s", + result.compile_to_assembly(Internal::get_test_tmp_dir() + "amx_matmul.s", {A_buf, B_buf}, Target{"x86-64-linux-avx512_sapphirerapids-no_asserts-no_runtime-no_bounds_query"}); return true; } @@ -211,28 +208,22 @@ bool matmul_bf16(int col, int row, int acc, int tile_x, int tile_y, int tile_r, RVar rri("rri"), rro("rro"); mm.compute_at(mm.in(), x) - .store_in(MemoryType::AMXTile) .update() .tile(x, y, rxi, ryi, tile_x, tile_y, TailStrategy::GuardWithIf) .split(r.x, rro, rri, tile_r) .reorder({rri, rxi, ryi, rro, x, y}) - .atomic() - .vectorize(rri) - .vectorize(rxi) - .vectorize(ryi); + .tile_matmul(rri, rxi, ryi); Var ixi("ixi"), iyi("iyi"); mm.compute_at(mm.in(), x) .tile(x, y, ixi, iyi, tile_x, tile_y) - .vectorize(ixi) - .vectorize(iyi); + .tile_init(ixi, iyi); // schedule the consumer Var mmxi("mmxi"), mmyi("mmyi"); mm.in() .tile(x, y, mmxi, mmyi, tile_x, tile_y) - .vectorize(mmxi) - .vectorize(mmyi); + .tile_store(mmxi, mmyi); Func result = mm.in(); @@ -242,15 +233,15 @@ bool matmul_bf16(int col, int row, int acc, int tile_x, int tile_y, int tile_r, Buffer out(col, row); // Uncomment to check the asm - // result.compile_to_llvm_assembly(Internal::get_test_tmp_dir() + "tiled_matmul_bf16.ll", {A, B}, target); - // result.compile_to_assembly(Internal::get_test_tmp_dir() + "tiled_matmul.s", {A, B}, target); + // result.compile_to_llvm_assembly(Internal::get_test_tmp_dir() + "amx_matmul_bf16.ll", {A, B}, target); + // result.compile_to_assembly(Internal::get_test_tmp_dir() + "amx_matmul.s", {A, B}, target); Target target = get_jit_target_from_environment(); if (target.has_feature(Target::AVX512_SapphireRapids)) { result.realize(out); } else { // Just compile it to see if anything crashes - result.compile_to_assembly(Internal::get_test_tmp_dir() + "tiled_matmul.s", {A, B}, Target{"x86-64-linux-avx512_sapphirerapids"}); + result.compile_to_assembly(Internal::get_test_tmp_dir() + "amx_matmul.s", {A, B}, Target{"x86-64-linux-avx512_sapphirerapids"}); return true; } diff --git a/test/correctness/gpu_async_copy.cpp b/test/correctness/gpu_async_copy.cpp new file mode 100644 index 000000000000..80bd32a2f20e --- /dev/null +++ b/test/correctness/gpu_async_copy.cpp @@ -0,0 +1,287 @@ +#include "Halide.h" +#include + +using namespace Halide; + +// A Func stored in GPUSharedAsync is staged into shared memory by the copy +// engine rather than by loading it into registers and storing it back out. +// The copy engine moves 4, 8 or 16 bytes per thread, so the cases below cover +// each of those widths at several element sizes, as well as the shapes a +// staged input tends to take: a two-dimensional tile, more than one input +// staged into the same kernel, a kernel that mixes the two staging modes, and +// a wrapper made with Func::in. + +namespace { + +int failures = 0; + +template +void check_result(const char *name, Buffer &result, Buffer &input) { + for (int y = 0; y < result.height(); y++) { + for (int x = 0; x < result.width(); x++) { + T correct = (T)(input(x, y) * 2); + if (result(x, y) != correct) { + printf("[%s] result(%d, %d) = %f instead of %f\n", + name, x, y, (double)result(x, y), (double)correct); + failures++; + return; + } + } + } + printf("[%s] OK\n", name); +} + +template +Buffer make_input(int w, int h) { + Buffer in(w, h); + in.fill([](int x, int y) { return (T)((x + y * 3) % 32); }); + return in; +} + +// Stage an input through shared memory with the copy engine, vectorized by +// `vec` elements, and check the consumer sees the right values. +template +void test_width(const char *name, int vec) { + const int W = 256, H = 32; + Buffer input = make_input(W, H); + + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); + stage(x, y) = input(x, y); + out(x, y) = cast(stage(x, y) * 2); + + out.gpu_tile(x, y, xi, yi, 64, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .gpu_threads(y) + .vectorize(x, vec); + + Buffer result(W, H); + out.realize(result); + result.copy_to_host(); + check_result(name, result, input); +} + +// The staged tile spread over threads in both dimensions. +void test_2d_tile() { + const int W = 256, H = 64; + Buffer input = make_input(W, H); + + Var x("x"), y("y"), xi("xi"), yi("yi"), xii("xii"); + Func stage("stage"), out("out"); + stage(x, y) = input(x, y); + out(x, y) = stage(x, y) * 2; + + out.gpu_tile(x, y, xi, yi, 64, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .split(x, x, xii, 4) + .vectorize(xii) + .gpu_threads(x, y); + + Buffer result(W, H); + out.realize(result); + result.copy_to_host(); + check_result("2d_tile", result, input); +} + +// Padding the rows to dodge bank conflicts, with a stride that stays a +// multiple of the vector width so the rows are still aligned. +void test_padded_storage() { + const int W = 256, H = 32; + Buffer input = make_input(W, H); + + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); + stage(x, y) = input(x, y); + out(x, y) = stage(x, y) * 2; + + out.gpu_tile(x, y, xi, yi, 64, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .align_storage(x, 8) + .gpu_threads(y) + .vectorize(x, 4); + + Buffer result(W, H); + out.realize(result); + result.copy_to_host(); + check_result("padded_storage", result, input); +} + +// Two inputs staged into the same kernel, so more than one copy is in flight +// before the barrier that waits for them. +void test_two_inputs() { + const int W = 256, H = 32; + Buffer a = make_input(W, H); + Buffer b = make_input(W, H); + + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func sa("sa"), sb("sb"), out("out"); + sa(x, y) = a(x, y); + sb(x, y) = b(x, y); + out(x, y) = sa(x, y) + sb(x, y); + + out.gpu_tile(x, y, xi, yi, 64, 8); + sa.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .gpu_threads(y) + .vectorize(x, 4); + sb.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .gpu_threads(y) + .vectorize(x, 4); + + Buffer result(W, H); + out.realize(result); + result.copy_to_host(); + check_result("two_inputs", result, a); +} + +// One input staged by the copy engine and one staged the ordinary way, in +// either order. The two get separate shared allocations, which have to be +// placed at different addresses within the block's shared memory. +void test_mixed(const char *name, MemoryType first, MemoryType second) { + const int W = 256, H = 32; + Buffer a = make_input(W, H); + Buffer b = make_input(W, H); + + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func sa("sa"), sb("sb"), out("out"); + sa(x, y) = a(x, y); + sb(x, y) = b(x, y); + out(x, y) = sa(x, y) + sb(x, y); + + out.gpu_tile(x, y, xi, yi, 64, 8); + sa.compute_at(out, x) + .store_in(first) + .gpu_threads(y) + .vectorize(x, 4); + sb.compute_at(out, x) + .store_in(second) + .gpu_threads(y) + .vectorize(x, 4); + + Buffer result(W, H); + out.realize(result); + result.copy_to_host(); + check_result(name, result, a); +} + +// Func::in is the idiomatic way to get a Func that is a plain copy, and is +// what the error message points users at. +void test_wrapper() { + const int W = 256, H = 32; + Buffer input = make_input(W, H); + + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func in_f("in_f"), out("out"); + in_f(x, y) = input(x, y); + out(x, y) = in_f(x, y) * 2; + + out.gpu_tile(x, y, xi, yi, 64, 8); + in_f.in() + .compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .gpu_threads(y) + .vectorize(x, 4); + in_f.compute_root(); + + Buffer result(W, H); + out.realize(result); + result.copy_to_host(); + check_result("wrapper", result, input); +} + +// The same staging in plain GPUShared must stay synchronous, so that a +// schedule can ask for either one and compare them. +void test_opt_out() { + const int W = 256, H = 32; + Buffer input = make_input(W, H); + + auto build = [&](MemoryType mt) { + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); + stage(x, y) = input(x, y); + out(x, y) = stage(x, y) * 2; + out.gpu_tile(x, y, xi, yi, 64, 8); + stage.compute_at(out, x) + .store_in(mt) + .gpu_threads(y) + .vectorize(x, 4); + Buffer result(W, H); + out.realize(result); + result.copy_to_host(); + return result; + }; + + Buffer shared = build(MemoryType::GPUShared); + check_result("opt_out_shared", shared, input); + Buffer async = build(MemoryType::GPUSharedAsync); + check_result("opt_out_async", async, input); +} + +// GPU APIs with no copy engine have to treat this memory type as ordinary +// shared memory rather than failing to compile. Only compiling is needed, so +// this runs whether or not such a device is present. +void test_other_gpu_apis() { + const int W = 256, H = 32; + Buffer input = make_input(W, H); + + for (Target::Feature api : {Target::OpenCL, Target::Metal}) { + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); + stage(x, y) = input(x, y); + out(x, y) = stage(x, y) * 2; + out.gpu_tile(x, y, xi, yi, 64, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .gpu_threads(y) + .vectorize(x, 4); + out.compile_to_module({}, "f", get_host_target().with_feature(api)); + } + printf("[other_gpu_apis] OK\n"); +} + +} // namespace + +int main(int argc, char **argv) { + test_other_gpu_apis(); + + Target target = get_jit_target_from_environment(); + if (!target.has_feature(Target::CUDA)) { + printf("[SKIP] No CUDA target enabled.\n"); + return 0; + } + if (target.get_cuda_capability_lower_bound() < 80) { + printf("[SKIP] Asynchronous copies need compute capability 8.0 or above.\n"); + return 0; + } + + // Each of the three copy widths the hardware supports, at several element + // sizes. + test_width("f32_x1_4b", 1); + test_width("f32_x2_8b", 2); + test_width("f32_x4_16b", 4); + test_width("u8_x4_4b", 4); + test_width("u8_x8_8b", 8); + test_width("u8_x16_16b", 16); + test_width("u16_x2_4b", 2); + test_width("u16_x8_16b", 8); + test_width("i32_x4_16b", 4); + + test_2d_tile(); + test_padded_storage(); + test_two_inputs(); + test_mixed("mixed_sync_first", MemoryType::GPUShared, MemoryType::GPUSharedAsync); + test_mixed("mixed_async_first", MemoryType::GPUSharedAsync, MemoryType::GPUShared); + test_wrapper(); + test_opt_out(); + + if (failures) { + printf("%d case(s) failed\n", failures); + return 1; + } + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/gpu_async_copy_errors.cpp b/test/correctness/gpu_async_copy_errors.cpp new file mode 100644 index 000000000000..cf533c16ec09 --- /dev/null +++ b/test/correctness/gpu_async_copy_errors.cpp @@ -0,0 +1,223 @@ +// Exercises the user-facing error paths for MemoryType::GPUSharedAsync. Asking +// for that memory type is a promise that every store to the allocation is a +// copy the copy engine can make, so each scenario below breaks one of those +// requirements and should be told exactly which one. +// +// The test verifies that each scenario produces a Halide::CompileError (a user +// error) rather than crashing, hitting an internal assert, or quietly falling +// back to a load and a store. + +#include "Halide.h" +#include + +#if HALIDE_WITH_EXCEPTIONS + +using namespace Halide; + +namespace { + +// Compiling is enough to reach the error - no device is needed. +Target async_target() { + return get_host_target() + .with_feature(Target::CUDA) + .with_feature(Target::CUDACapability80); +} + +// Run `body` and assert it produces a Halide user error mentioning `substring`. +template +bool expect_user_error(const char *name, const char *substring, F body) { + try { + body(); + } catch (const CompileError &e) { + std::string msg = e.what(); + if (msg.find(substring) == std::string::npos) { + printf("[%s] FAIL: error did not mention \"%s\":\n%s\n", + name, substring, msg.c_str()); + return false; + } + printf("[%s] OK\n", name); + return true; + } catch (...) { + printf("[%s] FAIL: expected a CompileError but got a different exception\n", name); + return false; + } + printf("[%s] FAIL: expected a user error but none was raised\n", name); + return false; +} + +Buffer input_f32() { + Buffer in(256, 64); + in.fill(0.f); + return in; +} + +// The copy engine moves bytes untouched, so the staged Func has to be a plain +// copy. Doing arithmetic in it means the value stored isn't a load at all. +void scenario_not_a_copy() { + Buffer in = input_f32(); + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); + stage(x, y) = in(x, y) * 2.f; + out(x, y) = stage(x, y); + out.gpu_tile(x, y, xi, yi, 64, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .gpu_threads(y) + .vectorize(x, 4); + out.compile_jit(async_target()); +} + +// The source has to live outside the kernel. Staging something already +// computed into registers or shared memory isn't a global-to-shared copy. +void scenario_source_inside_kernel() { + Buffer in = input_f32(); + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func producer("producer"), stage("stage"), out("out"); + producer(x, y) = in(x, y) + 1.f; + stage(x, y) = producer(x, y); + out(x, y) = stage(x, y); + out.gpu_tile(x, y, xi, yi, 64, 8); + producer.compute_at(out, x).gpu_threads(y); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .gpu_threads(y) + .vectorize(x, 4); + out.compile_jit(async_target()); +} + +// The hardware copies 4, 8 or 16 bytes per thread. Three floats is 12. +void scenario_bad_vector_width() { + Buffer in = input_f32(); + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); + stage(x, y) = in(x, y); + out(x, y) = stage(x, y); + out.gpu_tile(x, y, xi, yi, 48, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .gpu_threads(y) + .vectorize(x, 3); + out.compile_jit(async_target()); +} + +// A single byte per thread is below the minimum too. +void scenario_too_narrow() { + Buffer in(256, 64); + in.fill(0); + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); + stage(x, y) = in(x, y); + out(x, y) = stage(x, y); + out.gpu_tile(x, y, xi, yi, 64, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .gpu_threads(x, y); + out.compile_jit(async_target()); +} + +// Each copy has to be contiguous in both the source and the destination, so a +// strided read of the input can't be done this way. +void scenario_strided_source() { + Buffer in(512, 64); + in.fill(0.f); + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); + stage(x, y) = in(2 * x, y); + out(x, y) = stage(x, y); + out.gpu_tile(x, y, xi, yi, 64, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .gpu_threads(y) + .vectorize(x, 4); + out.compile_jit(async_target()); +} + +// A predicated tail means some lanes are masked off, which the copy engine +// can't express. +void scenario_predicated() { + Buffer in(256, 64); + in.fill(0.f); + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); + stage(x, y) = in(x, y); + out(x, y) = stage(x, y); + out.gpu_tile(x, y, xi, yi, 66, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .align_storage(x, 4) + .gpu_threads(y) + .vectorize(x, 4, TailStrategy::Predicate); + out.compile_jit(async_target()); +} + +// Each copy has to be dense at both ends. Storing the staged Func in the +// opposite order to the one it is read in leaves the source dense but the +// destination strided. +void scenario_strided_destination() { + Buffer in = input_f32(); + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); + stage(x, y) = in(x, y); + out(x, y) = stage(x, y); + out.gpu_tile(x, y, xi, yi, 64, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .reorder_storage(y, x) + .gpu_threads(y) + .vectorize(x, 4); + out.compile_jit(async_target()); +} + +// The destination address of each copy has to be aligned to its width. Padding +// the rows to a stride that isn't a multiple of the vector width breaks that +// for every row after the first. +void scenario_misaligned_destination() { + Buffer in = input_f32(); + Var x("x"), y("y"), xi("xi"), yi("yi"); + Func stage("stage"), out("out"); + stage(x, y) = in(x, y); + out(x, y) = stage(x, y); + out.gpu_tile(x, y, xi, yi, 64, 8); + stage.compute_at(out, x) + .store_in(MemoryType::GPUSharedAsync) + .align_storage(x, 6) + .gpu_threads(y) + .vectorize(x, 4); + out.compile_jit(async_target()); +} + +} // namespace + +int main(int argc, char **argv) { + if (!Halide::exceptions_enabled()) { + printf("[SKIP] Halide was compiled without exceptions.\n"); + return 0; + } + + int failures = 0; + + failures += !expect_user_error("not_a_copy", "not a load", scenario_not_a_copy); + failures += !expect_user_error("source_inside_kernel", "another allocation inside the", scenario_source_inside_kernel); + failures += !expect_user_error("bad_vector_width", "4, 8 or 16 bytes", scenario_bad_vector_width); + failures += !expect_user_error("too_narrow", "4, 8 or 16 bytes", scenario_too_narrow); + failures += !expect_user_error("strided_source", "not read densely", scenario_strided_source); + failures += !expect_user_error("predicated", "predicated", scenario_predicated); + failures += !expect_user_error("strided_destination", "densely", scenario_strided_destination); + failures += !expect_user_error("misaligned_destination", "aligned", scenario_misaligned_destination); + + if (failures != 0) { + printf("%d scenario(s) did not produce the expected user error\n", failures); + return 1; + } + printf("Success!\n"); + return 0; +} + +#else // HALIDE_WITH_EXCEPTIONS + +int main(int argc, char **argv) { + printf("[SKIP] Halide was compiled without exceptions.\n"); + return 0; +} + +#endif // HALIDE_WITH_EXCEPTIONS diff --git a/test/correctness/multiramp.cpp b/test/correctness/multiramp.cpp index c7b8aae195db..da7a37c8cc08 100644 --- a/test/correctness/multiramp.cpp +++ b/test/correctness/multiramp.cpp @@ -1,6 +1,7 @@ #include "Halide.h" #include +#include #include #include @@ -583,6 +584,169 @@ void check_reject_non_multiramp_sum() { CHECK(!is_multiramp(sum, scope, &m), "reject coprime-shape add"); } +// ---- Shuffles ------------------------------------------------------------ + +// Mirror Shuffle::make_transpose: view v as a row-major matrix with `cols` +// columns and transpose it. +std::vector transpose_vec(const std::vector &v, int cols) { + int rows = (int)v.size() / cols; + std::vector result(v.size()); + for (int j = 0; j < cols; j++) { + for (int i = 0; i < rows; i++) { + result[j * rows + i] = v[i * cols + j]; + } + } + return result; +} + +// A mask whose strides are the prefix products of its lane counts, in some +// order, is a permutation of the lanes. Build one from a dim order. +MultiRamp permutation_mask(const std::vector &shape, const std::vector &order, + Type t) { + std::vector flat(shape.size()); + int64_t s = 1; + for (size_t i = 0; i < shape.size(); i++) { + flat[i] = s; + s *= shape[i]; + } + std::vector strides; + std::vector lanes; + for (int d : order) { + strides.push_back(make_const(t, flat[d])); + lanes.push_back(shape[d]); + } + return MultiRamp(make_zero(t), strides, lanes); +} + +// Apply a mask the slow way, for comparison. +std::vector apply_mask(const std::vector &v, const MultiRamp &mask) { + std::vector result; + for (int i : expand(mask)) { + result.push_back(v[i]); + } + return result; +} + +void check_shuffle_case(const MultiRamp &m, const std::vector &shape, + const std::vector &order, const char *msg) { + MultiRamp a = m; + MultiRamp mask = permutation_mask(shape, order, Int(32)); + auto want = apply_mask(expand(a), mask); + if (!a.shuffle(mask)) { + printf("FAIL: %s: shuffle returned false\n", msg); + failures++; + return; + } + check_seq(expand(a), want, msg, __LINE__); +} + +void check_shuffles() { + // The two-dimensional transpose the wmma pass does to an operand tile. + check_shuffle_case(MultiRamp{0, {1, 64}, {16, 16}}, {16, 16}, {1, 0}, + "transpose of a 2D tile"); + // A transpose that falls inside a dim, so it has to be split. + check_shuffle_case(MultiRamp{0, {1}, {8}}, {4, 2}, {1, 0}, + "transpose splitting a dim"); + // A permutation whose dims collapse to something coarser than the + // multiramp's, so both shapes need refining before they line up. + check_shuffle_case(MultiRamp{0, {1, 10, 100}, {2, 2, 3}}, {2, 2, 3}, {2, 0, 1}, + "permutation needing a common refinement"); + // A broadcast dim has stride zero, which never merges with its neighbour. + check_shuffle_case(MultiRamp{0, {1, 64, 0}, {16, 16, 16}}, {16, 256}, {1, 0}, + "transpose of a broadcast tile"); +} + +void check_shuffle_rejects_gather() { + // Strides that aren't the prefix products of the lane counts don't + // describe a permutation. + MultiRamp A{0, {1, 100}, {4, 4}}; + MultiRamp mask{0, {1, 2}, {4, 4}}; + CHECK(!A.shuffle(mask), "shuffle rejects a non-permutation mask"); +} + +void check_shuffle_rejects_coprime_shapes() { + // Lane counts of 3 and 2 innermost have no common refinement. + MultiRamp A{0, {1, 100}, {2, 3}}; + MultiRamp mask = permutation_mask({3, 2}, {1, 0}, Int(32)); + CHECK(!A.shuffle(mask), "shuffle rejects coprime shapes"); +} + +void check_shuffle_rejects_wrong_size() { + MultiRamp A{0, {1, 100}, {4, 4}}; + MultiRamp mask = permutation_mask({2, 2}, {1, 0}, Int(32)); + CHECK(!A.shuffle(mask), "shuffle rejects a mask of the wrong size"); +} + +void check_recognize_transpose_shuffle() { + Expr e = Shuffle::make_transpose(Ramp::make(Expr(0), Expr(1), 12), 4); + Scope scope; + MultiRamp m; + CHECK(is_multiramp(e, scope, &m), "recognize a transpose shuffle"); + std::vector in(12); + std::iota(in.begin(), in.end(), 0); + CHECK_SEQ(expand(m), transpose_vec(in, 4), "transpose shuffle values"); +} + +void check_recognize_reshaping_shuffle() { + // A shuffle of a 1D ramp whose lane indices are themselves a multiramp is + // a reshaping of the ramp rather than an arbitrary gather. Indices + // [0,2,4,6,1,3,5,7] have shape (4,2) and strides (2,1). + Expr r = Ramp::make(Expr(0), Expr(3), 8); + Expr e = Shuffle::make({r}, {0, 2, 4, 6, 1, 3, 5, 7}); + Scope scope; + MultiRamp m; + CHECK(is_multiramp(e, scope, &m), "recognize a reshaping shuffle"); + CHECK_SEQ_LIT(expand(m), "reshaping shuffle values", + 0, 6, 12, 18, 3, 9, 15, 21); +} + +void check_reject_gather_shuffle() { + // Lane indices that aren't a multiramp really are a gather. + Expr r = Ramp::make(Expr(0), Expr(3), 8); + Expr e = Shuffle::make({r}, {0, 3, 1, 7, 2, 5, 4, 6}); + Scope scope; + MultiRamp m; + CHECK(!is_multiramp(e, scope, &m), "reject a gather shuffle"); +} + +// ---- is_load_of_multiramp ------------------------------------------------ + +Expr make_test_load(int lanes) { + return Load::make(Int(16, lanes), "buf", Ramp::make(Expr(0), Expr(1), lanes)); +} + +void check_load_under_cast_and_broadcast() { + // A broadcast of a load is a load at a stride-zero outer dim. + Expr e = Broadcast::make(Cast::make(Int(32, 4), make_test_load(4)), 3); + Scope scope; + MultiRamp m; + const Load *load = is_load_of_multiramp(e, scope, &m); + CHECK(load && load->name == "buf", "see through cast and broadcast"); + if (load) { + CHECK_SEQ_LIT(expand(m), "broadcast load addresses", + 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3); + } +} + +void check_load_under_shuffle() { + // Permuting the loaded values permutes the addresses loaded from. + Expr e = Shuffle::make_transpose(make_test_load(4), 2); + Scope scope; + MultiRamp m; + const Load *load = is_load_of_multiramp(e, scope, &m); + CHECK(load && load->name == "buf", "see through a lane permutation"); + if (load) { + CHECK_SEQ_LIT(expand(m), "permuted load addresses", 0, 2, 1, 3); + } +} + +void check_reject_non_load() { + Scope scope; + MultiRamp m; + CHECK(!is_load_of_multiramp(Ramp::make(Expr(0), Expr(1), 4), scope, &m), + "reject an Expr with no load underneath"); +} + } // namespace int main(int argc, char **argv) { @@ -631,6 +795,19 @@ int main(int argc, char **argv) { check_roundtrips(); check_reject_non_multiramp_sum(); + check_shuffles(); + check_shuffle_rejects_gather(); + check_shuffle_rejects_coprime_shapes(); + check_shuffle_rejects_wrong_size(); + + check_recognize_transpose_shuffle(); + check_recognize_reshaping_shuffle(); + check_reject_gather_shuffle(); + + check_load_under_cast_and_broadcast(); + check_load_under_shuffle(); + check_reject_non_load(); + if (failures) { printf("%d failures\n", failures); return 1; diff --git a/test/correctness/tiled_matmul_errors.cpp b/test/correctness/tiled_matmul_errors.cpp index fd2e85a3f041..9714e2ffe422 100644 --- a/test/correctness/tiled_matmul_errors.cpp +++ b/test/correctness/tiled_matmul_errors.cpp @@ -19,10 +19,16 @@ const Target amx_target("x86-64-linux-avx512_sapphirerapids"); // Run `body` and assert it produces a Halide user error. template -bool expect_user_error(const char *name, F body) { +bool expect_user_error(const char *name, const char *substring, F body) { try { body(); } catch (const CompileError &e) { + std::string msg = e.what(); + if (msg.find(substring) == std::string::npos) { + printf("[%s] FAIL: error did not mention \"%s\":\n%s\n", + name, substring, msg.c_str()); + return false; + } printf("[%s] OK: %s\n", name, e.what()); return true; } catch (...) { @@ -39,27 +45,21 @@ void schedule_matmul(Func mm, RVar r, int tile_x, int tile_y, int tile_r) { Var x("x"), y("y"), rxi("rxi"), ryi("ryi"); RVar rri("rri"), rro("rro"); mm.compute_at(mm.in(), x) - .store_in(MemoryType::AMXTile) .update() .tile(x, y, rxi, ryi, tile_x, tile_y, TailStrategy::GuardWithIf) .split(r, rro, rri, tile_r) - .reorder(rri, rxi, ryi, rro, x, y) - .atomic() - .vectorize(rri) - .vectorize(rxi) - .vectorize(ryi); + .reorder(rro, x, y) + .tile_matmul(rri, rxi, ryi); Var ixi("ixi"), iyi("iyi"); mm.compute_at(mm.in(), x) .tile(x, y, ixi, iyi, tile_x, tile_y) - .vectorize(ixi) - .vectorize(iyi); + .tile_init(ixi, iyi); Var mmxi("mmxi"), mmyi("mmyi"); mm.in() .tile(x, y, mmxi, mmyi, tile_x, tile_y) - .vectorize(mmxi) - .vectorize(mmyi); + .tile_store(mmxi, mmyi); } // A tile too large for an AMX register (rows > 16). Triggers the explicit @@ -77,7 +77,7 @@ void scenario_too_large() { mm.in().compile_jit(amx_target); } -// AMXTile allocated for a non-i32/f32 result. AMX always accumulates into +// A tile allocated for a non-i32/f32 result. AMX always accumulates into // 32-bit registers, so we reject. Triggers the user_assert in // visit(Allocate). void scenario_bad_result_type() { @@ -113,7 +113,7 @@ void scenario_naive_rhs() { // A gather-style matmul with an indirect row index — natural for sparse / // pruned matmul, indirect attention, embedding lookups. The LHS load index // goes through a table lookup, so the multiramp lift fails. Triggers the -// "loads indices are not affine" path in convert_to_matmul. +// "not loads with affine indices" path in convert_to_matmul. void scenario_indirect() { Buffer A(64, 64); Buffer B(4, 64, 16); @@ -130,6 +130,26 @@ void scenario_indirect() { mm.in().compile_jit(amx_target); } +// A matmul over uint8 buffers whose values are reinterpreted as int8 before +// multiplying. AMX types its tile registers by what was loaded, and their +// signedness picks which of the four integer tdpb instructions runs, so +// accepting this would run an unsigned multiply for a signed algorithm. +// Triggers the "cast after being loaded" path in convert_to_matmul. +void scenario_sign_changing_cast() { + Buffer A(64, 64); + Buffer B(4, 64, 16); + Var x("x"), y("y"); + RDom r(0, 64, "r"); + + Func mm("matmul_sign_cast"); + mm(x, y) = cast(0); + mm(x, y) += + cast(cast(A(r, y))) * + cast(cast(B(r % 4, x, r / 4))); + schedule_matmul(mm, r.x, 8, 8, 8); + mm.in().compile_jit(amx_target); +} + // A 1D convolution of a 2D signal with per-row kernels, aggressively // vectorized. Structurally a sum-of-widening-multiplies with a contiguous // inner K, but the LHS depends on x, k, and y simultaneously (no broadcast @@ -151,7 +171,7 @@ void scenario_conv1d() { conv.in().compile_jit(amx_target); } -// A non-matmul value scheduled into an AMXTile allocation by mistake. +// A non-matmul value scheduled into a tile allocation by mistake. // Triggers the "no matrix multiply was found" assertion. void scenario_no_matmul() { Var x("x"), y("y"), xo("xo"), yo("yo"), xi("xi"), yi("yi"); @@ -159,7 +179,7 @@ void scenario_no_matmul() { Func f("f"); f(x, y) = 0; f.compute_at(f.in(), xo) - .store_in(MemoryType::AMXTile) + .store_in(MemoryType::Tile) .vectorize(x, 8) .vectorize(y, 8); @@ -185,8 +205,40 @@ void scenario_widening_16bit() { mm.in().compile_jit(amx_target); } +// Two matmuls into the same allocation with tiles of the same rank but +// different extents, so they disagree about the shape of a tile register. +void scenario_mismatched_strides() { + Buffer A(64, 64), C(64, 64); + Buffer B(4, 64, 16), D(4, 64, 16); + Var x("x"), y("y"); + RDom r1(0, 64, "r1"), r2(0, 64, "r2"); + + Func mm("matmul_mismatched"); + mm(x, y) = cast(0); + mm(x, y) += cast(A(r1, y)) * cast(B(r1 % 4, x, r1 / 4)); + mm(x, y) += cast(C(r2, y)) * cast(D(r2 % 4, x, r2 / 4)); + + Var rxi("rxi"), ryi("ryi"); + RVar rri("rri"), rro("rro"); + mm.update(0) + .tile(x, y, rxi, ryi, 8, 4, TailStrategy::GuardWithIf) + .split(r1.x, rro, rri, 8) + .reorder(rro, x, y) + .tile_matmul(rri, rxi, ryi); + mm.update(1) + .tile(x, y, rxi, ryi, 4, 8, TailStrategy::GuardWithIf) + .split(r2.x, rro, rri, 8) + .reorder(rro, x, y) + .tile_matmul(rri, rxi, ryi); + Var ixi("ixi"), iyi("iyi"); + mm.compute_at(mm.in(), x).tile(x, y, ixi, iyi, 8, 8).tile_init(ixi, iyi); + Var mmxi("mmxi"), mmyi("mmyi"); + mm.in().tile(x, y, mmxi, mmyi, 8, 8).tile_store(mmxi, mmyi); + mm.in().compile_jit(amx_target); +} + // A user gives the same Func two update definitions that each store into -// the same AMXTile allocation but with different tile sizes (e.g. a fast +// the same tile allocation but with different tile sizes (e.g. a fast // path for the bulk of K and a smaller fallback). The matcher requires // every matmul touching a given allocation to agree on tile dimensions. void scenario_inconsistent_tiles() { @@ -203,40 +255,31 @@ void scenario_inconsistent_tiles() { Var rxi("rxi"), ryi("ryi"); RVar rri("rri"), rro("rro"); - mm.compute_at(mm.in(), x).store_in(MemoryType::AMXTile); mm.update(0) .tile(x, y, rxi, ryi, 8, 8, TailStrategy::GuardWithIf) .split(r1.x, rro, rri, 8) - .reorder(rri, rxi, ryi, rro, x, y) - .atomic() - .vectorize(rri) - .vectorize(rxi) - .vectorize(ryi); + .reorder(rro, x, y) + .tile_matmul(rri, rxi, ryi); mm.update(1) .tile(x, y, rxi, ryi, 4, 4, TailStrategy::GuardWithIf) .split(r2.x, rro, rri, 8) - .reorder(rri, rxi, ryi, rro, x, y) - .atomic() - .vectorize(rri) - .vectorize(rxi) - .vectorize(ryi); + .reorder(rro, x, y) + .tile_matmul(rri, rxi, ryi); Var ixi("ixi"), iyi("iyi"); mm.compute_at(mm.in(), x) .tile(x, y, ixi, iyi, 8, 8) - .vectorize(ixi) - .vectorize(iyi); + .tile_init(ixi, iyi); Var mmxi("mmxi"), mmyi("mmyi"); mm.in() .tile(x, y, mmxi, mmyi, 8, 8) - .vectorize(mmxi) - .vectorize(mmyi); + .tile_store(mmxi, mmyi); mm.in().compile_jit(amx_target); } -// A reduction inside an AMXTile that's a sum-of-something-else, not a +// A reduction inside a tile that's a sum-of-something-else, not a // vector_reduce_add of a widening multiply. Here we accumulate a // non-multiplied value, which produces a Store whose RHS is not a // vector-reduce-of-multiply. @@ -253,27 +296,21 @@ void scenario_not_a_matmul_pattern() { Var rxi("rxi"), ryi("ryi"); RVar rri("rri"), rro("rro"); mm.compute_at(mm.in(), x) - .store_in(MemoryType::AMXTile) .update() .tile(x, y, rxi, ryi, 8, 8, TailStrategy::GuardWithIf) .split(r.x, rro, rri, 8) - .reorder(rri, rxi, ryi, rro, x, y) - .atomic() - .vectorize(rri) - .vectorize(rxi) - .vectorize(ryi); + .reorder(rro, x, y) + .tile_matmul(rri, rxi, ryi); Var ixi("ixi"), iyi("iyi"); mm.compute_at(mm.in(), x) .tile(x, y, ixi, iyi, 8, 8) - .vectorize(ixi) - .vectorize(iyi); + .tile_init(ixi, iyi); Var mmxi("mmxi"), mmyi("mmyi"); mm.in() .tile(x, y, mmxi, mmyi, 8, 8) - .vectorize(mmxi) - .vectorize(mmyi); + .tile_store(mmxi, mmyi); mm.in().compile_jit(amx_target); } @@ -303,16 +340,18 @@ int main(int argc, char **argv) { int failures = 0; - failures += !expect_user_error("too_large", scenario_too_large); - failures += !expect_user_error("bad_result_type", scenario_bad_result_type); - failures += !expect_user_error("naive_rhs", scenario_naive_rhs); - failures += !expect_user_error("indirect", scenario_indirect); - failures += !expect_user_error("conv1d", scenario_conv1d); - failures += !expect_user_error("no_matmul", scenario_no_matmul); - failures += !expect_user_error("widening_16bit", scenario_widening_16bit); - failures += !expect_user_error("inconsistent_tiles", scenario_inconsistent_tiles); - failures += !expect_user_error("not_a_matmul_pattern", scenario_not_a_matmul_pattern); - failures += !expect_user_error("matmul_by_constant", scenario_matmul_by_constant); + failures += !expect_user_error("too_large", "too large to fit in", scenario_too_large); + failures += !expect_user_error("bad_result_type", "must yield 32-bit integers", scenario_bad_result_type); + failures += !expect_user_error("naive_rhs", "storage layout for a matrix multiply operand is unsupported", scenario_naive_rhs); + failures += !expect_user_error("indirect", "not loads with affine indices", scenario_indirect); + failures += !expect_user_error("sign_changing_cast", "cast after being loaded", scenario_sign_changing_cast); + failures += !expect_user_error("conv1d", "storage layout for a matrix multiply operand is unsupported", scenario_conv1d); + failures += !expect_user_error("no_matmul", "no matrix multiply operation was found", scenario_no_matmul); + failures += !expect_user_error("widening_16bit", "operand or result types are not supported", scenario_widening_16bit); + failures += !expect_user_error("mismatched_strides", "has different size and strides", scenario_mismatched_strides); + failures += !expect_user_error("inconsistent_tiles", "does not have the same shape", scenario_inconsistent_tiles); + failures += !expect_user_error("not_a_matmul_pattern", "operand or result types are not supported", scenario_not_a_matmul_pattern); + failures += !expect_user_error("matmul_by_constant", "not loads with affine indices", scenario_matmul_by_constant); if (failures != 0) { printf("%d scenario(s) failed to produce a user-facing CompileError\n", failures); diff --git a/test/correctness/wmma_matmul.cpp b/test/correctness/wmma_matmul.cpp new file mode 100644 index 000000000000..cb4779ab865f --- /dev/null +++ b/test/correctness/wmma_matmul.cpp @@ -0,0 +1,593 @@ +#include "Halide.h" +#include + +using namespace Halide; + +namespace { + +struct Params { + // The type of the two matrices being multiplied. The accumulator type + // follows from it: the integer instructions accumulate into int32, and the + // float ones into float32 unless half_accumulator asks otherwise. + Type operand = Float(16); + // The size of the matrices. + int M = 64, N = 64, K = 64; + // The tensor core tile shape to use. + int tile_m = 16, tile_n = 16, tile_k = 16; + // How many tiles of accumulator each warp holds, in each dimension. + int tiles_m = 1, tiles_n = 1; + // How many warps per block. + int warps = 1; + // Whether each input matrix is stored with its rows or its columns dense. + bool a_transposed = false, b_transposed = false; + // Whether to accumulate in half precision instead of single precision. + // Only available with half precision operands. + bool half_accumulator = false; + // Whether to stage the operand tiles through shared memory inside the + // reduction loop. + bool stage_in_shared = false; + // Whether to start the accumulator from a matrix in memory rather than + // from zero. + bool init_from_memory = false; + // Whether the output is stored with its columns dense rather than its rows. + bool out_transposed = false; +}; + +std::ostream &operator<<(std::ostream &s, const Params &p) { + return s << p.M << "x" << p.N << "x" << p.K + << " in " << p.tile_m << "x" << p.tile_n << "x" << p.tile_k + << " tiles, " << p.tiles_m << "x" << p.tiles_n << " tiles per warp, " + << p.warps << " warps, a" << (p.a_transposed ? "T" : "") + << " b" << (p.b_transposed ? "T" : "") + << " of " << p.operand + << (p.half_accumulator ? ", f16 accumulator" : "") + << (p.stage_in_shared ? ", staged through shared" : "") + << (p.init_from_memory ? ", accumulator initialized from memory" : "") + << (p.out_transposed ? ", transposed output" : ""); +} + +// Small non-negative integers, which every type here represents exactly, and +// which keep the dot products small enough to be exact in the accumulator too. +// That lets the results be checked for equality rather than to a tolerance. +void fill(Buffer<> buf) { + Type t = buf.type(); + for (int y = 0; y < buf.height(); y++) { + for (int x = 0; x < buf.width(); x++) { + int v = rand() & 3; + if (t == Float(16)) { + buf.as()(x, y) = float16_t((float)v); + } else if (t == BFloat(16)) { + buf.as()(x, y) = bfloat16_t((float)v); + } else if (t == Int(8)) { + buf.as()(x, y) = (int8_t)v; + } else if (t == UInt(8)) { + buf.as()(x, y) = (uint8_t)v; + } else { + assert(false && "unhandled operand type"); + } + } + } +} + +// Read an element of a buffer of any of the types this test uses. +double read(const Buffer<> &buf, int x, int y) { + Type t = buf.type(); + if (t == Float(16)) { + return (double)Buffer(buf)(x, y); + } else if (t == BFloat(16)) { + return (double)Buffer(buf)(x, y); + } else if (t == Float(32)) { + return Buffer(buf)(x, y); + } else if (t == Int(8)) { + return Buffer(buf)(x, y); + } else if (t == UInt(8)) { + return Buffer(buf)(x, y); + } else if (t == Int(32)) { + return Buffer(buf)(x, y); + } + assert(false && "unhandled buffer type"); + return 0; +} + +// The accumulator the hardware pairs with a given operand type. +Type accumulator_for(const Params &p) { + if (p.operand.is_int() || p.operand.is_uint()) { + return Int(32); + } + return p.half_accumulator ? Float(16) : Float(32); +} + +bool test(const Params &p) { + // Halide indexes matrices with the dense dimension first, so A(k, y) is a + // row-major M x K matrix, and A(y, k) is a column-major one. + Buffer<> A(p.operand, p.a_transposed ? p.M : p.K, p.a_transposed ? p.K : p.M); + Buffer<> B(p.operand, p.b_transposed ? p.K : p.N, p.b_transposed ? p.N : p.K); + fill(A); + fill(B); + + Var x("x"), y("y"), kk("kk"), xx("xx"), yy("yy"); + RDom k(0, p.K, "k"); + Func prod("prod"), out("out"), Af("Af"), Bf("Bf"), init("init"); + + // These are inlined unless we're staging through shared memory. + Af(kk, yy) = p.a_transposed ? A(yy, kk) : A(kk, yy); + Bf(xx, kk) = p.b_transposed ? B(kk, xx) : B(xx, kk); + Expr a = Af(k, y); + Expr b = Bf(x, k); + // With a half-precision accumulator the tile is copied out to memory as + // float16, so the output has to be float16 too. + // The accumulator either starts at zero, or at a matrix that already exists + // in memory, which the hardware can load straight into the fragments. + Type acc_type = accumulator_for(p); + init(x, y) = cast(acc_type, (x * 3 + y) % 7); + prod(x, y) = p.init_from_memory ? init(x, y) : cast(acc_type, 0); + prod(x, y) += cast(acc_type, a) * cast(acc_type, b); + out(x, y) = prod(x, y); + + Var xi("xi"), yi("yi"), xt("xt"), mmxi("mmxi"), mmyi("mmyi"); + Var rxi("rxi"), ryi("ryi"); + RVar rro("rro"), rri("rri"); + + // Each block computes a tile of the output. Within a block, each warp + // computes a strip of that tile, and each warp's strip is broken into + // tensor core tiles, which are the innermost two (vectorized) dimensions. + out.bound(x, 0, p.N) + .bound(y, 0, p.M) + .split(x, x, xi, p.tile_n * p.tiles_n * p.warps) + .split(xi, xt, xi, p.tile_n * p.tiles_n) + .split(y, y, yi, p.tile_m * p.tiles_m) + .tile(xi, yi, mmxi, mmyi, p.tile_n, p.tile_m) + .gpu_blocks(x, y) + .reorder(xi, yi, xt, x, y) + .unroll(xi) + .unroll(yi) + .tile_store(mmxi, mmyi); + if (p.warps > 1) { + // With one warp there's no need for a loop over warps, and leaving it + // serial keeps anything computed inside it out of the thread loops. + out.gpu_threads(xt); + } + + prod.compute_at(out, xt) + .tile(x, y, rxi, ryi, p.tile_n, p.tile_m) + .unroll(x) + .unroll(y); + // The accumulator either starts at a constant or is loaded from memory, + // and those are different tile instructions. + if (p.init_from_memory) { + prod.tile_load(rxi, ryi); + } else { + prod.tile_init(rxi, ryi); + } + + prod.update() + .tile(x, y, rxi, ryi, p.tile_n, p.tile_m) + .split(k, rro, rri, p.tile_k) + .reorder(x, y, rro) + .unroll(x) + .unroll(y) + .tile_matmul(rri, rxi, ryi); + + if (p.init_from_memory) { + Var ixi("ixi"), iyi("iyi"); + init.compute_root().gpu_tile(x, y, ixi, iyi, 16, 16); + } + + if (p.stage_in_shared) { + Var t("t"), ti("ti"), to("to"); + Af.compute_at(prod, rro) + .store_in(MemoryType::GPUShared) + .fuse(kk, yy, t) + .split(t, to, ti, 32) + .gpu_lanes(ti); + Bf.compute_at(prod, rro) + .store_in(MemoryType::GPUShared) + .fuse(xx, kk, t) + .split(t, to, ti, 32) + .gpu_lanes(ti); + } + + if (p.out_transposed) { + out.output_buffer().dim(0).set_stride(p.M).dim(1).set_stride(1); + } + + // A transposed output is a view of a buffer with the dimensions swapped, so + // its columns are dense in memory instead of its rows. + Buffer<> result_storage(acc_type, p.out_transposed ? p.M : p.N, + p.out_transposed ? p.N : p.M); + Buffer<> result = + p.out_transposed ? result_storage.transposed(0, 1) : result_storage; + out.realize(result); + result.copy_to_host(); + + for (int j = 0; j < p.M; j++) { + for (int i = 0; i < p.N; i++) { + double ref = p.init_from_memory ? (double)((i * 3 + j) % 7) : 0.0; + for (int l = 0; l < p.K; l++) { + ref += read(A, p.a_transposed ? j : l, p.a_transposed ? l : j) * + read(B, p.b_transposed ? l : i, p.b_transposed ? i : l); + } + // The operands are small integers and the dot products stay small + // enough to be exact in every accumulator type here, so the answer + // should be exact however the GPU orders the accumulation. + if (read(result, i, j) != ref) { + std::cerr << "Mismatch at " << i << ", " << j << ": " + << read(result, i, j) << " != " << ref << "\n" + << "For matmul of " << p << "\n"; + return false; + } + } + } + return true; +} + +// The accumulators can also live outside the loop over warps, which is what +// lets the reduction loop sit above it, so that the operand panels can be +// staged into shared memory once per block and reused by every warp. +bool test_block_level_accumulator() { + const int M = 128, N = 128, K = 128; + const int tile = 16, tiles_x = 2, tiles_y = 2, warps = 2, bk = 32; + const int block_x = tile * tiles_x * warps, block_y = tile * tiles_y; + + Buffer A(K, M), B(N, K); + fill(A); + fill(B); + + Var x("x"), y("y"), kk("kk"), yy("yy"), xx("xx"); + RDom k(0, K, "k"); + Func prod("prod"), out("out"), As("As"), Bs("Bs"); + + As(kk, yy) = A(kk, yy); + Bs(xx, kk) = B(xx, kk); + prod(x, y) = 0.f; + prod(x, y) += cast(As(k, y)) * cast(Bs(x, k)); + out(x, y) = prod(x, y); + + Var xi("xi"), xt("xt"), yi("yi"), mmxi("mmxi"), mmyi("mmyi"); + Var rxi("rxi"), ryi("ryi"), xw("xw"), t("t"), ti("ti"), tw("tw"), to("to"); + Var kko("kko"), kki("kki"), xxo("xxo"), xxi("xxi"); + RVar ko("ko"), ki("ki"), rri("rri"); + + out.bound(x, 0, N) + .bound(y, 0, M) + .tile(x, y, xi, yi, block_x, block_y) + .split(xi, xt, xi, tile * tiles_x) + .tile(xi, yi, mmxi, mmyi, tile, tile) + .gpu_blocks(x, y) + .gpu_threads(xt) + .reorder(xi, yi, xt, x, y) + .unroll(xi) + .unroll(yi) + .tile_store(mmxi, mmyi); + + prod.compute_at(out, x) + .split(x, xw, xi, tile * tiles_x) + .split(xi, xi, rxi, tile) + .split(y, y, ryi, tile) + .reorder(xi, y, xw) + .gpu_threads(xw) + .tile_init(rxi, ryi) + .unroll(xi) + .unroll(y); + + prod.update() + .split(k, ko, ki, bk) + .split(x, xw, xi, tile * tiles_x) + .split(xi, xi, rxi, tile) + .split(y, y, ryi, tile) + .split(ki, ki, rri, tile) + .reorder(xi, y, ki, xw, ko) + .gpu_threads(xw) + .unroll(xi) + .unroll(y) + .unroll(ki) + .tile_matmul(rri, rxi, ryi); + + As.compute_at(prod, ko) + .store_in(MemoryType::GPUShared) + .split(kk, kko, kki, 8) + .fuse(kko, yy, t) + .split(t, t, ti, 32) + .split(t, to, tw, warps) + .gpu_lanes(ti) + .gpu_threads(tw) + .vectorize(kki); + Bs.compute_at(prod, ko) + .store_in(MemoryType::GPUShared) + .split(xx, xxo, xxi, 8) + .fuse(xxo, kk, t) + .split(t, t, ti, 32) + .split(t, to, tw, warps) + .gpu_lanes(ti) + .gpu_threads(tw) + .vectorize(xxi); + + Buffer result(N, M); + out.realize(result); + result.copy_to_host(); + + for (int j = 0; j < M; j++) { + for (int i = 0; i < N; i++) { + float ref = 0.f; + for (int l = 0; l < K; l++) { + ref += (float)A(l, j) * (float)B(i, l); + } + if (std::abs(result(i, j) - ref) > 1e-2f * std::max(1.f, std::abs(ref))) { + std::cerr << "Mismatch at " << i << ", " << j << ": " + << result(i, j) << " != " << ref << "\n" + << "For a block-level accumulator staged through shared memory\n"; + return false; + } + } + } + return true; +} + +// Stage the operand tiles into fragment registers, so that each fragment +// loaded feeds several multiplies. Where each staging happens says how much +// reuse we get out of it. +bool test_staged_operands() { + const int M = 128, N = 128, K = 64; + const int tile = 16, tiles_x = 2, tiles_y = 2, bk = 32; + + Buffer A(K, M), B(N, K); + fill(A); + fill(B); + + Var x("x"), y("y"), kk("kk"), yy("yy"), xx("xx"); + RDom k(0, K, "k"); + Func prod("prod"), out("out"), Am("Am"), Bm("Bm"); + + Am(kk, yy) = A(kk, yy); + Bm(xx, kk) = B(xx, kk); + prod(x, y) = 0.f; + prod(x, y) += cast(Am(k, y)) * cast(Bm(x, k)); + out(x, y) = prod(x, y); + + Var xi("xi"), yi("yi"), mmxi("mmxi"), mmyi("mmyi"), rxi("rxi"), ryi("ryi"); + Var kko("kko"), kki("kki"), yyo("yyo"), yyi("yyi"), xxo("xxo"), xxi("xxi"); + RVar ko("ko"), ki("ki"), rri("rri"); + + out.bound(x, 0, N) + .bound(y, 0, M) + .tile(x, y, xi, yi, tile * tiles_x, tile * tiles_y) + .tile(xi, yi, mmxi, mmyi, tile, tile) + .gpu_blocks(x, y) + .reorder(xi, yi, x, y) + .unroll(xi) + .unroll(yi) + .tile_store(mmxi, mmyi); + + prod.compute_at(out, x) + .tile(x, y, rxi, ryi, tile, tile) + .tile_init(rxi, ryi) + .unroll(x) + .unroll(y); + + // Loop nest of the update, outermost first: ko, ki, y, x. + prod.update() + .split(k, ko, ki, bk) + .tile(x, y, rxi, ryi, tile, tile) + .split(ki, ki, rri, tile) + .reorder(x, y, ki, ko) + .unroll(x) + .unroll(y) + .unroll(ki) + .tile_matmul(rri, rxi, ryi); + + // One a fragment per row of tiles, live across the loop over columns. + Am.compute_at(prod, y) + .split(kk, kko, kki, tile) + .split(yy, yyo, yyi, tile) + .reorder(kki, yyi, kko, yyo) + .unroll(kko) + .unroll(yyo) + .tile_load(kki, yyi); + + // All the b fragments at once, live across the loops over both. + Bm.compute_at(prod, ki) + .split(xx, xxo, xxi, tile) + .split(kk, kko, kki, tile) + .reorder(xxi, kki, xxo, kko) + .unroll(xxo) + .unroll(kko) + .tile_load(xxi, kki); + + Buffer result(N, M); + out.realize(result); + result.copy_to_host(); + + for (int j = 0; j < M; j++) { + for (int i = 0; i < N; i++) { + float ref = 0.f; + for (int l = 0; l < K; l++) { + ref += (float)A(l, j) * (float)B(i, l); + } + if (std::abs(result(i, j) - ref) > 1e-2f * std::max(1.f, std::abs(ref))) { + std::cerr << "Mismatch at " << i << ", " << j << ": " + << result(i, j) << " != " << ref << "\n" + << "For operands staged in fragment registers\n"; + return false; + } + } + } + return true; +} + +// A batch of matrix multiplies that all share the same left-hand side. Its +// fragments only need loading once for the whole batch, which is what staging +// them outside the loop over the batch does. The loop isn't unrolled, so +// nothing downstream of here could hoist them out of it. +bool test_operand_hoisted_out_of_loop() { + const int M = 32, N = 32, K = 16, batch = 4; + const int tile = 16; + + Buffer A(K, M), B(N, K, batch); + fill(A); + fill(B); + + Var x("x"), y("y"), n("n"), kk("kk"), yy("yy"); + RDom k(0, K, "k"); + Func prod("prod"), out("out"), Am("Am"); + + Am(kk, yy) = A(kk, yy); + prod(x, y, n) = 0.f; + prod(x, y, n) += cast(Am(k, y)) * cast(B(x, k, n)); + out(x, y, n) = prod(x, y, n); + + Var xi("xi"), yi("yi"), mmxi("mmxi"), mmyi("mmyi"), rxi("rxi"), ryi("ryi"); + Var kko("kko"), kki("kki"), yyo("yyo"), yyi("yyi"); + RVar rro("rro"), rri("rri"); + + out.bound(x, 0, N) + .bound(y, 0, M) + .bound(n, 0, batch) + .tile(x, y, xi, yi, N, M) + .tile(xi, yi, mmxi, mmyi, tile, tile) + .gpu_blocks(x, y) + .reorder(xi, yi, n, x, y) + .unroll(xi) + .unroll(yi) + .tile_store(mmxi, mmyi); + + prod.compute_at(out, n) + .tile(x, y, rxi, ryi, tile, tile) + .tile_init(rxi, ryi) + .unroll(x) + .unroll(y); + + prod.update() + .tile(x, y, rxi, ryi, tile, tile) + .split(k, rro, rri, tile) + .reorder(x, y, rro) + .unroll(x) + .unroll(y) + .tile_matmul(rri, rxi, ryi); + + Am.compute_at(out, x) + .split(kk, kko, kki, tile) + .split(yy, yyo, yyi, tile) + .reorder(kki, yyi, kko, yyo) + .unroll(kko) + .unroll(yyo) + .tile_load(kki, yyi); + + Buffer result(N, M, batch); + out.realize(result); + result.copy_to_host(); + + for (int b = 0; b < batch; b++) { + for (int j = 0; j < M; j++) { + for (int i = 0; i < N; i++) { + float ref = 0.f; + for (int l = 0; l < K; l++) { + ref += (float)A(l, j) * (float)B(i, l, b); + } + if (std::abs(result(i, j, b) - ref) > 1e-2f * std::max(1.f, std::abs(ref))) { + std::cerr << "Mismatch at " << i << ", " << j << ", " << b << ": " + << result(i, j, b) << " != " << ref << "\n" + << "For an operand staged outside a loop over a batch\n"; + return false; + } + } + } + } + return true; +} + +} // namespace + +int main(int argc, char **argv) { + Target target = get_jit_target_from_environment(); + if (!target.has_feature(Target::CUDA)) { + printf("[SKIP] WMMA matrix multiplies require CUDA.\n"); + return 0; + } + if (target.get_cuda_capability_lower_bound() < 70) { + printf("[SKIP] WMMA matrix multiplies require CUDA compute capability 7.0 or above.\n"); + return 0; + } + + std::vector params; + + // The simplest possible case. + params.push_back({}); + + // Each of the other supported tile shapes. + params.push_back({.tile_m = 32, .tile_n = 8}); + params.push_back({.tile_m = 8, .tile_n = 32}); + + // Each combination of input layouts. + params.push_back({.a_transposed = true}); + params.push_back({.b_transposed = true}); + params.push_back({.a_transposed = true, .b_transposed = true}); + + // Several accumulator fragments live at once, which is what you want in + // practice so that each operand load feeds more than one multiply. + params.push_back({.tiles_m = 2, .tiles_n = 2}); + + // More than one warp per block. + params.push_back({.warps = 2}); + params.push_back({.tiles_m = 2, .tiles_n = 2, .warps = 2}); + + // Accumulating in half precision, which uses half as many registers per + // accumulator fragment. + params.push_back({.half_accumulator = true}); + params.push_back({.tiles_m = 2, .tiles_n = 2, .half_accumulator = true}); + + // The other operand types the hardware multiplies. Brain floats accumulate + // into single precision, and eight-bit integers into 32-bit ones, which is + // the interesting case for imaging. + for (Type t : {BFloat(16), Int(8), UInt(8)}) { + params.push_back({.operand = t}); + params.push_back({.operand = t, .tile_m = 32, .tile_n = 8}); + params.push_back({.operand = t, .tile_m = 8, .tile_n = 32}); + params.push_back({.operand = t, .tiles_m = 2, .tiles_n = 2}); + params.push_back({.operand = t, .warps = 2}); + params.push_back({.operand = t, .a_transposed = true}); + params.push_back({.operand = t, .b_transposed = true}); + params.push_back({.operand = t, .stage_in_shared = true}); + params.push_back({.operand = t, .init_from_memory = true}); + params.push_back({.operand = t, .out_transposed = true}); + } + + // Operand tiles staged through shared memory inside the reduction loop. + // This only works if the loop over lanes wraps the individual wmma + // statements rather than the whole accumulator allocation. + params.push_back({.stage_in_shared = true}); + params.push_back({.tiles_m = 2, .tiles_n = 2, .stage_in_shared = true}); + params.push_back({.half_accumulator = true, .stage_in_shared = true}); + + // Accumulators that start from a matrix already in memory rather than from + // zero, which uses the load-into-the-accumulator-fragment instruction. + params.push_back({.init_from_memory = true}); + params.push_back({.tiles_m = 2, .tiles_n = 2, .init_from_memory = true}); + params.push_back({.half_accumulator = true, .init_from_memory = true}); + params.push_back({.stage_in_shared = true, .init_from_memory = true}); + + // A column-major output, which the accumulator is stored to with the other + // layout of the store instruction. + params.push_back({.out_transposed = true}); + params.push_back({.tiles_m = 2, .tiles_n = 2, .out_transposed = true}); + + // A reduction that isn't a whole number of tiles per unrolled step, and + // matrices that aren't square. + params.push_back({.M = 32, .N = 128, .K = 256}); + + for (const Params &p : params) { + if (!test(p)) { + printf("Failed!\n"); + return 1; + } + } + + if (!test_block_level_accumulator() || + !test_staged_operands() || + !test_operand_hoisted_out_of_loop()) { + printf("Failed!\n"); + return 1; + } + + printf("Success!\n"); + return 0; +} diff --git a/test/performance/tiled_matmul.cpp b/test/performance/tiled_matmul.cpp index 916aa354eb57..c990ca052e32 100644 --- a/test/performance/tiled_matmul.cpp +++ b/test/performance/tiled_matmul.cpp @@ -107,7 +107,7 @@ bool matmul(Halide::Target target) { Var rxi("rxi"), ryi("ryi"); RVar rri("rri"), rro("rro"); mm.compute_at(mm.in(), y) - .store_in(MemoryType::AMXTile) + .store_in(MemoryType::Tile) .update() // Split into (x,y) tile .tile(y, x, ryi, rxi, tile_y, tile_x, TailStrategy::GuardWithIf) @@ -196,7 +196,7 @@ bool matmul_bf16(Halide::Target target) { RVar rri("rri"), rro("rro"); mm.compute_at(mm.in(), x) - .store_in(MemoryType::AMXTile) + .store_in(MemoryType::Tile) .update() .tile(x, y, rxi, ryi, tile_x, tile_y, TailStrategy::GuardWithIf) .split(r.x, rro, rri, tile_r)