Merge branch 'dev' into feature/reduce

Author: junjihashimoto

Makefile

@@ -19,8 +19,37 @@ pch:
 	mkdir -p build && $(CXX) -std=c++17 $(INCLUDES) -x c++-header gpu.hpp -o build/gpu.hpp.pch
 
 # TODO(avh): change extension based on platform
-lib:
-	mkdir -p build && $(CXX) -std=c++17 $(INCLUDES) -L$(LIBDIR) -ldawn -ldl -shared -fPIC gpu.cpp -o build/libgpucpp.dylib
+# Get the current OS name
+OS = $(shell uname | tr -d '\n')
+# Set the specific variables for each platform
+LIB_PATH ?= /usr/lib
+HEADER_PATH ?= /usr/include
+ifeq ($(OS), Linux)
+OS_TYPE ?= Linux
+GPU_CPP_LIB_NAME ?= libgpucpp.so
+DAWN_LIB_NAME ?= libdawn.so
+else ifeq ($(OS), Darwin)
+OS_TYPE ?= macOS
+GPU_CPP_LIB_NAME ?= libgpucpp.dylib
+DAWN_LIB_NAME ?= libdawn.dylib
+else
+OS_TYPE ?= unknown
+endif
+
+lib: check-clang dawnlib
+	mkdir -p build && $(CXX) -std=c++17 $(INCLUDES) -L$(LIBDIR) -ldawn -ldl -shared -fPIC gpu.cpp -o build/$(GPU_CPP_LIB_NAME)
+	python3 build.py
+	cp third_party/lib/$(DAWN_LIB_NAME) build/
+
+install:
+	cp build/$(GPU_CPP_LIB_NAME) $(LIB_PATH)
+	cp build/$(DAWN_LIB_NAME) $(LIB_PATH)
+	cp build/gpu.hpp $(HEADER_PATH)
+
+uninstall:
+	rm $(LIB_PATH)/$(GPU_CPP_LIB_NAME)
+	rm $(LIB_PATH)/$(DAWN_LIB_NAME)
+	rm $(HEADER_PATH)/gpu.hpp
 
 examples/hello_world/build/hello_world: check-clang dawnlib examples/hello_world/run.cpp check-linux-vulkan
 	$(LIBSPEC) && cd examples/hello_world && make build/hello_world && ./build/hello_world
@@ -96,15 +125,24 @@ clean-all:
 # Checks
 ################################################################################
 
+# Check all
+check-all: check-os check-clang check-cmake check-python
+
+# check the os
+check-os:
+ifeq ($(OS_TYPE), unknown)
+$(error Unsupported operating system)
+endif
+
 # check for the existence of clang++ and cmake
 check-clang:
-	@command -v clang++ >/dev/null 2>&1 || { echo >&2 "Please install clang++ with 'sudo apt-get install clang' or 'brew install llvm'"; exit 1; }
+	@command -v clang++ >/dev/null 2>&1 || { echo -e >&2 "Clang++ is not installed. Please install clang++ to continue.\nOn Debian / Ubuntu: 'sudo apt-get install clang' or 'brew install llvm'\nOn Centos: 'sudo yum install clang'"; exit 1; }
 
 check-cmake:
-	@command -v cmake >/dev/null 2>&1 || { echo >&2 "Please install cmake with 'sudo apt-get install cmake' or 'brew install cmake'"; exit 1; }
+	@command -v cmake >/dev/null 2>&1 || { echo -e >&2 "Cmake is not installed. Please install cmake to continue.\nOn Debian / Ubuntu: 'sudo apt-get install cmake' or 'brew install cmake'\nOn Centos: 'sudo yum install cmake'"; exit 1; }
 
 check-python:
-	@command -v python3 >/dev/null 2>&1 || { echo >&2 "Python needs to be installed and in your path."; exit 1; } 
+	@command -v python3 >/dev/null 2>&1 || { echo -e >&2 "Python is not installed. Please install python to continue.\nOn Debian / Ubuntu: 'sudo apt-get install python'\nOn Centos: 'sudo yum install python'"; exit 1; } 
 
 check-linux-vulkan:
 	@echo "Checking system type and Vulkan availability..."
@@ -113,7 +151,7 @@ check-linux-vulkan:
 	        echo "Vulkan is installed."; \
 	        vulkaninfo; \
 	    else \
-		echo "Vulkan is not installed. Please install Vulkan drivers to continue. On Debian / Ubuntu: sudo apt install libvulkan1 mesa-vulkan-drivers vulkan-tools"; \
+		echo -e "Vulkan is not installed. Please install Vulkan drivers to continue.\nOn Debian / Ubuntu: 'sudo apt install libvulkan1 mesa-vulkan-drivers vulkan-tools'.\nOn Centos: 'sudo yum install vulkan vulkan-tools.'"; \
 	        exit 1; \
 	    fi \
 	else \

build.py

@@ -0,0 +1,32 @@
+# Dictionary of header files and their relative paths
+header_files = {
+    "#include \"webgpu/webgpu.h\"": "third_party/headers/webgpu/webgpu.h",
+    "#include \"numeric_types/half.hpp\"": "numeric_types/half.hpp",
+    "#include \"utils/logging.hpp\"": "utils/logging.hpp"
+}
+
+def main():
+    # File paths
+    source_file_path = "gpu.hpp"
+    output_file_path = "build/gpu.hpp"
+
+    # Open source file and read contents
+    with open(source_file_path, "r") as source:
+        file_contents = source.read()
+
+    # Ergodic over header files
+    for key, value in header_files.items():
+
+        # Replace header files
+        with open(value, "r") as header_file:
+            header_file_contents = header_file.read()
+        file_contents = file_contents.replace(key, header_file_contents)
+        
+
+    # Open output file
+    with open(output_file_path, "w") as output:
+        # Write contents to output file
+        output.write(file_contents)
+
+if __name__ == "__main__":
+    main()

experimental/kernels/Makefile

@@ -83,7 +83,7 @@ endef
 build/test_gpt2: llm.c build/unittest_kernels.o gpt2_124M.bin
 	mkdir -p build
 	$(call preprocess_file)
-	$(CC) $(CFLAGS) $(LDFLAGS) -o $@ llm.c/test_gpt2.c build/unittest_kernels.o
+	$(CC) $(CFLAGS) $(LDFLAGS) -o $@ llm.c/test_gpt2.c build/unittest_kernels.o -g
 
 build/test_gpt2_with_metal_profiler: llm.c build/unittest_kernels.o gpt2_124M.bin
 	mkdir -p build
@@ -94,12 +94,12 @@ build/test_gpt2_with_metal_profiler: llm.c build/unittest_kernels.o gpt2_124M.bi
 build/train_gpt2: llm.c build/unittest_kernels.o gpt2_124M.bin
 	mkdir -p build
 	$(call preprocess_file)
-	$(CC) $(CFLAGS) $(LDFLAGS) -o $@ llm.c/train_gpt2.c build/unittest_kernels.o
+	$(CC) $(CFLAGS) $(LDFLAGS) -o $@ llm.c/train_gpt2.c build/unittest_kernels.o -g
 
 build/ops.o: ops.cpp ops.hpp kernels.h llm.c
 	mkdir -p build && $(CXX) $(CXXFLAGS) -c -o $@ $<
 
-build/gpt2_webgpu: llm.c gpt2_124M.bin llm.c
+build/gpt2_webgpu: llm.c gpt2_124M.bin llm.c gpt2_webgpu.cpp ops.cpp
 	mkdir -p build
 	$(CC) $(CXXFLAGS) -Illm.c $(LDFLAGS) -o $@ gpt2_webgpu.cpp ops.cpp
 

experimental/kernels/kernels.h

@@ -309,6 +309,104 @@ fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
       }
     }
 }
+
+)";
+
+
+static const char *kShaderMatmul2DTiling = R"(
+@group(0) @binding(0) var<storage, read_write> inp : array<{{precision}}>;
+@group(0) @binding(1) var<storage, read_write> weight : array<{{precision}}>;
+@group(0) @binding(2) var<storage, read_write> bias : array<{{precision}}>;
+@group(0) @binding(3) var<storage, read_write> out : array<{{precision}}>;
+@group(0) @binding(4) var<uniform> params : Params;
+struct Params {
+    B: u32,
+    T: u32,
+    C: u32,
+    OC: u32,
+};
+var<workgroup> tileInp: array<{{precision}}, {{BT}} * {{BC}}>;
+var<workgroup> tileWeight: array<{{precision}}, {{BOC}} * {{BC}}>;
+
+@compute @workgroup_size({{workgroupSize}})
+fn main(
+    @builtin(local_invocation_id) localID : vec3<u32>,
+    @builtin(workgroup_id) groupid : vec3<u32>) {
+    let B : u32 = params.B;
+    let T : u32 = params.T;
+    let C : u32 = params.C;
+    let OC : u32 = params.OC;
+
+    var localT: array<{{precision}}, {{TT}}>;
+    var localOC: array<{{precision}}, {{TOC}}>;
+
+    let outB: u32 = groupid.x;
+    let outT: u32 = groupid.y;
+    let outOC: u32 = groupid.z;
+    let numThread: u32 = ({{BT}} * {{BOC}}) / ({{TT}} * {{TOC}});
+
+    // position of the first c element computed by the thread
+    let threadRow: u32 = (localID.x / ({{BOC}} / {{TOC}})) * {{TT}};
+    let threadCol: u32 = (localID.x % ({{BOC}} / {{TOC}})) * {{TOC}};
+
+    // inpPtr and weightPtr are the starting positions of the tiles in a and b,
+    // incremented in the bkidx loop. 
+    // outPtr is the starting position of the tile in c which is fixed.
+
+    var inpPtr = (outB * T + outT * {{BT}}) * C; // BTC 
+    var weightPtr = outOC * {{BOC}} * C; //OCC
+    var threadResults: array<{{precision}}, {{TT}} * {{TOC}}>;
+    let outPtr = (outB * T + outT * {{BT}}) * OC + outOC * {{BOC}}; //BTOC
+    let biasPtr = outOC * {{BOC}};
+
+    for (var bkidx: u32 = 0; bkidx < C; bkidx += {{BC}}) {
+      // Load BC x BOC by numThread(BT * BOC / (TT * TOC))
+      // The number of iteration == BC * BOC / (BT * BOC / (TT * TOC))
+      for (var idx: u32 = 0; idx < {{NUM_TILEW}}; idx++) {
+        tileWeight[localID.x + idx * numThread] = weight[weightPtr + ((localID.x + idx * numThread) / {{BC}}) * C + ((localID.x + idx * numThread) % {{BC}})];
+      }
+      weightPtr += {{BC}};
+    
+      // Load tile
+      // Load BT x BC by numThread(BT * BOC / (TT * TOC))
+      // The number of iteration == BT * BC / (BT * BOC / (TT * TOC))
+      for (var idx: u32 = 0; idx < {{NUM_TILEI}}; idx++) {
+        tileInp[localID.x + idx * numThread] = inp[inpPtr + ((localID.x + idx * numThread) / {{BC}}) * C + (localID.x + idx * numThread) % {{BC}}];
+      }
+      inpPtr += {{BC}};
+    
+      workgroupBarrier();
+      // Compute tile
+      for (var dotIdx: u32 = 0; dotIdx < {{BC}}; dotIdx = dotIdx + 1) {
+        for (var idx: u32 = 0; idx < {{TT}}; idx++) {
+          localT[idx] = tileInp[(threadRow + idx) * {{BC}} + dotIdx];
+        }
+        for (var idx: u32 = 0; idx < {{TOC}}; idx++) {
+          localOC[idx] = tileWeight[(threadCol + idx) * {{BC}} + dotIdx];
+        }
+        for (var resIdxT: u32 = 0; resIdxT < {{TT}}; resIdxT++) {
+          for (var resIdxOC: u32 = 0; resIdxOC < {{TOC}}; resIdxOC++) {
+            threadResults[resIdxT * {{TOC}} + resIdxOC] += localT[resIdxT] * localOC[resIdxOC];
+          }
+        }
+      }
+      workgroupBarrier();
+    }
+    
+    if (arrayLength(&bias) == 1) {
+      for (var resIdxT: u32 = 0; resIdxT < {{TT}}; resIdxT++) {
+        for (var resIdxOC: u32 = 0; resIdxOC < {{TOC}}; resIdxOC++) {
+          out[outPtr + (threadRow + resIdxT) * OC + threadCol + resIdxOC] = threadResults[resIdxT * {{TOC}} + resIdxOC];
+        }
+      }
+    } else {
+      for (var resIdxT: u32 = 0; resIdxT < {{TT}}; resIdxT++) {
+        for (var resIdxOC: u32 = 0; resIdxOC < {{TOC}}; resIdxOC++) {
+          out[outPtr + (threadRow + resIdxT) * OC + threadCol + resIdxOC] = threadResults[resIdxT * {{TOC}} + resIdxOC] + bias[biasPtr + threadCol + resIdxOC];
+        }
+      }
+    }
+}
 )";
 
 static const char *kShaderMatmulBackward = R"(