Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions CLAUDE.md
Original file line numberDiff line numberDiff line change
Expand Up@@ -87,9 +87,9 @@ uv run ruff format gpu_test/

- **Stack Type**: `!forth.stack` - untyped stack, programmer ensures type safety
- **Operations**: All take stack as input and produce stack as output (except `forth.stack`)
- **Supported Words**: literals, `DUP DROP SWAP OVER ROT NIP TUCK PICK ROLL`, `+ - * / MOD`, `AND OR XOR NOT LSHIFT RSHIFT`, `= < > <> <= >= 0=`, `@ !`, `CELLS`, `IF ELSE THEN`, `BEGIN UNTIL`, `BEGIN WHILE REPEAT`, `DO LOOP +LOOP I J K`, `LEAVE UNLOOP EXIT`, `TID-X/Y/Z BID-X/Y/Z BDIM-X/Y/Z GDIM-X/Y/Z GLOBAL-ID` (GPU indexing).
- **Supported Words**: literals, `DUP DROP SWAP OVER ROT NIP TUCK PICK ROLL`, `+ - * / MOD`, `AND OR XOR NOT LSHIFT RSHIFT`, `= < > <> <= >= 0=`, `@ !` (global memory), `S@ S!` (shared memory), `CELLS`, `IF ELSE THEN`, `BEGIN UNTIL`, `BEGIN WHILE REPEAT`, `DO LOOP +LOOP I J K`, `LEAVE UNLOOP EXIT`, `TID-X/Y/Z BID-X/Y/Z BDIM-X/Y/Z GDIM-X/Y/Z GLOBAL-ID` (GPU indexing).
- **Kernel Parameters**: Declared in the `\!` header. `\! kernel <name>` is required and must appear first. `\! param <name> i64[<N>]` becomes a `memref<Nxi64>` argument; `\! param <name> i64` becomes an `i64` argument. Using a param name in code emits `forth.param_ref` (arrays push address; scalars push value).
- **Shared Memory**: `\! shared <name> i64[<N>]` declares GPU shared (workgroup) memory. Emits a tagged `memref.alloca` at kernel entry; ForthToGPU converts it to a `gpu.func` workgroup attribution (`memref<Nxi64, #gpu.address_space<workgroup>>`). Using the shared name in code pushes its base address onto the stack. Cannot be referenced inside word definitions.
- **Shared Memory**: `\! shared <name> i64[<N>]` declares GPU shared (workgroup) memory. Emits a tagged `memref.alloca` at kernel entry; ForthToGPU converts it to a `gpu.func` workgroup attribution (`memref<Nxi64, #gpu.address_space<workgroup>>`). Using the shared name in code pushes its base address onto the stack. Use `S@`/`S!` for shared accesses. Cannot be referenced inside word definitions.
- **Conversion**: `!forth.stack` → `memref<256xi64>` with explicit stack pointer
- **GPU**: Functions wrapped in `gpu.module`, `main` gets `gpu.kernel` attribute, configured with bare pointers for NVVM conversion
- **User-defined Words**: Modeled as `func.func` with signature `(!forth.stack) -> !forth.stack`, called via `func.call`
Expand Down
62 changes: 62 additions & 0 deletions gpu_test/test_kernels.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -337,6 +337,68 @@ def test_naive_matmul_i64(kernel_runner: KernelRunner) -> None:
assert result == [12, 6, 9, 28, 14, 29]


def test_tiled_matmul_i64(kernel_runner: KernelRunner) -> None:
"""Tiled i64 matmul with shared memory: C = A(4x4) * B(4x4) -> C(4x4).

Uses 2x2 tiles, shared memory for A/B tiles, and BARRIER for sync.
Grid: (2,2,1), Block: (2,2,1) — 4 blocks of 4 threads each.
"""
result = kernel_runner.run(
forth_source=(
"\\! kernel main\n"
"\\! param A i64[16]\n"
"\\! param B i64[16]\n"
"\\! param C i64[16]\n"
"\\! shared SA i64[4]\n"
"\\! shared SB i64[4]\n"
"BID-Y 2 * TID-Y +\n"
"BID-X 2 * TID-X +\n"
"0\n"
"2 0 DO\n"
" 2 PICK 4 * I 2 * + TID-X + CELLS A + @\n"
" TID-Y 2 * TID-X + CELLS SA + S!\n"
" I 2 * TID-Y + 4 * 2 PICK + CELLS B + @\n"
" TID-Y 2 * TID-X + CELLS SB + S!\n"
" BARRIER\n"
" 2 0 DO\n"
" TID-Y 2 * I + CELLS SA + S@\n"
" I 2 * TID-X + CELLS SB + S@\n"
" * +\n"
" LOOP\n"
" BARRIER\n"
"LOOP\n"
"ROT 4 * ROT + CELLS C + !"
),
params={
"A": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
"B": [17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32],
},
grid=(2, 2, 1),
block=(2, 2, 1),
output_param=2,
output_count=16,
)
expected = [
250,
260,
270,
280,
618,
644,
670,
696,
986,
1028,
1070,
1112,
1354,
1412,
1470,
1528,
]
assert result == expected


# --- User-Defined Words ---


Expand Down
18 changes: 18 additions & 0 deletions include/warpforth/Dialect/Forth/ForthOps.td
Original file line numberDiff line numberDiff line change
Expand Up@@ -250,6 +250,24 @@ def Forth_StoreOp : Forth_StackOpBase<"store"> {
}];
}

def Forth_SharedLoadOp : Forth_StackOpBase<"shared_load"> {
let summary = "Load value from shared memory buffer";
let description = [{
Pops an address from the stack, loads a value from shared/workgroup memory at
that address, and pushes the loaded value onto the stack.
Forth semantics: ( addr -- value )
}];
}

def Forth_SharedStoreOp : Forth_StackOpBase<"shared_store"> {
let summary = "Store value to shared memory buffer";
let description = [{
Pops an address and value from the stack, stores the value to shared/workgroup
memory at the specified address.
Forth semantics: ( x addr -- )
}];
}

def Forth_ParamRefOp : Forth_Op<"param_ref", [Pure]> {
let summary = "Push kernel parameter address onto stack";
let description = [{
Expand Down
85 changes: 83 additions & 2 deletions lib/Conversion/ForthToMemRef/ForthToMemRef.cpp
Original file line numberDiff line numberDiff line change
Expand Up@@ -9,6 +9,7 @@
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Func/Transforms/FuncConversions.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/BuiltinTypes.h"
Expand All@@ -27,6 +28,8 @@ namespace {

// Stack configuration constants
constexpr int64_t kStackSize = 256;
constexpr unsigned kWorkgroupAddressSpace =
static_cast<unsigned>(gpu::AddressSpace::Workgroup);

/// Type converter for forth.stack -> memref + index
class ForthToMemRefTypeConverter : public TypeConverter {
Expand DownExpand Up@@ -725,6 +728,83 @@ struct StoreOpConversion : public OpConversionPattern<forth::StoreOp> {
}
};

/// Conversion pattern for forth.shared_load operation (S@).
/// Pops address from stack, loads value via shared/workgroup pointer, pushes
/// value: ( addr -- value )
struct SharedLoadOpConversion
: public OpConversionPattern<forth::SharedLoadOp> {
SharedLoadOpConversion(const TypeConverter &typeConverter,
MLIRContext *context)
: OpConversionPattern<forth::SharedLoadOp>(typeConverter, context) {}
using OneToNOpAdaptor = OpConversionPattern::OneToNOpAdaptor;

LogicalResult
matchAndRewrite(forth::SharedLoadOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
ValueRange inputStack = adaptor.getOperands()[0];
Value memref = inputStack[0];
Value stackPtr = inputStack[1];

auto i64Type = rewriter.getI64Type();
auto ptrType = LLVM::LLVMPointerType::get(rewriter.getContext(),
kWorkgroupAddressSpace);

// Load address from stack.
Value addrValue = rewriter.create<memref::LoadOp>(loc, memref, stackPtr);

// Load value from shared memory via address-space-qualified pointer.
Value ptr = rewriter.create<LLVM::IntToPtrOp>(loc, ptrType, addrValue);
Value loadedValue = rewriter.create<LLVM::LoadOp>(loc, i64Type, ptr);

// Store loaded value back at same position (replaces address).
rewriter.create<memref::StoreOp>(loc, loadedValue, memref, stackPtr);

rewriter.replaceOpWithMultiple(op, {{memref, stackPtr}});
return success();
}
};

/// Conversion pattern for forth.shared_store operation (S!).
/// Pops address and value from stack, stores value via shared/workgroup
/// pointer: ( x addr -- )
struct SharedStoreOpConversion
: public OpConversionPattern<forth::SharedStoreOp> {
SharedStoreOpConversion(const TypeConverter &typeConverter,
MLIRContext *context)
: OpConversionPattern<forth::SharedStoreOp>(typeConverter, context) {}
using OneToNOpAdaptor = OpConversionPattern::OneToNOpAdaptor;

LogicalResult
matchAndRewrite(forth::SharedStoreOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
ValueRange inputStack = adaptor.getOperands()[0];
Value memref = inputStack[0];
Value stackPtr = inputStack[1];

auto ptrType = LLVM::LLVMPointerType::get(rewriter.getContext(),
kWorkgroupAddressSpace);

// Pop address from stack.
Value addrValue = rewriter.create<memref::LoadOp>(loc, memref, stackPtr);

// Pop value from stack.
Value one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
Value spMinus1 = rewriter.create<arith::SubIOp>(loc, stackPtr, one);
Value value = rewriter.create<memref::LoadOp>(loc, memref, spMinus1);

// Store value to shared memory via address-space-qualified pointer.
Value ptr = rewriter.create<LLVM::IntToPtrOp>(loc, ptrType, addrValue);
rewriter.create<LLVM::StoreOp>(loc, value, ptr);

// New stack pointer is SP-2 (popped both address and value).
Value spMinus2 = rewriter.create<arith::SubIOp>(loc, spMinus1, one);
rewriter.replaceOpWithMultiple(op, {{memref, spMinus2}});
return success();
}
};

/// Template for converting GPU indexing ops to intrinsic ops.
/// Creates an intrinsic op with the specified name and pushes the value onto
/// the stack.
Expand DownExpand Up@@ -1026,8 +1106,9 @@ struct ConvertForthToMemRefPass
NotOpConversion, LshiftOpConversion, RshiftOpConversion, EqOpConversion,
LtOpConversion, GtOpConversion, NeOpConversion, LeOpConversion,
GeOpConversion, ZeroEqOpConversion, ParamRefOpConversion,
LoadOpConversion, StoreOpConversion, PopFlagOpConversion,
PopOpConversion, PushValueOpConversion>(typeConverter, context);
LoadOpConversion, StoreOpConversion, SharedLoadOpConversion,
SharedStoreOpConversion, PopFlagOpConversion, PopOpConversion,
PushValueOpConversion>(typeConverter, context);

// Add GPU indexing op conversion patterns
patterns.add<IntrinsicOpConversion<forth::ThreadIdXOp>>(typeConverter,
Expand Down
6 changes: 6 additions & 0 deletions lib/Translation/ForthToMLIR/ForthToMLIR.cpp
Original file line numberDiff line numberDiff line change
Expand Up@@ -487,6 +487,12 @@ Value ForthParser::emitOperation(StringRef word, Value inputStack,
} else if (word == "!") {
return builder.create<forth::StoreOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "S@") {
return builder.create<forth::SharedLoadOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "S!") {
return builder.create<forth::SharedStoreOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "TID-X") {
return builder.create<forth::ThreadIdXOp>(loc, stackType, inputStack)
.getResult();
Expand Down
13 changes: 13 additions & 0 deletions test/Conversion/ForthToMemRef/memory-ops.mlir
Original file line numberDiff line numberDiff line change
Expand Up@@ -15,6 +15,14 @@
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr
// CHECK: llvm.store %{{.*}}, %{{.*}} : i64, !llvm.ptr

// shared load (S@): pop address, inttoptr shared addrspace, llvm.load
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr<{{[1-9][0-9]*}}>
// CHECK: llvm.load %{{.*}} : !llvm.ptr<{{[1-9][0-9]*}}> -> i64

// shared store (S!): pop address + value, inttoptr shared addrspace, llvm.store
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr<{{[1-9][0-9]*}}>
// CHECK: llvm.store %{{.*}}, %{{.*}} : i64, !llvm.ptr<{{[1-9][0-9]*}}>

module {
func.func private @main() {
%0 = forth.stack !forth.stack
Expand All@@ -23,6 +31,11 @@ module {
%3 = forth.literal %2 42 : !forth.stack -> !forth.stack
%4 = forth.literal %3 100 : !forth.stack -> !forth.stack
%5 = forth.store %4 : !forth.stack -> !forth.stack
%6 = forth.literal %5 2 : !forth.stack -> !forth.stack
%7 = forth.shared_load %6 : !forth.stack -> !forth.stack
%8 = forth.literal %7 9 : !forth.stack -> !forth.stack
%9 = forth.literal %8 3 : !forth.stack -> !forth.stack
%10 = forth.shared_store %9 : !forth.stack -> !forth.stack
return
}
}
8 changes: 7 additions & 1 deletion test/Translation/Forth/memory-ops.forth
Original file line numberDiff line numberDiff line change
Expand Up@@ -6,9 +6,15 @@
\ Test ! produces forth.store
\ CHECK: forth.store %{{.*}} : !forth.stack -> !forth.stack

\ Test S@ produces forth.shared_load
\ CHECK: forth.shared_load %{{.*}} : !forth.stack -> !forth.stack

\ Test S! produces forth.shared_store
\ CHECK: forth.shared_store %{{.*}} : !forth.stack -> !forth.stack

\ Test CELLS produces literal 8 + mul
\ CHECK: forth.literal %{{.*}} 8
\ CHECK-NEXT: forth.mul
\! kernel main
1 @ 2 3 !
1 @ 2 3 ! 4 S@ 5 6 S!
4 CELLS
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions CLAUDE.md
Original file line numberDiff line numberDiff line change
Expand Up@@ -87,9 +87,9 @@ uv run ruff format gpu_test/

- **Stack Type**: `!forth.stack` - untyped stack, programmer ensures type safety
- **Operations**: All take stack as input and produce stack as output (except `forth.stack`)
- **Supported Words**: literals, `DUP DROP SWAP OVER ROT NIP TUCK PICK ROLL`, `+ - * / MOD`, `AND OR XOR NOT LSHIFT RSHIFT`, `= < > <> <= >= 0=`, `@ !`, `CELLS`, `IF ELSE THEN`, `BEGIN UNTIL`, `BEGIN WHILE REPEAT`, `DO LOOP +LOOP I J K`, `LEAVE UNLOOP EXIT`, `TID-X/Y/Z BID-X/Y/Z BDIM-X/Y/Z GDIM-X/Y/Z GLOBAL-ID` (GPU indexing).
- **Supported Words**: literals, `DUP DROP SWAP OVER ROT NIP TUCK PICK ROLL`, `+ - * / MOD`, `AND OR XOR NOT LSHIFT RSHIFT`, `= < > <> <= >= 0=`, `@ !` (global memory), `S@ S!` (shared memory), `CELLS`, `IF ELSE THEN`, `BEGIN UNTIL`, `BEGIN WHILE REPEAT`, `DO LOOP +LOOP I J K`, `LEAVE UNLOOP EXIT`, `TID-X/Y/Z BID-X/Y/Z BDIM-X/Y/Z GDIM-X/Y/Z GLOBAL-ID` (GPU indexing).
- **Kernel Parameters**: Declared in the `\!` header. `\! kernel <name>` is required and must appear first. `\! param <name> i64[<N>]` becomes a `memref<Nxi64>` argument; `\! param <name> i64` becomes an `i64` argument. Using a param name in code emits `forth.param_ref` (arrays push address; scalars push value).
- **Shared Memory**: `\! shared <name> i64[<N>]` declares GPU shared (workgroup) memory. Emits a tagged `memref.alloca` at kernel entry; ForthToGPU converts it to a `gpu.func` workgroup attribution (`memref<Nxi64, #gpu.address_space<workgroup>>`). Using the shared name in code pushes its base address onto the stack. Cannot be referenced inside word definitions.
- **Shared Memory**: `\! shared <name> i64[<N>]` declares GPU shared (workgroup) memory. Emits a tagged `memref.alloca` at kernel entry; ForthToGPU converts it to a `gpu.func` workgroup attribution (`memref<Nxi64, #gpu.address_space<workgroup>>`). Using the shared name in code pushes its base address onto the stack. Use `S@`/`S!` for shared accesses. Cannot be referenced inside word definitions.
- **Conversion**: `!forth.stack` → `memref<256xi64>` with explicit stack pointer
- **GPU**: Functions wrapped in `gpu.module`, `main` gets `gpu.kernel` attribute, configured with bare pointers for NVVM conversion
- **User-defined Words**: Modeled as `func.func` with signature `(!forth.stack) -> !forth.stack`, called via `func.call`
Expand Down
62 changes: 62 additions & 0 deletions gpu_test/test_kernels.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -337,6 +337,68 @@ def test_naive_matmul_i64(kernel_runner: KernelRunner) -> None:
assert result == [12, 6, 9, 28, 14, 29]


def test_tiled_matmul_i64(kernel_runner: KernelRunner) -> None:
"""Tiled i64 matmul with shared memory: C = A(4x4) * B(4x4) -> C(4x4).

Uses 2x2 tiles, shared memory for A/B tiles, and BARRIER for sync.
Grid: (2,2,1), Block: (2,2,1) — 4 blocks of 4 threads each.
"""
result = kernel_runner.run(
forth_source=(
"\\! kernel main\n"
"\\! param A i64[16]\n"
"\\! param B i64[16]\n"
"\\! param C i64[16]\n"
"\\! shared SA i64[4]\n"
"\\! shared SB i64[4]\n"
"BID-Y 2 * TID-Y +\n"
"BID-X 2 * TID-X +\n"
"0\n"
"2 0 DO\n"
" 2 PICK 4 * I 2 * + TID-X + CELLS A + @\n"
" TID-Y 2 * TID-X + CELLS SA + S!\n"
" I 2 * TID-Y + 4 * 2 PICK + CELLS B + @\n"
" TID-Y 2 * TID-X + CELLS SB + S!\n"
" BARRIER\n"
" 2 0 DO\n"
" TID-Y 2 * I + CELLS SA + S@\n"
" I 2 * TID-X + CELLS SB + S@\n"
" * +\n"
" LOOP\n"
" BARRIER\n"
"LOOP\n"
"ROT 4 * ROT + CELLS C + !"
),
params={
"A": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
"B": [17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32],
},
grid=(2, 2, 1),
block=(2, 2, 1),
output_param=2,
output_count=16,
)
expected = [
250,
260,
270,
280,
618,
644,
670,
696,
986,
1028,
1070,
1112,
1354,
1412,
1470,
1528,
]
assert result == expected


# --- User-Defined Words ---


Expand Down
18 changes: 18 additions & 0 deletions include/warpforth/Dialect/Forth/ForthOps.td
Original file line numberDiff line numberDiff line change
Expand Up@@ -250,6 +250,24 @@ def Forth_StoreOp : Forth_StackOpBase<"store"> {
}];
}

def Forth_SharedLoadOp : Forth_StackOpBase<"shared_load"> {
let summary = "Load value from shared memory buffer";
let description = [{
Pops an address from the stack, loads a value from shared/workgroup memory at
that address, and pushes the loaded value onto the stack.
Forth semantics: ( addr -- value )
}];
}

def Forth_SharedStoreOp : Forth_StackOpBase<"shared_store"> {
let summary = "Store value to shared memory buffer";
let description = [{
Pops an address and value from the stack, stores the value to shared/workgroup
memory at the specified address.
Forth semantics: ( x addr -- )
}];
}

def Forth_ParamRefOp : Forth_Op<"param_ref", [Pure]> {
let summary = "Push kernel parameter address onto stack";
let description = [{
Expand Down
85 changes: 83 additions & 2 deletions lib/Conversion/ForthToMemRef/ForthToMemRef.cpp
Original file line numberDiff line numberDiff line change
Expand Up@@ -9,6 +9,7 @@
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Func/Transforms/FuncConversions.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/BuiltinTypes.h"
Expand All@@ -27,6 +28,8 @@ namespace {

// Stack configuration constants
constexpr int64_t kStackSize = 256;
constexpr unsigned kWorkgroupAddressSpace =
static_cast<unsigned>(gpu::AddressSpace::Workgroup);

/// Type converter for forth.stack -> memref + index
class ForthToMemRefTypeConverter : public TypeConverter {
Expand DownExpand Up@@ -725,6 +728,83 @@ struct StoreOpConversion : public OpConversionPattern<forth::StoreOp> {
}
};

/// Conversion pattern for forth.shared_load operation (S@).
/// Pops address from stack, loads value via shared/workgroup pointer, pushes
/// value: ( addr -- value )
struct SharedLoadOpConversion
: public OpConversionPattern<forth::SharedLoadOp> {
SharedLoadOpConversion(const TypeConverter &typeConverter,
MLIRContext *context)
: OpConversionPattern<forth::SharedLoadOp>(typeConverter, context) {}
using OneToNOpAdaptor = OpConversionPattern::OneToNOpAdaptor;

LogicalResult
matchAndRewrite(forth::SharedLoadOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
ValueRange inputStack = adaptor.getOperands()[0];
Value memref = inputStack[0];
Value stackPtr = inputStack[1];

auto i64Type = rewriter.getI64Type();
auto ptrType = LLVM::LLVMPointerType::get(rewriter.getContext(),
kWorkgroupAddressSpace);

// Load address from stack.
Value addrValue = rewriter.create<memref::LoadOp>(loc, memref, stackPtr);

// Load value from shared memory via address-space-qualified pointer.
Value ptr = rewriter.create<LLVM::IntToPtrOp>(loc, ptrType, addrValue);
Value loadedValue = rewriter.create<LLVM::LoadOp>(loc, i64Type, ptr);

// Store loaded value back at same position (replaces address).
rewriter.create<memref::StoreOp>(loc, loadedValue, memref, stackPtr);

rewriter.replaceOpWithMultiple(op, {{memref, stackPtr}});
return success();
}
};

/// Conversion pattern for forth.shared_store operation (S!).
/// Pops address and value from stack, stores value via shared/workgroup
/// pointer: ( x addr -- )
struct SharedStoreOpConversion
: public OpConversionPattern<forth::SharedStoreOp> {
SharedStoreOpConversion(const TypeConverter &typeConverter,
MLIRContext *context)
: OpConversionPattern<forth::SharedStoreOp>(typeConverter, context) {}
using OneToNOpAdaptor = OpConversionPattern::OneToNOpAdaptor;

LogicalResult
matchAndRewrite(forth::SharedStoreOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
ValueRange inputStack = adaptor.getOperands()[0];
Value memref = inputStack[0];
Value stackPtr = inputStack[1];

auto ptrType = LLVM::LLVMPointerType::get(rewriter.getContext(),
kWorkgroupAddressSpace);

// Pop address from stack.
Value addrValue = rewriter.create<memref::LoadOp>(loc, memref, stackPtr);

// Pop value from stack.
Value one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
Value spMinus1 = rewriter.create<arith::SubIOp>(loc, stackPtr, one);
Value value = rewriter.create<memref::LoadOp>(loc, memref, spMinus1);

// Store value to shared memory via address-space-qualified pointer.
Value ptr = rewriter.create<LLVM::IntToPtrOp>(loc, ptrType, addrValue);
rewriter.create<LLVM::StoreOp>(loc, value, ptr);

// New stack pointer is SP-2 (popped both address and value).
Value spMinus2 = rewriter.create<arith::SubIOp>(loc, spMinus1, one);
rewriter.replaceOpWithMultiple(op, {{memref, spMinus2}});
return success();
}
};

/// Template for converting GPU indexing ops to intrinsic ops.
/// Creates an intrinsic op with the specified name and pushes the value onto
/// the stack.
Expand DownExpand Up@@ -1026,8 +1106,9 @@ struct ConvertForthToMemRefPass
NotOpConversion, LshiftOpConversion, RshiftOpConversion, EqOpConversion,
LtOpConversion, GtOpConversion, NeOpConversion, LeOpConversion,
GeOpConversion, ZeroEqOpConversion, ParamRefOpConversion,
LoadOpConversion, StoreOpConversion, PopFlagOpConversion,
PopOpConversion, PushValueOpConversion>(typeConverter, context);
LoadOpConversion, StoreOpConversion, SharedLoadOpConversion,
SharedStoreOpConversion, PopFlagOpConversion, PopOpConversion,
PushValueOpConversion>(typeConverter, context);

// Add GPU indexing op conversion patterns
patterns.add<IntrinsicOpConversion<forth::ThreadIdXOp>>(typeConverter,
Expand Down
6 changes: 6 additions & 0 deletions lib/Translation/ForthToMLIR/ForthToMLIR.cpp
Original file line numberDiff line numberDiff line change
Expand Up@@ -487,6 +487,12 @@ Value ForthParser::emitOperation(StringRef word, Value inputStack,
} else if (word == "!") {
return builder.create<forth::StoreOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "S@") {
return builder.create<forth::SharedLoadOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "S!") {
return builder.create<forth::SharedStoreOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "TID-X") {
return builder.create<forth::ThreadIdXOp>(loc, stackType, inputStack)
.getResult();
Expand Down
13 changes: 13 additions & 0 deletions test/Conversion/ForthToMemRef/memory-ops.mlir
Original file line numberDiff line numberDiff line change
Expand Up@@ -15,6 +15,14 @@
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr
// CHECK: llvm.store %{{.*}}, %{{.*}} : i64, !llvm.ptr

// shared load (S@): pop address, inttoptr shared addrspace, llvm.load
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr<{{[1-9][0-9]*}}>
// CHECK: llvm.load %{{.*}} : !llvm.ptr<{{[1-9][0-9]*}}> -> i64

// shared store (S!): pop address + value, inttoptr shared addrspace, llvm.store
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr<{{[1-9][0-9]*}}>
// CHECK: llvm.store %{{.*}}, %{{.*}} : i64, !llvm.ptr<{{[1-9][0-9]*}}>

module {
func.func private @main() {
%0 = forth.stack !forth.stack
Expand All@@ -23,6 +31,11 @@ module {
%3 = forth.literal %2 42 : !forth.stack -> !forth.stack
%4 = forth.literal %3 100 : !forth.stack -> !forth.stack
%5 = forth.store %4 : !forth.stack -> !forth.stack
%6 = forth.literal %5 2 : !forth.stack -> !forth.stack
%7 = forth.shared_load %6 : !forth.stack -> !forth.stack
%8 = forth.literal %7 9 : !forth.stack -> !forth.stack
%9 = forth.literal %8 3 : !forth.stack -> !forth.stack
%10 = forth.shared_store %9 : !forth.stack -> !forth.stack
return
}
}
8 changes: 7 additions & 1 deletion test/Translation/Forth/memory-ops.forth
Original file line numberDiff line numberDiff line change
Expand Up@@ -6,9 +6,15 @@
\ Test ! produces forth.store
\ CHECK: forth.store %{{.*}} : !forth.stack -> !forth.stack

\ Test S@ produces forth.shared_load
\ CHECK: forth.shared_load %{{.*}} : !forth.stack -> !forth.stack

\ Test S! produces forth.shared_store
\ CHECK: forth.shared_store %{{.*}} : !forth.stack -> !forth.stack

\ Test CELLS produces literal 8 + mul
\ CHECK: forth.literal %{{.*}} 8
\ CHECK-NEXT: forth.mul
\! kernel main
1 @ 2 3 !
1 @ 2 3 ! 4 S@ 5 6 S!
4 CELLS
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix"); } } catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions CLAUDE.md
Original file line numberDiff line numberDiff line change
Expand Up@@ -87,9 +87,9 @@ uv run ruff format gpu_test/

- **Stack Type**: `!forth.stack` - untyped stack, programmer ensures type safety
- **Operations**: All take stack as input and produce stack as output (except `forth.stack`)
- **Supported Words**: literals, `DUP DROP SWAP OVER ROT NIP TUCK PICK ROLL`, `+ - * / MOD`, `AND OR XOR NOT LSHIFT RSHIFT`, `= < > <> <= >= 0=`, `@ !`, `CELLS`, `IF ELSE THEN`, `BEGIN UNTIL`, `BEGIN WHILE REPEAT`, `DO LOOP +LOOP I J K`, `LEAVE UNLOOP EXIT`, `TID-X/Y/Z BID-X/Y/Z BDIM-X/Y/Z GDIM-X/Y/Z GLOBAL-ID` (GPU indexing).
- **Supported Words**: literals, `DUP DROP SWAP OVER ROT NIP TUCK PICK ROLL`, `+ - * / MOD`, `AND OR XOR NOT LSHIFT RSHIFT`, `= < > <> <= >= 0=`, `@ !` (global memory), `S@ S!` (shared memory), `CELLS`, `IF ELSE THEN`, `BEGIN UNTIL`, `BEGIN WHILE REPEAT`, `DO LOOP +LOOP I J K`, `LEAVE UNLOOP EXIT`, `TID-X/Y/Z BID-X/Y/Z BDIM-X/Y/Z GDIM-X/Y/Z GLOBAL-ID` (GPU indexing).
- **Kernel Parameters**: Declared in the `\!` header. `\! kernel <name>` is required and must appear first. `\! param <name> i64[<N>]` becomes a `memref<Nxi64>` argument; `\! param <name> i64` becomes an `i64` argument. Using a param name in code emits `forth.param_ref` (arrays push address; scalars push value).
- **Shared Memory**: `\! shared <name> i64[<N>]` declares GPU shared (workgroup) memory. Emits a tagged `memref.alloca` at kernel entry; ForthToGPU converts it to a `gpu.func` workgroup attribution (`memref<Nxi64, #gpu.address_space<workgroup>>`). Using the shared name in code pushes its base address onto the stack. Cannot be referenced inside word definitions.
- **Shared Memory**: `\! shared <name> i64[<N>]` declares GPU shared (workgroup) memory. Emits a tagged `memref.alloca` at kernel entry; ForthToGPU converts it to a `gpu.func` workgroup attribution (`memref<Nxi64, #gpu.address_space<workgroup>>`). Using the shared name in code pushes its base address onto the stack. Use `S@`/`S!` for shared accesses. Cannot be referenced inside word definitions.
- **Conversion**: `!forth.stack` → `memref<256xi64>` with explicit stack pointer
- **GPU**: Functions wrapped in `gpu.module`, `main` gets `gpu.kernel` attribute, configured with bare pointers for NVVM conversion
- **User-defined Words**: Modeled as `func.func` with signature `(!forth.stack) -> !forth.stack`, called via `func.call`
Expand Down
62 changes: 62 additions & 0 deletions gpu_test/test_kernels.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -337,6 +337,68 @@ def test_naive_matmul_i64(kernel_runner: KernelRunner) -> None:
assert result == [12, 6, 9, 28, 14, 29]


def test_tiled_matmul_i64(kernel_runner: KernelRunner) -> None:
"""Tiled i64 matmul with shared memory: C = A(4x4) * B(4x4) -> C(4x4).

Uses 2x2 tiles, shared memory for A/B tiles, and BARRIER for sync.
Grid: (2,2,1), Block: (2,2,1) — 4 blocks of 4 threads each.
"""
result = kernel_runner.run(
forth_source=(
"\\! kernel main\n"
"\\! param A i64[16]\n"
"\\! param B i64[16]\n"
"\\! param C i64[16]\n"
"\\! shared SA i64[4]\n"
"\\! shared SB i64[4]\n"
"BID-Y 2 * TID-Y +\n"
"BID-X 2 * TID-X +\n"
"0\n"
"2 0 DO\n"
" 2 PICK 4 * I 2 * + TID-X + CELLS A + @\n"
" TID-Y 2 * TID-X + CELLS SA + S!\n"
" I 2 * TID-Y + 4 * 2 PICK + CELLS B + @\n"
" TID-Y 2 * TID-X + CELLS SB + S!\n"
" BARRIER\n"
" 2 0 DO\n"
" TID-Y 2 * I + CELLS SA + S@\n"
" I 2 * TID-X + CELLS SB + S@\n"
" * +\n"
" LOOP\n"
" BARRIER\n"
"LOOP\n"
"ROT 4 * ROT + CELLS C + !"
),
params={
"A": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
"B": [17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32],
},
grid=(2, 2, 1),
block=(2, 2, 1),
output_param=2,
output_count=16,
)
expected = [
250,
260,
270,
280,
618,
644,
670,
696,
986,
1028,
1070,
1112,
1354,
1412,
1470,
1528,
]
assert result == expected


# --- User-Defined Words ---


Expand Down
18 changes: 18 additions & 0 deletions include/warpforth/Dialect/Forth/ForthOps.td
Original file line numberDiff line numberDiff line change
Expand Up@@ -250,6 +250,24 @@ def Forth_StoreOp : Forth_StackOpBase<"store"> {
}];
}

def Forth_SharedLoadOp : Forth_StackOpBase<"shared_load"> {
let summary = "Load value from shared memory buffer";
let description = [{
Pops an address from the stack, loads a value from shared/workgroup memory at
that address, and pushes the loaded value onto the stack.
Forth semantics: ( addr -- value )
}];
}

def Forth_SharedStoreOp : Forth_StackOpBase<"shared_store"> {
let summary = "Store value to shared memory buffer";
let description = [{
Pops an address and value from the stack, stores the value to shared/workgroup
memory at the specified address.
Forth semantics: ( x addr -- )
}];
}

def Forth_ParamRefOp : Forth_Op<"param_ref", [Pure]> {
let summary = "Push kernel parameter address onto stack";
let description = [{
Expand Down
85 changes: 83 additions & 2 deletions lib/Conversion/ForthToMemRef/ForthToMemRef.cpp
Original file line numberDiff line numberDiff line change
Expand Up@@ -9,6 +9,7 @@
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Func/Transforms/FuncConversions.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/BuiltinTypes.h"
Expand All@@ -27,6 +28,8 @@ namespace {

// Stack configuration constants
constexpr int64_t kStackSize = 256;
constexpr unsigned kWorkgroupAddressSpace =
static_cast<unsigned>(gpu::AddressSpace::Workgroup);

/// Type converter for forth.stack -> memref + index
class ForthToMemRefTypeConverter : public TypeConverter {
Expand DownExpand Up@@ -725,6 +728,83 @@ struct StoreOpConversion : public OpConversionPattern<forth::StoreOp> {
}
};

/// Conversion pattern for forth.shared_load operation (S@).
/// Pops address from stack, loads value via shared/workgroup pointer, pushes
/// value: ( addr -- value )
struct SharedLoadOpConversion
: public OpConversionPattern<forth::SharedLoadOp> {
SharedLoadOpConversion(const TypeConverter &typeConverter,
MLIRContext *context)
: OpConversionPattern<forth::SharedLoadOp>(typeConverter, context) {}
using OneToNOpAdaptor = OpConversionPattern::OneToNOpAdaptor;

LogicalResult
matchAndRewrite(forth::SharedLoadOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
ValueRange inputStack = adaptor.getOperands()[0];
Value memref = inputStack[0];
Value stackPtr = inputStack[1];

auto i64Type = rewriter.getI64Type();
auto ptrType = LLVM::LLVMPointerType::get(rewriter.getContext(),
kWorkgroupAddressSpace);

// Load address from stack.
Value addrValue = rewriter.create<memref::LoadOp>(loc, memref, stackPtr);

// Load value from shared memory via address-space-qualified pointer.
Value ptr = rewriter.create<LLVM::IntToPtrOp>(loc, ptrType, addrValue);
Value loadedValue = rewriter.create<LLVM::LoadOp>(loc, i64Type, ptr);

// Store loaded value back at same position (replaces address).
rewriter.create<memref::StoreOp>(loc, loadedValue, memref, stackPtr);

rewriter.replaceOpWithMultiple(op, {{memref, stackPtr}});
return success();
}
};

/// Conversion pattern for forth.shared_store operation (S!).
/// Pops address and value from stack, stores value via shared/workgroup
/// pointer: ( x addr -- )
struct SharedStoreOpConversion
: public OpConversionPattern<forth::SharedStoreOp> {
SharedStoreOpConversion(const TypeConverter &typeConverter,
MLIRContext *context)
: OpConversionPattern<forth::SharedStoreOp>(typeConverter, context) {}
using OneToNOpAdaptor = OpConversionPattern::OneToNOpAdaptor;

LogicalResult
matchAndRewrite(forth::SharedStoreOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
ValueRange inputStack = adaptor.getOperands()[0];
Value memref = inputStack[0];
Value stackPtr = inputStack[1];

auto ptrType = LLVM::LLVMPointerType::get(rewriter.getContext(),
kWorkgroupAddressSpace);

// Pop address from stack.
Value addrValue = rewriter.create<memref::LoadOp>(loc, memref, stackPtr);

// Pop value from stack.
Value one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
Value spMinus1 = rewriter.create<arith::SubIOp>(loc, stackPtr, one);
Value value = rewriter.create<memref::LoadOp>(loc, memref, spMinus1);

// Store value to shared memory via address-space-qualified pointer.
Value ptr = rewriter.create<LLVM::IntToPtrOp>(loc, ptrType, addrValue);
rewriter.create<LLVM::StoreOp>(loc, value, ptr);

// New stack pointer is SP-2 (popped both address and value).
Value spMinus2 = rewriter.create<arith::SubIOp>(loc, spMinus1, one);
rewriter.replaceOpWithMultiple(op, {{memref, spMinus2}});
return success();
}
};

/// Template for converting GPU indexing ops to intrinsic ops.
/// Creates an intrinsic op with the specified name and pushes the value onto
/// the stack.
Expand DownExpand Up@@ -1026,8 +1106,9 @@ struct ConvertForthToMemRefPass
NotOpConversion, LshiftOpConversion, RshiftOpConversion, EqOpConversion,
LtOpConversion, GtOpConversion, NeOpConversion, LeOpConversion,
GeOpConversion, ZeroEqOpConversion, ParamRefOpConversion,
LoadOpConversion, StoreOpConversion, PopFlagOpConversion,
PopOpConversion, PushValueOpConversion>(typeConverter, context);
LoadOpConversion, StoreOpConversion, SharedLoadOpConversion,
SharedStoreOpConversion, PopFlagOpConversion, PopOpConversion,
PushValueOpConversion>(typeConverter, context);

// Add GPU indexing op conversion patterns
patterns.add<IntrinsicOpConversion<forth::ThreadIdXOp>>(typeConverter,
Expand Down
6 changes: 6 additions & 0 deletions lib/Translation/ForthToMLIR/ForthToMLIR.cpp
Original file line numberDiff line numberDiff line change
Expand Up@@ -487,6 +487,12 @@ Value ForthParser::emitOperation(StringRef word, Value inputStack,
} else if (word == "!") {
return builder.create<forth::StoreOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "S@") {
return builder.create<forth::SharedLoadOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "S!") {
return builder.create<forth::SharedStoreOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "TID-X") {
return builder.create<forth::ThreadIdXOp>(loc, stackType, inputStack)
.getResult();
Expand Down
13 changes: 13 additions & 0 deletions test/Conversion/ForthToMemRef/memory-ops.mlir
Original file line numberDiff line numberDiff line change
Expand Up@@ -15,6 +15,14 @@
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr
// CHECK: llvm.store %{{.*}}, %{{.*}} : i64, !llvm.ptr

// shared load (S@): pop address, inttoptr shared addrspace, llvm.load
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr<{{[1-9][0-9]*}}>
// CHECK: llvm.load %{{.*}} : !llvm.ptr<{{[1-9][0-9]*}}> -> i64

// shared store (S!): pop address + value, inttoptr shared addrspace, llvm.store
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr<{{[1-9][0-9]*}}>
// CHECK: llvm.store %{{.*}}, %{{.*}} : i64, !llvm.ptr<{{[1-9][0-9]*}}>

module {
func.func private @main() {
%0 = forth.stack !forth.stack
Expand All@@ -23,6 +31,11 @@ module {
%3 = forth.literal %2 42 : !forth.stack -> !forth.stack
%4 = forth.literal %3 100 : !forth.stack -> !forth.stack
%5 = forth.store %4 : !forth.stack -> !forth.stack
%6 = forth.literal %5 2 : !forth.stack -> !forth.stack
%7 = forth.shared_load %6 : !forth.stack -> !forth.stack
%8 = forth.literal %7 9 : !forth.stack -> !forth.stack
%9 = forth.literal %8 3 : !forth.stack -> !forth.stack
%10 = forth.shared_store %9 : !forth.stack -> !forth.stack
return
}
}
8 changes: 7 additions & 1 deletion test/Translation/Forth/memory-ops.forth
Original file line numberDiff line numberDiff line change
Expand Up@@ -6,9 +6,15 @@
\ Test ! produces forth.store
\ CHECK: forth.store %{{.*}} : !forth.stack -> !forth.stack

\ Test S@ produces forth.shared_load
\ CHECK: forth.shared_load %{{.*}} : !forth.stack -> !forth.stack

\ Test S! produces forth.shared_store
\ CHECK: forth.shared_store %{{.*}} : !forth.stack -> !forth.stack

\ Test CELLS produces literal 8 + mul
\ CHECK: forth.literal %{{.*}} 8
\ CHECK-NEXT: forth.mul
\! kernel main
1 @ 2 3 !
1 @ 2 3 ! 4 S@ 5 6 S!
4 CELLS
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms"); } } catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions CLAUDE.md
Original file line numberDiff line numberDiff line change
Expand Up@@ -87,9 +87,9 @@ uv run ruff format gpu_test/

- **Stack Type**: `!forth.stack` - untyped stack, programmer ensures type safety
- **Operations**: All take stack as input and produce stack as output (except `forth.stack`)
- **Supported Words**: literals, `DUP DROP SWAP OVER ROT NIP TUCK PICK ROLL`, `+ - * / MOD`, `AND OR XOR NOT LSHIFT RSHIFT`, `= < > <> <= >= 0=`, `@ !`, `CELLS`, `IF ELSE THEN`, `BEGIN UNTIL`, `BEGIN WHILE REPEAT`, `DO LOOP +LOOP I J K`, `LEAVE UNLOOP EXIT`, `TID-X/Y/Z BID-X/Y/Z BDIM-X/Y/Z GDIM-X/Y/Z GLOBAL-ID` (GPU indexing).
- **Supported Words**: literals, `DUP DROP SWAP OVER ROT NIP TUCK PICK ROLL`, `+ - * / MOD`, `AND OR XOR NOT LSHIFT RSHIFT`, `= < > <> <= >= 0=`, `@ !` (global memory), `S@ S!` (shared memory), `CELLS`, `IF ELSE THEN`, `BEGIN UNTIL`, `BEGIN WHILE REPEAT`, `DO LOOP +LOOP I J K`, `LEAVE UNLOOP EXIT`, `TID-X/Y/Z BID-X/Y/Z BDIM-X/Y/Z GDIM-X/Y/Z GLOBAL-ID` (GPU indexing).
- **Kernel Parameters**: Declared in the `\!` header. `\! kernel <name>` is required and must appear first. `\! param <name> i64[<N>]` becomes a `memref<Nxi64>` argument; `\! param <name> i64` becomes an `i64` argument. Using a param name in code emits `forth.param_ref` (arrays push address; scalars push value).
- **Shared Memory**: `\! shared <name> i64[<N>]` declares GPU shared (workgroup) memory. Emits a tagged `memref.alloca` at kernel entry; ForthToGPU converts it to a `gpu.func` workgroup attribution (`memref<Nxi64, #gpu.address_space<workgroup>>`). Using the shared name in code pushes its base address onto the stack. Cannot be referenced inside word definitions.
- **Shared Memory**: `\! shared <name> i64[<N>]` declares GPU shared (workgroup) memory. Emits a tagged `memref.alloca` at kernel entry; ForthToGPU converts it to a `gpu.func` workgroup attribution (`memref<Nxi64, #gpu.address_space<workgroup>>`). Using the shared name in code pushes its base address onto the stack. Use `S@`/`S!` for shared accesses. Cannot be referenced inside word definitions.
- **Conversion**: `!forth.stack` → `memref<256xi64>` with explicit stack pointer
- **GPU**: Functions wrapped in `gpu.module`, `main` gets `gpu.kernel` attribute, configured with bare pointers for NVVM conversion
- **User-defined Words**: Modeled as `func.func` with signature `(!forth.stack) -> !forth.stack`, called via `func.call`
Expand Down
62 changes: 62 additions & 0 deletions gpu_test/test_kernels.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -337,6 +337,68 @@ def test_naive_matmul_i64(kernel_runner: KernelRunner) -> None:
assert result == [12, 6, 9, 28, 14, 29]


def test_tiled_matmul_i64(kernel_runner: KernelRunner) -> None:
"""Tiled i64 matmul with shared memory: C = A(4x4) * B(4x4) -> C(4x4).

Uses 2x2 tiles, shared memory for A/B tiles, and BARRIER for sync.
Grid: (2,2,1), Block: (2,2,1) — 4 blocks of 4 threads each.
"""
result = kernel_runner.run(
forth_source=(
"\\! kernel main\n"
"\\! param A i64[16]\n"
"\\! param B i64[16]\n"
"\\! param C i64[16]\n"
"\\! shared SA i64[4]\n"
"\\! shared SB i64[4]\n"
"BID-Y 2 * TID-Y +\n"
"BID-X 2 * TID-X +\n"
"0\n"
"2 0 DO\n"
" 2 PICK 4 * I 2 * + TID-X + CELLS A + @\n"
" TID-Y 2 * TID-X + CELLS SA + S!\n"
" I 2 * TID-Y + 4 * 2 PICK + CELLS B + @\n"
" TID-Y 2 * TID-X + CELLS SB + S!\n"
" BARRIER\n"
" 2 0 DO\n"
" TID-Y 2 * I + CELLS SA + S@\n"
" I 2 * TID-X + CELLS SB + S@\n"
" * +\n"
" LOOP\n"
" BARRIER\n"
"LOOP\n"
"ROT 4 * ROT + CELLS C + !"
),
params={
"A": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
"B": [17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32],
},
grid=(2, 2, 1),
block=(2, 2, 1),
output_param=2,
output_count=16,
)
expected = [
250,
260,
270,
280,
618,
644,
670,
696,
986,
1028,
1070,
1112,
1354,
1412,
1470,
1528,
]
assert result == expected


# --- User-Defined Words ---


Expand Down
18 changes: 18 additions & 0 deletions include/warpforth/Dialect/Forth/ForthOps.td
Original file line numberDiff line numberDiff line change
Expand Up@@ -250,6 +250,24 @@ def Forth_StoreOp : Forth_StackOpBase<"store"> {
}];
}

def Forth_SharedLoadOp : Forth_StackOpBase<"shared_load"> {
let summary = "Load value from shared memory buffer";
let description = [{
Pops an address from the stack, loads a value from shared/workgroup memory at
that address, and pushes the loaded value onto the stack.
Forth semantics: ( addr -- value )
}];
}

def Forth_SharedStoreOp : Forth_StackOpBase<"shared_store"> {
let summary = "Store value to shared memory buffer";
let description = [{
Pops an address and value from the stack, stores the value to shared/workgroup
memory at the specified address.
Forth semantics: ( x addr -- )
}];
}

def Forth_ParamRefOp : Forth_Op<"param_ref", [Pure]> {
let summary = "Push kernel parameter address onto stack";
let description = [{
Expand Down
85 changes: 83 additions & 2 deletions lib/Conversion/ForthToMemRef/ForthToMemRef.cpp
Original file line numberDiff line numberDiff line change
Expand Up@@ -9,6 +9,7 @@
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Func/Transforms/FuncConversions.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/BuiltinTypes.h"
Expand All@@ -27,6 +28,8 @@ namespace {

// Stack configuration constants
constexpr int64_t kStackSize = 256;
constexpr unsigned kWorkgroupAddressSpace =
static_cast<unsigned>(gpu::AddressSpace::Workgroup);

/// Type converter for forth.stack -> memref + index
class ForthToMemRefTypeConverter : public TypeConverter {
Expand DownExpand Up@@ -725,6 +728,83 @@ struct StoreOpConversion : public OpConversionPattern<forth::StoreOp> {
}
};

/// Conversion pattern for forth.shared_load operation (S@).
/// Pops address from stack, loads value via shared/workgroup pointer, pushes
/// value: ( addr -- value )
struct SharedLoadOpConversion
: public OpConversionPattern<forth::SharedLoadOp> {
SharedLoadOpConversion(const TypeConverter &typeConverter,
MLIRContext *context)
: OpConversionPattern<forth::SharedLoadOp>(typeConverter, context) {}
using OneToNOpAdaptor = OpConversionPattern::OneToNOpAdaptor;

LogicalResult
matchAndRewrite(forth::SharedLoadOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
ValueRange inputStack = adaptor.getOperands()[0];
Value memref = inputStack[0];
Value stackPtr = inputStack[1];

auto i64Type = rewriter.getI64Type();
auto ptrType = LLVM::LLVMPointerType::get(rewriter.getContext(),
kWorkgroupAddressSpace);

// Load address from stack.
Value addrValue = rewriter.create<memref::LoadOp>(loc, memref, stackPtr);

// Load value from shared memory via address-space-qualified pointer.
Value ptr = rewriter.create<LLVM::IntToPtrOp>(loc, ptrType, addrValue);
Value loadedValue = rewriter.create<LLVM::LoadOp>(loc, i64Type, ptr);

// Store loaded value back at same position (replaces address).
rewriter.create<memref::StoreOp>(loc, loadedValue, memref, stackPtr);

rewriter.replaceOpWithMultiple(op, {{memref, stackPtr}});
return success();
}
};

/// Conversion pattern for forth.shared_store operation (S!).
/// Pops address and value from stack, stores value via shared/workgroup
/// pointer: ( x addr -- )
struct SharedStoreOpConversion
: public OpConversionPattern<forth::SharedStoreOp> {
SharedStoreOpConversion(const TypeConverter &typeConverter,
MLIRContext *context)
: OpConversionPattern<forth::SharedStoreOp>(typeConverter, context) {}
using OneToNOpAdaptor = OpConversionPattern::OneToNOpAdaptor;

LogicalResult
matchAndRewrite(forth::SharedStoreOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
ValueRange inputStack = adaptor.getOperands()[0];
Value memref = inputStack[0];
Value stackPtr = inputStack[1];

auto ptrType = LLVM::LLVMPointerType::get(rewriter.getContext(),
kWorkgroupAddressSpace);

// Pop address from stack.
Value addrValue = rewriter.create<memref::LoadOp>(loc, memref, stackPtr);

// Pop value from stack.
Value one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
Value spMinus1 = rewriter.create<arith::SubIOp>(loc, stackPtr, one);
Value value = rewriter.create<memref::LoadOp>(loc, memref, spMinus1);

// Store value to shared memory via address-space-qualified pointer.
Value ptr = rewriter.create<LLVM::IntToPtrOp>(loc, ptrType, addrValue);
rewriter.create<LLVM::StoreOp>(loc, value, ptr);

// New stack pointer is SP-2 (popped both address and value).
Value spMinus2 = rewriter.create<arith::SubIOp>(loc, spMinus1, one);
rewriter.replaceOpWithMultiple(op, {{memref, spMinus2}});
return success();
}
};

/// Template for converting GPU indexing ops to intrinsic ops.
/// Creates an intrinsic op with the specified name and pushes the value onto
/// the stack.
Expand DownExpand Up@@ -1026,8 +1106,9 @@ struct ConvertForthToMemRefPass
NotOpConversion, LshiftOpConversion, RshiftOpConversion, EqOpConversion,
LtOpConversion, GtOpConversion, NeOpConversion, LeOpConversion,
GeOpConversion, ZeroEqOpConversion, ParamRefOpConversion,
LoadOpConversion, StoreOpConversion, PopFlagOpConversion,
PopOpConversion, PushValueOpConversion>(typeConverter, context);
LoadOpConversion, StoreOpConversion, SharedLoadOpConversion,
SharedStoreOpConversion, PopFlagOpConversion, PopOpConversion,
PushValueOpConversion>(typeConverter, context);

// Add GPU indexing op conversion patterns
patterns.add<IntrinsicOpConversion<forth::ThreadIdXOp>>(typeConverter,
Expand Down
6 changes: 6 additions & 0 deletions lib/Translation/ForthToMLIR/ForthToMLIR.cpp
Original file line numberDiff line numberDiff line change
Expand Up@@ -487,6 +487,12 @@ Value ForthParser::emitOperation(StringRef word, Value inputStack,
} else if (word == "!") {
return builder.create<forth::StoreOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "S@") {
return builder.create<forth::SharedLoadOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "S!") {
return builder.create<forth::SharedStoreOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "TID-X") {
return builder.create<forth::ThreadIdXOp>(loc, stackType, inputStack)
.getResult();
Expand Down
13 changes: 13 additions & 0 deletions test/Conversion/ForthToMemRef/memory-ops.mlir
Original file line numberDiff line numberDiff line change
Expand Up@@ -15,6 +15,14 @@
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr
// CHECK: llvm.store %{{.*}}, %{{.*}} : i64, !llvm.ptr

// shared load (S@): pop address, inttoptr shared addrspace, llvm.load
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr<{{[1-9][0-9]*}}>
// CHECK: llvm.load %{{.*}} : !llvm.ptr<{{[1-9][0-9]*}}> -> i64

// shared store (S!): pop address + value, inttoptr shared addrspace, llvm.store
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr<{{[1-9][0-9]*}}>
// CHECK: llvm.store %{{.*}}, %{{.*}} : i64, !llvm.ptr<{{[1-9][0-9]*}}>

module {
func.func private @main() {
%0 = forth.stack !forth.stack
Expand All@@ -23,6 +31,11 @@ module {
%3 = forth.literal %2 42 : !forth.stack -> !forth.stack
%4 = forth.literal %3 100 : !forth.stack -> !forth.stack
%5 = forth.store %4 : !forth.stack -> !forth.stack
%6 = forth.literal %5 2 : !forth.stack -> !forth.stack
%7 = forth.shared_load %6 : !forth.stack -> !forth.stack
%8 = forth.literal %7 9 : !forth.stack -> !forth.stack
%9 = forth.literal %8 3 : !forth.stack -> !forth.stack
%10 = forth.shared_store %9 : !forth.stack -> !forth.stack
return
}
}
8 changes: 7 additions & 1 deletion test/Translation/Forth/memory-ops.forth
Original file line numberDiff line numberDiff line change
Expand Up@@ -6,9 +6,15 @@
\ Test ! produces forth.store
\ CHECK: forth.store %{{.*}} : !forth.stack -> !forth.stack

\ Test S@ produces forth.shared_load
\ CHECK: forth.shared_load %{{.*}} : !forth.stack -> !forth.stack

\ Test S! produces forth.shared_store
\ CHECK: forth.shared_store %{{.*}} : !forth.stack -> !forth.stack

\ Test CELLS produces literal 8 + mul
\ CHECK: forth.literal %{{.*}} 8
\ CHECK-NEXT: forth.mul
\! kernel main
1 @ 2 3 !
1 @ 2 3 ! 4 S@ 5 6 S!
4 CELLS
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links"); } } catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); } })(); (function(){ try { var __m = "youtube.com"; var __re = new RegExp('^' + "youtube\\.com" + '
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions CLAUDE.md
Original file line numberDiff line numberDiff line change
Expand Up@@ -87,9 +87,9 @@ uv run ruff format gpu_test/

- **Stack Type**: `!forth.stack` - untyped stack, programmer ensures type safety
- **Operations**: All take stack as input and produce stack as output (except `forth.stack`)
- **Supported Words**: literals, `DUP DROP SWAP OVER ROT NIP TUCK PICK ROLL`, `+ - * / MOD`, `AND OR XOR NOT LSHIFT RSHIFT`, `= < > <> <= >= 0=`, `@ !`, `CELLS`, `IF ELSE THEN`, `BEGIN UNTIL`, `BEGIN WHILE REPEAT`, `DO LOOP +LOOP I J K`, `LEAVE UNLOOP EXIT`, `TID-X/Y/Z BID-X/Y/Z BDIM-X/Y/Z GDIM-X/Y/Z GLOBAL-ID` (GPU indexing).
- **Supported Words**: literals, `DUP DROP SWAP OVER ROT NIP TUCK PICK ROLL`, `+ - * / MOD`, `AND OR XOR NOT LSHIFT RSHIFT`, `= < > <> <= >= 0=`, `@ !` (global memory), `S@ S!` (shared memory), `CELLS`, `IF ELSE THEN`, `BEGIN UNTIL`, `BEGIN WHILE REPEAT`, `DO LOOP +LOOP I J K`, `LEAVE UNLOOP EXIT`, `TID-X/Y/Z BID-X/Y/Z BDIM-X/Y/Z GDIM-X/Y/Z GLOBAL-ID` (GPU indexing).
- **Kernel Parameters**: Declared in the `\!` header. `\! kernel <name>` is required and must appear first. `\! param <name> i64[<N>]` becomes a `memref<Nxi64>` argument; `\! param <name> i64` becomes an `i64` argument. Using a param name in code emits `forth.param_ref` (arrays push address; scalars push value).
- **Shared Memory**: `\! shared <name> i64[<N>]` declares GPU shared (workgroup) memory. Emits a tagged `memref.alloca` at kernel entry; ForthToGPU converts it to a `gpu.func` workgroup attribution (`memref<Nxi64, #gpu.address_space<workgroup>>`). Using the shared name in code pushes its base address onto the stack. Cannot be referenced inside word definitions.
- **Shared Memory**: `\! shared <name> i64[<N>]` declares GPU shared (workgroup) memory. Emits a tagged `memref.alloca` at kernel entry; ForthToGPU converts it to a `gpu.func` workgroup attribution (`memref<Nxi64, #gpu.address_space<workgroup>>`). Using the shared name in code pushes its base address onto the stack. Use `S@`/`S!` for shared accesses. Cannot be referenced inside word definitions.
- **Conversion**: `!forth.stack` → `memref<256xi64>` with explicit stack pointer
- **GPU**: Functions wrapped in `gpu.module`, `main` gets `gpu.kernel` attribute, configured with bare pointers for NVVM conversion
- **User-defined Words**: Modeled as `func.func` with signature `(!forth.stack) -> !forth.stack`, called via `func.call`
Expand Down
62 changes: 62 additions & 0 deletions gpu_test/test_kernels.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -337,6 +337,68 @@ def test_naive_matmul_i64(kernel_runner: KernelRunner) -> None:
assert result == [12, 6, 9, 28, 14, 29]


def test_tiled_matmul_i64(kernel_runner: KernelRunner) -> None:
"""Tiled i64 matmul with shared memory: C = A(4x4) * B(4x4) -> C(4x4).

Uses 2x2 tiles, shared memory for A/B tiles, and BARRIER for sync.
Grid: (2,2,1), Block: (2,2,1) — 4 blocks of 4 threads each.
"""
result = kernel_runner.run(
forth_source=(
"\\! kernel main\n"
"\\! param A i64[16]\n"
"\\! param B i64[16]\n"
"\\! param C i64[16]\n"
"\\! shared SA i64[4]\n"
"\\! shared SB i64[4]\n"
"BID-Y 2 * TID-Y +\n"
"BID-X 2 * TID-X +\n"
"0\n"
"2 0 DO\n"
" 2 PICK 4 * I 2 * + TID-X + CELLS A + @\n"
" TID-Y 2 * TID-X + CELLS SA + S!\n"
" I 2 * TID-Y + 4 * 2 PICK + CELLS B + @\n"
" TID-Y 2 * TID-X + CELLS SB + S!\n"
" BARRIER\n"
" 2 0 DO\n"
" TID-Y 2 * I + CELLS SA + S@\n"
" I 2 * TID-X + CELLS SB + S@\n"
" * +\n"
" LOOP\n"
" BARRIER\n"
"LOOP\n"
"ROT 4 * ROT + CELLS C + !"
),
params={
"A": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
"B": [17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32],
},
grid=(2, 2, 1),
block=(2, 2, 1),
output_param=2,
output_count=16,
)
expected = [
250,
260,
270,
280,
618,
644,
670,
696,
986,
1028,
1070,
1112,
1354,
1412,
1470,
1528,
]
assert result == expected


# --- User-Defined Words ---


Expand Down
18 changes: 18 additions & 0 deletions include/warpforth/Dialect/Forth/ForthOps.td
Original file line numberDiff line numberDiff line change
Expand Up@@ -250,6 +250,24 @@ def Forth_StoreOp : Forth_StackOpBase<"store"> {
}];
}

def Forth_SharedLoadOp : Forth_StackOpBase<"shared_load"> {
let summary = "Load value from shared memory buffer";
let description = [{
Pops an address from the stack, loads a value from shared/workgroup memory at
that address, and pushes the loaded value onto the stack.
Forth semantics: ( addr -- value )
}];
}

def Forth_SharedStoreOp : Forth_StackOpBase<"shared_store"> {
let summary = "Store value to shared memory buffer";
let description = [{
Pops an address and value from the stack, stores the value to shared/workgroup
memory at the specified address.
Forth semantics: ( x addr -- )
}];
}

def Forth_ParamRefOp : Forth_Op<"param_ref", [Pure]> {
let summary = "Push kernel parameter address onto stack";
let description = [{
Expand Down
85 changes: 83 additions & 2 deletions lib/Conversion/ForthToMemRef/ForthToMemRef.cpp
Original file line numberDiff line numberDiff line change
Expand Up@@ -9,6 +9,7 @@
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Func/Transforms/FuncConversions.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/BuiltinTypes.h"
Expand All@@ -27,6 +28,8 @@ namespace {

// Stack configuration constants
constexpr int64_t kStackSize = 256;
constexpr unsigned kWorkgroupAddressSpace =
static_cast<unsigned>(gpu::AddressSpace::Workgroup);

/// Type converter for forth.stack -> memref + index
class ForthToMemRefTypeConverter : public TypeConverter {
Expand DownExpand Up@@ -725,6 +728,83 @@ struct StoreOpConversion : public OpConversionPattern<forth::StoreOp> {
}
};

/// Conversion pattern for forth.shared_load operation (S@).
/// Pops address from stack, loads value via shared/workgroup pointer, pushes
/// value: ( addr -- value )
struct SharedLoadOpConversion
: public OpConversionPattern<forth::SharedLoadOp> {
SharedLoadOpConversion(const TypeConverter &typeConverter,
MLIRContext *context)
: OpConversionPattern<forth::SharedLoadOp>(typeConverter, context) {}
using OneToNOpAdaptor = OpConversionPattern::OneToNOpAdaptor;

LogicalResult
matchAndRewrite(forth::SharedLoadOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
ValueRange inputStack = adaptor.getOperands()[0];
Value memref = inputStack[0];
Value stackPtr = inputStack[1];

auto i64Type = rewriter.getI64Type();
auto ptrType = LLVM::LLVMPointerType::get(rewriter.getContext(),
kWorkgroupAddressSpace);

// Load address from stack.
Value addrValue = rewriter.create<memref::LoadOp>(loc, memref, stackPtr);

// Load value from shared memory via address-space-qualified pointer.
Value ptr = rewriter.create<LLVM::IntToPtrOp>(loc, ptrType, addrValue);
Value loadedValue = rewriter.create<LLVM::LoadOp>(loc, i64Type, ptr);

// Store loaded value back at same position (replaces address).
rewriter.create<memref::StoreOp>(loc, loadedValue, memref, stackPtr);

rewriter.replaceOpWithMultiple(op, {{memref, stackPtr}});
return success();
}
};

/// Conversion pattern for forth.shared_store operation (S!).
/// Pops address and value from stack, stores value via shared/workgroup
/// pointer: ( x addr -- )
struct SharedStoreOpConversion
: public OpConversionPattern<forth::SharedStoreOp> {
SharedStoreOpConversion(const TypeConverter &typeConverter,
MLIRContext *context)
: OpConversionPattern<forth::SharedStoreOp>(typeConverter, context) {}
using OneToNOpAdaptor = OpConversionPattern::OneToNOpAdaptor;

LogicalResult
matchAndRewrite(forth::SharedStoreOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
ValueRange inputStack = adaptor.getOperands()[0];
Value memref = inputStack[0];
Value stackPtr = inputStack[1];

auto ptrType = LLVM::LLVMPointerType::get(rewriter.getContext(),
kWorkgroupAddressSpace);

// Pop address from stack.
Value addrValue = rewriter.create<memref::LoadOp>(loc, memref, stackPtr);

// Pop value from stack.
Value one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
Value spMinus1 = rewriter.create<arith::SubIOp>(loc, stackPtr, one);
Value value = rewriter.create<memref::LoadOp>(loc, memref, spMinus1);

// Store value to shared memory via address-space-qualified pointer.
Value ptr = rewriter.create<LLVM::IntToPtrOp>(loc, ptrType, addrValue);
rewriter.create<LLVM::StoreOp>(loc, value, ptr);

// New stack pointer is SP-2 (popped both address and value).
Value spMinus2 = rewriter.create<arith::SubIOp>(loc, spMinus1, one);
rewriter.replaceOpWithMultiple(op, {{memref, spMinus2}});
return success();
}
};

/// Template for converting GPU indexing ops to intrinsic ops.
/// Creates an intrinsic op with the specified name and pushes the value onto
/// the stack.
Expand DownExpand Up@@ -1026,8 +1106,9 @@ struct ConvertForthToMemRefPass
NotOpConversion, LshiftOpConversion, RshiftOpConversion, EqOpConversion,
LtOpConversion, GtOpConversion, NeOpConversion, LeOpConversion,
GeOpConversion, ZeroEqOpConversion, ParamRefOpConversion,
LoadOpConversion, StoreOpConversion, PopFlagOpConversion,
PopOpConversion, PushValueOpConversion>(typeConverter, context);
LoadOpConversion, StoreOpConversion, SharedLoadOpConversion,
SharedStoreOpConversion, PopFlagOpConversion, PopOpConversion,
PushValueOpConversion>(typeConverter, context);

// Add GPU indexing op conversion patterns
patterns.add<IntrinsicOpConversion<forth::ThreadIdXOp>>(typeConverter,
Expand Down
6 changes: 6 additions & 0 deletions lib/Translation/ForthToMLIR/ForthToMLIR.cpp
Original file line numberDiff line numberDiff line change
Expand Up@@ -487,6 +487,12 @@ Value ForthParser::emitOperation(StringRef word, Value inputStack,
} else if (word == "!") {
return builder.create<forth::StoreOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "S@") {
return builder.create<forth::SharedLoadOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "S!") {
return builder.create<forth::SharedStoreOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "TID-X") {
return builder.create<forth::ThreadIdXOp>(loc, stackType, inputStack)
.getResult();
Expand Down
13 changes: 13 additions & 0 deletions test/Conversion/ForthToMemRef/memory-ops.mlir
Original file line numberDiff line numberDiff line change
Expand Up@@ -15,6 +15,14 @@
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr
// CHECK: llvm.store %{{.*}}, %{{.*}} : i64, !llvm.ptr

// shared load (S@): pop address, inttoptr shared addrspace, llvm.load
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr<{{[1-9][0-9]*}}>
// CHECK: llvm.load %{{.*}} : !llvm.ptr<{{[1-9][0-9]*}}> -> i64

// shared store (S!): pop address + value, inttoptr shared addrspace, llvm.store
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr<{{[1-9][0-9]*}}>
// CHECK: llvm.store %{{.*}}, %{{.*}} : i64, !llvm.ptr<{{[1-9][0-9]*}}>

module {
func.func private @main() {
%0 = forth.stack !forth.stack
Expand All@@ -23,6 +31,11 @@ module {
%3 = forth.literal %2 42 : !forth.stack -> !forth.stack
%4 = forth.literal %3 100 : !forth.stack -> !forth.stack
%5 = forth.store %4 : !forth.stack -> !forth.stack
%6 = forth.literal %5 2 : !forth.stack -> !forth.stack
%7 = forth.shared_load %6 : !forth.stack -> !forth.stack
%8 = forth.literal %7 9 : !forth.stack -> !forth.stack
%9 = forth.literal %8 3 : !forth.stack -> !forth.stack
%10 = forth.shared_store %9 : !forth.stack -> !forth.stack
return
}
}
8 changes: 7 additions & 1 deletion test/Translation/Forth/memory-ops.forth
Original file line numberDiff line numberDiff line change
Expand Up@@ -6,9 +6,15 @@
\ Test ! produces forth.store
\ CHECK: forth.store %{{.*}} : !forth.stack -> !forth.stack

\ Test S@ produces forth.shared_load
\ CHECK: forth.shared_load %{{.*}} : !forth.stack -> !forth.stack

\ Test S! produces forth.shared_store
\ CHECK: forth.shared_store %{{.*}} : !forth.stack -> !forth.stack

\ Test CELLS produces literal 8 + mul
\ CHECK: forth.literal %{{.*}} 8
\ CHECK-NEXT: forth.mul
\! kernel main
1 @ 2 3 !
1 @ 2 3 ! 4 S@ 5 6 S!
4 CELLS
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default"); } } catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions CLAUDE.md
Original file line numberDiff line numberDiff line change
Expand Up@@ -87,9 +87,9 @@ uv run ruff format gpu_test/

- **Stack Type**: `!forth.stack` - untyped stack, programmer ensures type safety
- **Operations**: All take stack as input and produce stack as output (except `forth.stack`)
- **Supported Words**: literals, `DUP DROP SWAP OVER ROT NIP TUCK PICK ROLL`, `+ - * / MOD`, `AND OR XOR NOT LSHIFT RSHIFT`, `= < > <> <= >= 0=`, `@ !`, `CELLS`, `IF ELSE THEN`, `BEGIN UNTIL`, `BEGIN WHILE REPEAT`, `DO LOOP +LOOP I J K`, `LEAVE UNLOOP EXIT`, `TID-X/Y/Z BID-X/Y/Z BDIM-X/Y/Z GDIM-X/Y/Z GLOBAL-ID` (GPU indexing).
- **Supported Words**: literals, `DUP DROP SWAP OVER ROT NIP TUCK PICK ROLL`, `+ - * / MOD`, `AND OR XOR NOT LSHIFT RSHIFT`, `= < > <> <= >= 0=`, `@ !` (global memory), `S@ S!` (shared memory), `CELLS`, `IF ELSE THEN`, `BEGIN UNTIL`, `BEGIN WHILE REPEAT`, `DO LOOP +LOOP I J K`, `LEAVE UNLOOP EXIT`, `TID-X/Y/Z BID-X/Y/Z BDIM-X/Y/Z GDIM-X/Y/Z GLOBAL-ID` (GPU indexing).
- **Kernel Parameters**: Declared in the `\!` header. `\! kernel <name>` is required and must appear first. `\! param <name> i64[<N>]` becomes a `memref<Nxi64>` argument; `\! param <name> i64` becomes an `i64` argument. Using a param name in code emits `forth.param_ref` (arrays push address; scalars push value).
- **Shared Memory**: `\! shared <name> i64[<N>]` declares GPU shared (workgroup) memory. Emits a tagged `memref.alloca` at kernel entry; ForthToGPU converts it to a `gpu.func` workgroup attribution (`memref<Nxi64, #gpu.address_space<workgroup>>`). Using the shared name in code pushes its base address onto the stack. Cannot be referenced inside word definitions.
- **Shared Memory**: `\! shared <name> i64[<N>]` declares GPU shared (workgroup) memory. Emits a tagged `memref.alloca` at kernel entry; ForthToGPU converts it to a `gpu.func` workgroup attribution (`memref<Nxi64, #gpu.address_space<workgroup>>`). Using the shared name in code pushes its base address onto the stack. Use `S@`/`S!` for shared accesses. Cannot be referenced inside word definitions.
- **Conversion**: `!forth.stack` → `memref<256xi64>` with explicit stack pointer
- **GPU**: Functions wrapped in `gpu.module`, `main` gets `gpu.kernel` attribute, configured with bare pointers for NVVM conversion
- **User-defined Words**: Modeled as `func.func` with signature `(!forth.stack) -> !forth.stack`, called via `func.call`
Expand Down
62 changes: 62 additions & 0 deletions gpu_test/test_kernels.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -337,6 +337,68 @@ def test_naive_matmul_i64(kernel_runner: KernelRunner) -> None:
assert result == [12, 6, 9, 28, 14, 29]


def test_tiled_matmul_i64(kernel_runner: KernelRunner) -> None:
"""Tiled i64 matmul with shared memory: C = A(4x4) * B(4x4) -> C(4x4).

Uses 2x2 tiles, shared memory for A/B tiles, and BARRIER for sync.
Grid: (2,2,1), Block: (2,2,1) — 4 blocks of 4 threads each.
"""
result = kernel_runner.run(
forth_source=(
"\\! kernel main\n"
"\\! param A i64[16]\n"
"\\! param B i64[16]\n"
"\\! param C i64[16]\n"
"\\! shared SA i64[4]\n"
"\\! shared SB i64[4]\n"
"BID-Y 2 * TID-Y +\n"
"BID-X 2 * TID-X +\n"
"0\n"
"2 0 DO\n"
" 2 PICK 4 * I 2 * + TID-X + CELLS A + @\n"
" TID-Y 2 * TID-X + CELLS SA + S!\n"
" I 2 * TID-Y + 4 * 2 PICK + CELLS B + @\n"
" TID-Y 2 * TID-X + CELLS SB + S!\n"
" BARRIER\n"
" 2 0 DO\n"
" TID-Y 2 * I + CELLS SA + S@\n"
" I 2 * TID-X + CELLS SB + S@\n"
" * +\n"
" LOOP\n"
" BARRIER\n"
"LOOP\n"
"ROT 4 * ROT + CELLS C + !"
),
params={
"A": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
"B": [17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32],
},
grid=(2, 2, 1),
block=(2, 2, 1),
output_param=2,
output_count=16,
)
expected = [
250,
260,
270,
280,
618,
644,
670,
696,
986,
1028,
1070,
1112,
1354,
1412,
1470,
1528,
]
assert result == expected


# --- User-Defined Words ---


Expand Down
18 changes: 18 additions & 0 deletions include/warpforth/Dialect/Forth/ForthOps.td
Original file line numberDiff line numberDiff line change
Expand Up@@ -250,6 +250,24 @@ def Forth_StoreOp : Forth_StackOpBase<"store"> {
}];
}

def Forth_SharedLoadOp : Forth_StackOpBase<"shared_load"> {
let summary = "Load value from shared memory buffer";
let description = [{
Pops an address from the stack, loads a value from shared/workgroup memory at
that address, and pushes the loaded value onto the stack.
Forth semantics: ( addr -- value )
}];
}

def Forth_SharedStoreOp : Forth_StackOpBase<"shared_store"> {
let summary = "Store value to shared memory buffer";
let description = [{
Pops an address and value from the stack, stores the value to shared/workgroup
memory at the specified address.
Forth semantics: ( x addr -- )
}];
}

def Forth_ParamRefOp : Forth_Op<"param_ref", [Pure]> {
let summary = "Push kernel parameter address onto stack";
let description = [{
Expand Down
85 changes: 83 additions & 2 deletions lib/Conversion/ForthToMemRef/ForthToMemRef.cpp
Original file line numberDiff line numberDiff line change
Expand Up@@ -9,6 +9,7 @@
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Func/Transforms/FuncConversions.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/BuiltinTypes.h"
Expand All@@ -27,6 +28,8 @@ namespace {

// Stack configuration constants
constexpr int64_t kStackSize = 256;
constexpr unsigned kWorkgroupAddressSpace =
static_cast<unsigned>(gpu::AddressSpace::Workgroup);

/// Type converter for forth.stack -> memref + index
class ForthToMemRefTypeConverter : public TypeConverter {
Expand DownExpand Up@@ -725,6 +728,83 @@ struct StoreOpConversion : public OpConversionPattern<forth::StoreOp> {
}
};

/// Conversion pattern for forth.shared_load operation (S@).
/// Pops address from stack, loads value via shared/workgroup pointer, pushes
/// value: ( addr -- value )
struct SharedLoadOpConversion
: public OpConversionPattern<forth::SharedLoadOp> {
SharedLoadOpConversion(const TypeConverter &typeConverter,
MLIRContext *context)
: OpConversionPattern<forth::SharedLoadOp>(typeConverter, context) {}
using OneToNOpAdaptor = OpConversionPattern::OneToNOpAdaptor;

LogicalResult
matchAndRewrite(forth::SharedLoadOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
ValueRange inputStack = adaptor.getOperands()[0];
Value memref = inputStack[0];
Value stackPtr = inputStack[1];

auto i64Type = rewriter.getI64Type();
auto ptrType = LLVM::LLVMPointerType::get(rewriter.getContext(),
kWorkgroupAddressSpace);

// Load address from stack.
Value addrValue = rewriter.create<memref::LoadOp>(loc, memref, stackPtr);

// Load value from shared memory via address-space-qualified pointer.
Value ptr = rewriter.create<LLVM::IntToPtrOp>(loc, ptrType, addrValue);
Value loadedValue = rewriter.create<LLVM::LoadOp>(loc, i64Type, ptr);

// Store loaded value back at same position (replaces address).
rewriter.create<memref::StoreOp>(loc, loadedValue, memref, stackPtr);

rewriter.replaceOpWithMultiple(op, {{memref, stackPtr}});
return success();
}
};

/// Conversion pattern for forth.shared_store operation (S!).
/// Pops address and value from stack, stores value via shared/workgroup
/// pointer: ( x addr -- )
struct SharedStoreOpConversion
: public OpConversionPattern<forth::SharedStoreOp> {
SharedStoreOpConversion(const TypeConverter &typeConverter,
MLIRContext *context)
: OpConversionPattern<forth::SharedStoreOp>(typeConverter, context) {}
using OneToNOpAdaptor = OpConversionPattern::OneToNOpAdaptor;

LogicalResult
matchAndRewrite(forth::SharedStoreOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
ValueRange inputStack = adaptor.getOperands()[0];
Value memref = inputStack[0];
Value stackPtr = inputStack[1];

auto ptrType = LLVM::LLVMPointerType::get(rewriter.getContext(),
kWorkgroupAddressSpace);

// Pop address from stack.
Value addrValue = rewriter.create<memref::LoadOp>(loc, memref, stackPtr);

// Pop value from stack.
Value one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
Value spMinus1 = rewriter.create<arith::SubIOp>(loc, stackPtr, one);
Value value = rewriter.create<memref::LoadOp>(loc, memref, spMinus1);

// Store value to shared memory via address-space-qualified pointer.
Value ptr = rewriter.create<LLVM::IntToPtrOp>(loc, ptrType, addrValue);
rewriter.create<LLVM::StoreOp>(loc, value, ptr);

// New stack pointer is SP-2 (popped both address and value).
Value spMinus2 = rewriter.create<arith::SubIOp>(loc, spMinus1, one);
rewriter.replaceOpWithMultiple(op, {{memref, spMinus2}});
return success();
}
};

/// Template for converting GPU indexing ops to intrinsic ops.
/// Creates an intrinsic op with the specified name and pushes the value onto
/// the stack.
Expand DownExpand Up@@ -1026,8 +1106,9 @@ struct ConvertForthToMemRefPass
NotOpConversion, LshiftOpConversion, RshiftOpConversion, EqOpConversion,
LtOpConversion, GtOpConversion, NeOpConversion, LeOpConversion,
GeOpConversion, ZeroEqOpConversion, ParamRefOpConversion,
LoadOpConversion, StoreOpConversion, PopFlagOpConversion,
PopOpConversion, PushValueOpConversion>(typeConverter, context);
LoadOpConversion, StoreOpConversion, SharedLoadOpConversion,
SharedStoreOpConversion, PopFlagOpConversion, PopOpConversion,
PushValueOpConversion>(typeConverter, context);

// Add GPU indexing op conversion patterns
patterns.add<IntrinsicOpConversion<forth::ThreadIdXOp>>(typeConverter,
Expand Down
6 changes: 6 additions & 0 deletions lib/Translation/ForthToMLIR/ForthToMLIR.cpp
Original file line numberDiff line numberDiff line change
Expand Up@@ -487,6 +487,12 @@ Value ForthParser::emitOperation(StringRef word, Value inputStack,
} else if (word == "!") {
return builder.create<forth::StoreOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "S@") {
return builder.create<forth::SharedLoadOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "S!") {
return builder.create<forth::SharedStoreOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "TID-X") {
return builder.create<forth::ThreadIdXOp>(loc, stackType, inputStack)
.getResult();
Expand Down
13 changes: 13 additions & 0 deletions test/Conversion/ForthToMemRef/memory-ops.mlir
Original file line numberDiff line numberDiff line change
Expand Up@@ -15,6 +15,14 @@
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr
// CHECK: llvm.store %{{.*}}, %{{.*}} : i64, !llvm.ptr

// shared load (S@): pop address, inttoptr shared addrspace, llvm.load
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr<{{[1-9][0-9]*}}>
// CHECK: llvm.load %{{.*}} : !llvm.ptr<{{[1-9][0-9]*}}> -> i64

// shared store (S!): pop address + value, inttoptr shared addrspace, llvm.store
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr<{{[1-9][0-9]*}}>
// CHECK: llvm.store %{{.*}}, %{{.*}} : i64, !llvm.ptr<{{[1-9][0-9]*}}>

module {
func.func private @main() {
%0 = forth.stack !forth.stack
Expand All@@ -23,6 +31,11 @@ module {
%3 = forth.literal %2 42 : !forth.stack -> !forth.stack
%4 = forth.literal %3 100 : !forth.stack -> !forth.stack
%5 = forth.store %4 : !forth.stack -> !forth.stack
%6 = forth.literal %5 2 : !forth.stack -> !forth.stack
%7 = forth.shared_load %6 : !forth.stack -> !forth.stack
%8 = forth.literal %7 9 : !forth.stack -> !forth.stack
%9 = forth.literal %8 3 : !forth.stack -> !forth.stack
%10 = forth.shared_store %9 : !forth.stack -> !forth.stack
return
}
}
8 changes: 7 additions & 1 deletion test/Translation/Forth/memory-ops.forth
Original file line numberDiff line numberDiff line change
Expand Up@@ -6,9 +6,15 @@
\ Test ! produces forth.store
\ CHECK: forth.store %{{.*}} : !forth.stack -> !forth.stack

\ Test S@ produces forth.shared_load
\ CHECK: forth.shared_load %{{.*}} : !forth.stack -> !forth.stack

\ Test S! produces forth.shared_store
\ CHECK: forth.shared_store %{{.*}} : !forth.stack -> !forth.stack

\ Test CELLS produces literal 8 + mul
\ CHECK: forth.literal %{{.*}} 8
\ CHECK-NEXT: forth.mul
\! kernel main
1 @ 2 3 !
1 @ 2 3 ! 4 S@ 5 6 S!
4 CELLS
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers"); } } catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions CLAUDE.md
Original file line numberDiff line numberDiff line change
Expand Up@@ -87,9 +87,9 @@ uv run ruff format gpu_test/

- **Stack Type**: `!forth.stack` - untyped stack, programmer ensures type safety
- **Operations**: All take stack as input and produce stack as output (except `forth.stack`)
- **Supported Words**: literals, `DUP DROP SWAP OVER ROT NIP TUCK PICK ROLL`, `+ - * / MOD`, `AND OR XOR NOT LSHIFT RSHIFT`, `= < > <> <= >= 0=`, `@ !`, `CELLS`, `IF ELSE THEN`, `BEGIN UNTIL`, `BEGIN WHILE REPEAT`, `DO LOOP +LOOP I J K`, `LEAVE UNLOOP EXIT`, `TID-X/Y/Z BID-X/Y/Z BDIM-X/Y/Z GDIM-X/Y/Z GLOBAL-ID` (GPU indexing).
- **Supported Words**: literals, `DUP DROP SWAP OVER ROT NIP TUCK PICK ROLL`, `+ - * / MOD`, `AND OR XOR NOT LSHIFT RSHIFT`, `= < > <> <= >= 0=`, `@ !` (global memory), `S@ S!` (shared memory), `CELLS`, `IF ELSE THEN`, `BEGIN UNTIL`, `BEGIN WHILE REPEAT`, `DO LOOP +LOOP I J K`, `LEAVE UNLOOP EXIT`, `TID-X/Y/Z BID-X/Y/Z BDIM-X/Y/Z GDIM-X/Y/Z GLOBAL-ID` (GPU indexing).
- **Kernel Parameters**: Declared in the `\!` header. `\! kernel <name>` is required and must appear first. `\! param <name> i64[<N>]` becomes a `memref<Nxi64>` argument; `\! param <name> i64` becomes an `i64` argument. Using a param name in code emits `forth.param_ref` (arrays push address; scalars push value).
- **Shared Memory**: `\! shared <name> i64[<N>]` declares GPU shared (workgroup) memory. Emits a tagged `memref.alloca` at kernel entry; ForthToGPU converts it to a `gpu.func` workgroup attribution (`memref<Nxi64, #gpu.address_space<workgroup>>`). Using the shared name in code pushes its base address onto the stack. Cannot be referenced inside word definitions.
- **Shared Memory**: `\! shared <name> i64[<N>]` declares GPU shared (workgroup) memory. Emits a tagged `memref.alloca` at kernel entry; ForthToGPU converts it to a `gpu.func` workgroup attribution (`memref<Nxi64, #gpu.address_space<workgroup>>`). Using the shared name in code pushes its base address onto the stack. Use `S@`/`S!` for shared accesses. Cannot be referenced inside word definitions.
- **Conversion**: `!forth.stack` → `memref<256xi64>` with explicit stack pointer
- **GPU**: Functions wrapped in `gpu.module`, `main` gets `gpu.kernel` attribute, configured with bare pointers for NVVM conversion
- **User-defined Words**: Modeled as `func.func` with signature `(!forth.stack) -> !forth.stack`, called via `func.call`
Expand Down
62 changes: 62 additions & 0 deletions gpu_test/test_kernels.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -337,6 +337,68 @@ def test_naive_matmul_i64(kernel_runner: KernelRunner) -> None:
assert result == [12, 6, 9, 28, 14, 29]


def test_tiled_matmul_i64(kernel_runner: KernelRunner) -> None:
"""Tiled i64 matmul with shared memory: C = A(4x4) * B(4x4) -> C(4x4).

Uses 2x2 tiles, shared memory for A/B tiles, and BARRIER for sync.
Grid: (2,2,1), Block: (2,2,1) — 4 blocks of 4 threads each.
"""
result = kernel_runner.run(
forth_source=(
"\\! kernel main\n"
"\\! param A i64[16]\n"
"\\! param B i64[16]\n"
"\\! param C i64[16]\n"
"\\! shared SA i64[4]\n"
"\\! shared SB i64[4]\n"
"BID-Y 2 * TID-Y +\n"
"BID-X 2 * TID-X +\n"
"0\n"
"2 0 DO\n"
" 2 PICK 4 * I 2 * + TID-X + CELLS A + @\n"
" TID-Y 2 * TID-X + CELLS SA + S!\n"
" I 2 * TID-Y + 4 * 2 PICK + CELLS B + @\n"
" TID-Y 2 * TID-X + CELLS SB + S!\n"
" BARRIER\n"
" 2 0 DO\n"
" TID-Y 2 * I + CELLS SA + S@\n"
" I 2 * TID-X + CELLS SB + S@\n"
" * +\n"
" LOOP\n"
" BARRIER\n"
"LOOP\n"
"ROT 4 * ROT + CELLS C + !"
),
params={
"A": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
"B": [17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32],
},
grid=(2, 2, 1),
block=(2, 2, 1),
output_param=2,
output_count=16,
)
expected = [
250,
260,
270,
280,
618,
644,
670,
696,
986,
1028,
1070,
1112,
1354,
1412,
1470,
1528,
]
assert result == expected


# --- User-Defined Words ---


Expand Down
18 changes: 18 additions & 0 deletions include/warpforth/Dialect/Forth/ForthOps.td
Original file line numberDiff line numberDiff line change
Expand Up@@ -250,6 +250,24 @@ def Forth_StoreOp : Forth_StackOpBase<"store"> {
}];
}

def Forth_SharedLoadOp : Forth_StackOpBase<"shared_load"> {
let summary = "Load value from shared memory buffer";
let description = [{
Pops an address from the stack, loads a value from shared/workgroup memory at
that address, and pushes the loaded value onto the stack.
Forth semantics: ( addr -- value )
}];
}

def Forth_SharedStoreOp : Forth_StackOpBase<"shared_store"> {
let summary = "Store value to shared memory buffer";
let description = [{
Pops an address and value from the stack, stores the value to shared/workgroup
memory at the specified address.
Forth semantics: ( x addr -- )
}];
}

def Forth_ParamRefOp : Forth_Op<"param_ref", [Pure]> {
let summary = "Push kernel parameter address onto stack";
let description = [{
Expand Down
85 changes: 83 additions & 2 deletions lib/Conversion/ForthToMemRef/ForthToMemRef.cpp
Original file line numberDiff line numberDiff line change
Expand Up@@ -9,6 +9,7 @@
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Func/Transforms/FuncConversions.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/BuiltinTypes.h"
Expand All@@ -27,6 +28,8 @@ namespace {

// Stack configuration constants
constexpr int64_t kStackSize = 256;
constexpr unsigned kWorkgroupAddressSpace =
static_cast<unsigned>(gpu::AddressSpace::Workgroup);

/// Type converter for forth.stack -> memref + index
class ForthToMemRefTypeConverter : public TypeConverter {
Expand DownExpand Up@@ -725,6 +728,83 @@ struct StoreOpConversion : public OpConversionPattern<forth::StoreOp> {
}
};

/// Conversion pattern for forth.shared_load operation (S@).
/// Pops address from stack, loads value via shared/workgroup pointer, pushes
/// value: ( addr -- value )
struct SharedLoadOpConversion
: public OpConversionPattern<forth::SharedLoadOp> {
SharedLoadOpConversion(const TypeConverter &typeConverter,
MLIRContext *context)
: OpConversionPattern<forth::SharedLoadOp>(typeConverter, context) {}
using OneToNOpAdaptor = OpConversionPattern::OneToNOpAdaptor;

LogicalResult
matchAndRewrite(forth::SharedLoadOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
ValueRange inputStack = adaptor.getOperands()[0];
Value memref = inputStack[0];
Value stackPtr = inputStack[1];

auto i64Type = rewriter.getI64Type();
auto ptrType = LLVM::LLVMPointerType::get(rewriter.getContext(),
kWorkgroupAddressSpace);

// Load address from stack.
Value addrValue = rewriter.create<memref::LoadOp>(loc, memref, stackPtr);

// Load value from shared memory via address-space-qualified pointer.
Value ptr = rewriter.create<LLVM::IntToPtrOp>(loc, ptrType, addrValue);
Value loadedValue = rewriter.create<LLVM::LoadOp>(loc, i64Type, ptr);

// Store loaded value back at same position (replaces address).
rewriter.create<memref::StoreOp>(loc, loadedValue, memref, stackPtr);

rewriter.replaceOpWithMultiple(op, {{memref, stackPtr}});
return success();
}
};

/// Conversion pattern for forth.shared_store operation (S!).
/// Pops address and value from stack, stores value via shared/workgroup
/// pointer: ( x addr -- )
struct SharedStoreOpConversion
: public OpConversionPattern<forth::SharedStoreOp> {
SharedStoreOpConversion(const TypeConverter &typeConverter,
MLIRContext *context)
: OpConversionPattern<forth::SharedStoreOp>(typeConverter, context) {}
using OneToNOpAdaptor = OpConversionPattern::OneToNOpAdaptor;

LogicalResult
matchAndRewrite(forth::SharedStoreOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
ValueRange inputStack = adaptor.getOperands()[0];
Value memref = inputStack[0];
Value stackPtr = inputStack[1];

auto ptrType = LLVM::LLVMPointerType::get(rewriter.getContext(),
kWorkgroupAddressSpace);

// Pop address from stack.
Value addrValue = rewriter.create<memref::LoadOp>(loc, memref, stackPtr);

// Pop value from stack.
Value one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
Value spMinus1 = rewriter.create<arith::SubIOp>(loc, stackPtr, one);
Value value = rewriter.create<memref::LoadOp>(loc, memref, spMinus1);

// Store value to shared memory via address-space-qualified pointer.
Value ptr = rewriter.create<LLVM::IntToPtrOp>(loc, ptrType, addrValue);
rewriter.create<LLVM::StoreOp>(loc, value, ptr);

// New stack pointer is SP-2 (popped both address and value).
Value spMinus2 = rewriter.create<arith::SubIOp>(loc, spMinus1, one);
rewriter.replaceOpWithMultiple(op, {{memref, spMinus2}});
return success();
}
};

/// Template for converting GPU indexing ops to intrinsic ops.
/// Creates an intrinsic op with the specified name and pushes the value onto
/// the stack.
Expand DownExpand Up@@ -1026,8 +1106,9 @@ struct ConvertForthToMemRefPass
NotOpConversion, LshiftOpConversion, RshiftOpConversion, EqOpConversion,
LtOpConversion, GtOpConversion, NeOpConversion, LeOpConversion,
GeOpConversion, ZeroEqOpConversion, ParamRefOpConversion,
LoadOpConversion, StoreOpConversion, PopFlagOpConversion,
PopOpConversion, PushValueOpConversion>(typeConverter, context);
LoadOpConversion, StoreOpConversion, SharedLoadOpConversion,
SharedStoreOpConversion, PopFlagOpConversion, PopOpConversion,
PushValueOpConversion>(typeConverter, context);

// Add GPU indexing op conversion patterns
patterns.add<IntrinsicOpConversion<forth::ThreadIdXOp>>(typeConverter,
Expand Down
6 changes: 6 additions & 0 deletions lib/Translation/ForthToMLIR/ForthToMLIR.cpp
Original file line numberDiff line numberDiff line change
Expand Up@@ -487,6 +487,12 @@ Value ForthParser::emitOperation(StringRef word, Value inputStack,
} else if (word == "!") {
return builder.create<forth::StoreOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "S@") {
return builder.create<forth::SharedLoadOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "S!") {
return builder.create<forth::SharedStoreOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "TID-X") {
return builder.create<forth::ThreadIdXOp>(loc, stackType, inputStack)
.getResult();
Expand Down
13 changes: 13 additions & 0 deletions test/Conversion/ForthToMemRef/memory-ops.mlir
Original file line numberDiff line numberDiff line change
Expand Up@@ -15,6 +15,14 @@
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr
// CHECK: llvm.store %{{.*}}, %{{.*}} : i64, !llvm.ptr

// shared load (S@): pop address, inttoptr shared addrspace, llvm.load
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr<{{[1-9][0-9]*}}>
// CHECK: llvm.load %{{.*}} : !llvm.ptr<{{[1-9][0-9]*}}> -> i64

// shared store (S!): pop address + value, inttoptr shared addrspace, llvm.store
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr<{{[1-9][0-9]*}}>
// CHECK: llvm.store %{{.*}}, %{{.*}} : i64, !llvm.ptr<{{[1-9][0-9]*}}>

module {
func.func private @main() {
%0 = forth.stack !forth.stack
Expand All@@ -23,6 +31,11 @@ module {
%3 = forth.literal %2 42 : !forth.stack -> !forth.stack
%4 = forth.literal %3 100 : !forth.stack -> !forth.stack
%5 = forth.store %4 : !forth.stack -> !forth.stack
%6 = forth.literal %5 2 : !forth.stack -> !forth.stack
%7 = forth.shared_load %6 : !forth.stack -> !forth.stack
%8 = forth.literal %7 9 : !forth.stack -> !forth.stack
%9 = forth.literal %8 3 : !forth.stack -> !forth.stack
%10 = forth.shared_store %9 : !forth.stack -> !forth.stack
return
}
}
8 changes: 7 additions & 1 deletion test/Translation/Forth/memory-ops.forth
Original file line numberDiff line numberDiff line change
Expand Up@@ -6,9 +6,15 @@
\ Test ! produces forth.store
\ CHECK: forth.store %{{.*}} : !forth.stack -> !forth.stack

\ Test S@ produces forth.shared_load
\ CHECK: forth.shared_load %{{.*}} : !forth.stack -> !forth.stack

\ Test S! produces forth.shared_store
\ CHECK: forth.shared_store %{{.*}} : !forth.stack -> !forth.stack

\ Test CELLS produces literal 8 + mul
\ CHECK: forth.literal %{{.*}} 8
\ CHECK-NEXT: forth.mul
\! kernel main
1 @ 2 3 !
1 @ 2 3 ! 4 S@ 5 6 S!
4 CELLS
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode"); } } catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); } })(); })();
Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions CLAUDE.md
Original file line numberDiff line numberDiff line change
Expand Up@@ -87,9 +87,9 @@ uv run ruff format gpu_test/

- **Stack Type**: `!forth.stack` - untyped stack, programmer ensures type safety
- **Operations**: All take stack as input and produce stack as output (except `forth.stack`)
- **Supported Words**: literals, `DUP DROP SWAP OVER ROT NIP TUCK PICK ROLL`, `+ - * / MOD`, `AND OR XOR NOT LSHIFT RSHIFT`, `= < > <> <= >= 0=`, `@ !`, `CELLS`, `IF ELSE THEN`, `BEGIN UNTIL`, `BEGIN WHILE REPEAT`, `DO LOOP +LOOP I J K`, `LEAVE UNLOOP EXIT`, `TID-X/Y/Z BID-X/Y/Z BDIM-X/Y/Z GDIM-X/Y/Z GLOBAL-ID` (GPU indexing).
- **Supported Words**: literals, `DUP DROP SWAP OVER ROT NIP TUCK PICK ROLL`, `+ - * / MOD`, `AND OR XOR NOT LSHIFT RSHIFT`, `= < > <> <= >= 0=`, `@ !` (global memory), `S@ S!` (shared memory), `CELLS`, `IF ELSE THEN`, `BEGIN UNTIL`, `BEGIN WHILE REPEAT`, `DO LOOP +LOOP I J K`, `LEAVE UNLOOP EXIT`, `TID-X/Y/Z BID-X/Y/Z BDIM-X/Y/Z GDIM-X/Y/Z GLOBAL-ID` (GPU indexing).
- **Kernel Parameters**: Declared in the `\!` header. `\! kernel <name>` is required and must appear first. `\! param <name> i64[<N>]` becomes a `memref<Nxi64>` argument; `\! param <name> i64` becomes an `i64` argument. Using a param name in code emits `forth.param_ref` (arrays push address; scalars push value).
- **Shared Memory**: `\! shared <name> i64[<N>]` declares GPU shared (workgroup) memory. Emits a tagged `memref.alloca` at kernel entry; ForthToGPU converts it to a `gpu.func` workgroup attribution (`memref<Nxi64, #gpu.address_space<workgroup>>`). Using the shared name in code pushes its base address onto the stack. Cannot be referenced inside word definitions.
- **Shared Memory**: `\! shared <name> i64[<N>]` declares GPU shared (workgroup) memory. Emits a tagged `memref.alloca` at kernel entry; ForthToGPU converts it to a `gpu.func` workgroup attribution (`memref<Nxi64, #gpu.address_space<workgroup>>`). Using the shared name in code pushes its base address onto the stack. Use `S@`/`S!` for shared accesses. Cannot be referenced inside word definitions.
- **Conversion**: `!forth.stack` → `memref<256xi64>` with explicit stack pointer
- **GPU**: Functions wrapped in `gpu.module`, `main` gets `gpu.kernel` attribute, configured with bare pointers for NVVM conversion
- **User-defined Words**: Modeled as `func.func` with signature `(!forth.stack) -> !forth.stack`, called via `func.call`
Expand Down
62 changes: 62 additions & 0 deletions gpu_test/test_kernels.py
Original file line numberDiff line numberDiff line change
Expand Up@@ -337,6 +337,68 @@ def test_naive_matmul_i64(kernel_runner: KernelRunner) -> None:
assert result == [12, 6, 9, 28, 14, 29]


def test_tiled_matmul_i64(kernel_runner: KernelRunner) -> None:
"""Tiled i64 matmul with shared memory: C = A(4x4) * B(4x4) -> C(4x4).

Uses 2x2 tiles, shared memory for A/B tiles, and BARRIER for sync.
Grid: (2,2,1), Block: (2,2,1) — 4 blocks of 4 threads each.
"""
result = kernel_runner.run(
forth_source=(
"\\! kernel main\n"
"\\! param A i64[16]\n"
"\\! param B i64[16]\n"
"\\! param C i64[16]\n"
"\\! shared SA i64[4]\n"
"\\! shared SB i64[4]\n"
"BID-Y 2 * TID-Y +\n"
"BID-X 2 * TID-X +\n"
"0\n"
"2 0 DO\n"
" 2 PICK 4 * I 2 * + TID-X + CELLS A + @\n"
" TID-Y 2 * TID-X + CELLS SA + S!\n"
" I 2 * TID-Y + 4 * 2 PICK + CELLS B + @\n"
" TID-Y 2 * TID-X + CELLS SB + S!\n"
" BARRIER\n"
" 2 0 DO\n"
" TID-Y 2 * I + CELLS SA + S@\n"
" I 2 * TID-X + CELLS SB + S@\n"
" * +\n"
" LOOP\n"
" BARRIER\n"
"LOOP\n"
"ROT 4 * ROT + CELLS C + !"
),
params={
"A": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16],
"B": [17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32],
},
grid=(2, 2, 1),
block=(2, 2, 1),
output_param=2,
output_count=16,
)
expected = [
250,
260,
270,
280,
618,
644,
670,
696,
986,
1028,
1070,
1112,
1354,
1412,
1470,
1528,
]
assert result == expected


# --- User-Defined Words ---


Expand Down
18 changes: 18 additions & 0 deletions include/warpforth/Dialect/Forth/ForthOps.td
Original file line numberDiff line numberDiff line change
Expand Up@@ -250,6 +250,24 @@ def Forth_StoreOp : Forth_StackOpBase<"store"> {
}];
}

def Forth_SharedLoadOp : Forth_StackOpBase<"shared_load"> {
let summary = "Load value from shared memory buffer";
let description = [{
Pops an address from the stack, loads a value from shared/workgroup memory at
that address, and pushes the loaded value onto the stack.
Forth semantics: ( addr -- value )
}];
}

def Forth_SharedStoreOp : Forth_StackOpBase<"shared_store"> {
let summary = "Store value to shared memory buffer";
let description = [{
Pops an address and value from the stack, stores the value to shared/workgroup
memory at the specified address.
Forth semantics: ( x addr -- )
}];
}

def Forth_ParamRefOp : Forth_Op<"param_ref", [Pure]> {
let summary = "Push kernel parameter address onto stack";
let description = [{
Expand Down
85 changes: 83 additions & 2 deletions lib/Conversion/ForthToMemRef/ForthToMemRef.cpp
Original file line numberDiff line numberDiff line change
Expand Up@@ -9,6 +9,7 @@
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Func/Transforms/FuncConversions.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/BuiltinTypes.h"
Expand All@@ -27,6 +28,8 @@ namespace {

// Stack configuration constants
constexpr int64_t kStackSize = 256;
constexpr unsigned kWorkgroupAddressSpace =
static_cast<unsigned>(gpu::AddressSpace::Workgroup);

/// Type converter for forth.stack -> memref + index
class ForthToMemRefTypeConverter : public TypeConverter {
Expand DownExpand Up@@ -725,6 +728,83 @@ struct StoreOpConversion : public OpConversionPattern<forth::StoreOp> {
}
};

/// Conversion pattern for forth.shared_load operation (S@).
/// Pops address from stack, loads value via shared/workgroup pointer, pushes
/// value: ( addr -- value )
struct SharedLoadOpConversion
: public OpConversionPattern<forth::SharedLoadOp> {
SharedLoadOpConversion(const TypeConverter &typeConverter,
MLIRContext *context)
: OpConversionPattern<forth::SharedLoadOp>(typeConverter, context) {}
using OneToNOpAdaptor = OpConversionPattern::OneToNOpAdaptor;

LogicalResult
matchAndRewrite(forth::SharedLoadOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
ValueRange inputStack = adaptor.getOperands()[0];
Value memref = inputStack[0];
Value stackPtr = inputStack[1];

auto i64Type = rewriter.getI64Type();
auto ptrType = LLVM::LLVMPointerType::get(rewriter.getContext(),
kWorkgroupAddressSpace);

// Load address from stack.
Value addrValue = rewriter.create<memref::LoadOp>(loc, memref, stackPtr);

// Load value from shared memory via address-space-qualified pointer.
Value ptr = rewriter.create<LLVM::IntToPtrOp>(loc, ptrType, addrValue);
Value loadedValue = rewriter.create<LLVM::LoadOp>(loc, i64Type, ptr);

// Store loaded value back at same position (replaces address).
rewriter.create<memref::StoreOp>(loc, loadedValue, memref, stackPtr);

rewriter.replaceOpWithMultiple(op, {{memref, stackPtr}});
return success();
}
};

/// Conversion pattern for forth.shared_store operation (S!).
/// Pops address and value from stack, stores value via shared/workgroup
/// pointer: ( x addr -- )
struct SharedStoreOpConversion
: public OpConversionPattern<forth::SharedStoreOp> {
SharedStoreOpConversion(const TypeConverter &typeConverter,
MLIRContext *context)
: OpConversionPattern<forth::SharedStoreOp>(typeConverter, context) {}
using OneToNOpAdaptor = OpConversionPattern::OneToNOpAdaptor;

LogicalResult
matchAndRewrite(forth::SharedStoreOp op, OneToNOpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
ValueRange inputStack = adaptor.getOperands()[0];
Value memref = inputStack[0];
Value stackPtr = inputStack[1];

auto ptrType = LLVM::LLVMPointerType::get(rewriter.getContext(),
kWorkgroupAddressSpace);

// Pop address from stack.
Value addrValue = rewriter.create<memref::LoadOp>(loc, memref, stackPtr);

// Pop value from stack.
Value one = rewriter.create<arith::ConstantIndexOp>(loc, 1);
Value spMinus1 = rewriter.create<arith::SubIOp>(loc, stackPtr, one);
Value value = rewriter.create<memref::LoadOp>(loc, memref, spMinus1);

// Store value to shared memory via address-space-qualified pointer.
Value ptr = rewriter.create<LLVM::IntToPtrOp>(loc, ptrType, addrValue);
rewriter.create<LLVM::StoreOp>(loc, value, ptr);

// New stack pointer is SP-2 (popped both address and value).
Value spMinus2 = rewriter.create<arith::SubIOp>(loc, spMinus1, one);
rewriter.replaceOpWithMultiple(op, {{memref, spMinus2}});
return success();
}
};

/// Template for converting GPU indexing ops to intrinsic ops.
/// Creates an intrinsic op with the specified name and pushes the value onto
/// the stack.
Expand DownExpand Up@@ -1026,8 +1106,9 @@ struct ConvertForthToMemRefPass
NotOpConversion, LshiftOpConversion, RshiftOpConversion, EqOpConversion,
LtOpConversion, GtOpConversion, NeOpConversion, LeOpConversion,
GeOpConversion, ZeroEqOpConversion, ParamRefOpConversion,
LoadOpConversion, StoreOpConversion, PopFlagOpConversion,
PopOpConversion, PushValueOpConversion>(typeConverter, context);
LoadOpConversion, StoreOpConversion, SharedLoadOpConversion,
SharedStoreOpConversion, PopFlagOpConversion, PopOpConversion,
PushValueOpConversion>(typeConverter, context);

// Add GPU indexing op conversion patterns
patterns.add<IntrinsicOpConversion<forth::ThreadIdXOp>>(typeConverter,
Expand Down
6 changes: 6 additions & 0 deletions lib/Translation/ForthToMLIR/ForthToMLIR.cpp
Original file line numberDiff line numberDiff line change
Expand Up@@ -487,6 +487,12 @@ Value ForthParser::emitOperation(StringRef word, Value inputStack,
} else if (word == "!") {
return builder.create<forth::StoreOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "S@") {
return builder.create<forth::SharedLoadOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "S!") {
return builder.create<forth::SharedStoreOp>(loc, stackType, inputStack)
.getResult();
} else if (word == "TID-X") {
return builder.create<forth::ThreadIdXOp>(loc, stackType, inputStack)
.getResult();
Expand Down
13 changes: 13 additions & 0 deletions test/Conversion/ForthToMemRef/memory-ops.mlir
Original file line numberDiff line numberDiff line change
Expand Up@@ -15,6 +15,14 @@
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr
// CHECK: llvm.store %{{.*}}, %{{.*}} : i64, !llvm.ptr

// shared load (S@): pop address, inttoptr shared addrspace, llvm.load
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr<{{[1-9][0-9]*}}>
// CHECK: llvm.load %{{.*}} : !llvm.ptr<{{[1-9][0-9]*}}> -> i64

// shared store (S!): pop address + value, inttoptr shared addrspace, llvm.store
// CHECK: llvm.inttoptr %{{.*}} : i64 to !llvm.ptr<{{[1-9][0-9]*}}>
// CHECK: llvm.store %{{.*}}, %{{.*}} : i64, !llvm.ptr<{{[1-9][0-9]*}}>

module {
func.func private @main() {
%0 = forth.stack !forth.stack
Expand All@@ -23,6 +31,11 @@ module {
%3 = forth.literal %2 42 : !forth.stack -> !forth.stack
%4 = forth.literal %3 100 : !forth.stack -> !forth.stack
%5 = forth.store %4 : !forth.stack -> !forth.stack
%6 = forth.literal %5 2 : !forth.stack -> !forth.stack
%7 = forth.shared_load %6 : !forth.stack -> !forth.stack
%8 = forth.literal %7 9 : !forth.stack -> !forth.stack
%9 = forth.literal %8 3 : !forth.stack -> !forth.stack
%10 = forth.shared_store %9 : !forth.stack -> !forth.stack
return
}
}
8 changes: 7 additions & 1 deletion test/Translation/Forth/memory-ops.forth
Original file line numberDiff line numberDiff line change
Expand Up@@ -6,9 +6,15 @@
\ Test ! produces forth.store
\ CHECK: forth.store %{{.*}} : !forth.stack -> !forth.stack

\ Test S@ produces forth.shared_load
\ CHECK: forth.shared_load %{{.*}} : !forth.stack -> !forth.stack

\ Test S! produces forth.shared_store
\ CHECK: forth.shared_store %{{.*}} : !forth.stack -> !forth.stack

\ Test CELLS produces literal 8 + mul
\ CHECK: forth.literal %{{.*}} 8
\ CHECK-NEXT: forth.mul
\! kernel main
1 @ 2 3 !
1 @ 2 3 ! 4 S@ 5 6 S!
4 CELLS