From 93ca21caeb4807ef3ee795889e2ef9676c520c8b Mon Sep 17 00:00:00 2001
From: Zane Fink <finkzane@gmail.com>
Date: Thu, 2 Oct 2025 17:46:47 -0700
Subject: [PATCH 01/12] pj-dsl support

---
 benchmarks/hecbench/dsl/attention/Makefile   |  67 ++++
 benchmarks/hecbench/dsl/attention/main.cpp   | 315 +++++++++++++++++++
 benchmarks/hecbench/hip/feynman-kac/Makefile |   2 +-
 driver.py                                    |  10 +-
 hecbench.toml                                |  14 +
 presets/dsl.toml                             |   3 +
 presets/proteus.toml                         |  84 ++---
 7 files changed, 446 insertions(+), 49 deletions(-)
 create mode 100644 benchmarks/hecbench/dsl/attention/Makefile
 create mode 100644 benchmarks/hecbench/dsl/attention/main.cpp
 create mode 100644 presets/dsl.toml

diff --git a/benchmarks/hecbench/dsl/attention/Makefile b/benchmarks/hecbench/dsl/attention/Makefile
new file mode 100644
index 0000000..63bba36
--- /dev/null
+++ b/benchmarks/hecbench/dsl/attention/Makefile
@@ -0,0 +1,67 @@
+#===============================================================================
+# User Options
+#===============================================================================
+
+# Compiler can be set below, or via environment variable
+CC        = ${PROTEUS_CC}
+OPTIMIZE  = yes
+DEBUG     = no
+PROTEUS_PATH  ?=/path/to/proteus/install
+
+#===============================================================================
+# Program name & source code list
+#===============================================================================
+
+SUFFIX = "-proteus"
+
+program = attention$(SUFFIX).x
+
+source = main.cpp
+obj = $(source:.cpp=$(SUFFIX).o)
+
+#===============================================================================
+# Sets Flags
+#===============================================================================
+
+# Standard Flags
+CFLAGS := $(EXTRA_CFLAGS) -std=c++17 -Wall
+
+# Linker Flags
+LDFLAGS =
+
+# Debug Flags
+ifeq ($(DEBUG),yes)
+	CFLAGS += -g
+	LDFLAGS  += -g
+endif
+
+# Optimization Flags
+ifeq ($(OPTIMIZE),yes)
+	CFLAGS += -O3
+endif
+
+CFLAGS += -fpass-plugin=${PROTEUS_PATH}/lib64/libProteusPass.so -I${PROTEUS_PATH}/include
+LDFLAGS += -Wl,-rpath,${PROTEUS_PATH}/lib64 -L${PROTEUS_PATH}/lib64/ -lproteus \
+    -L${ROCM_PATH}/lib -L${ROCM_PATH}/llvm/lib \
+    -Wl,--start-group \
+    $(shell ls ${ROCM_PATH}/llvm/lib/libclang*.a) \
+    $(shell ${ROCM_PATH}/llvm/bin/llvm-config --libs) \
+    -Wl,--end-group \
+    $(shell ${ROCM_PATH}/llvm/bin/llvm-config --system-libs) \
+    -llldCommon -llldELF
+
+#===============================================================================
+# Targets to Build
+#===============================================================================
+
+$(program): $(obj) Makefile
+	$(CC) $(CFLAGS) $(obj) -o $@ $(LDFLAGS)
+
+%$(SUFFIX).o: %.cpp Makefile
+	$(CC) $(CFLAGS) -x hip -c $< -o $@
+
+clean:
+	rm -rf *.x *.o *.ll *.bc .proteus
+
+run: $(program)
+	./$(program) 65536 2048 100
diff --git a/benchmarks/hecbench/dsl/attention/main.cpp b/benchmarks/hecbench/dsl/attention/main.cpp
new file mode 100644
index 0000000..d09b969
--- /dev/null
+++ b/benchmarks/hecbench/dsl/attention/main.cpp
@@ -0,0 +1,315 @@
+#include <stdlib.h>
+#include <stdio.h>
+#include <math.h>
+#include <chrono>
+#include <random>
+#include <memory>
+#include <hip/hip_runtime.h>
+#include <proteus/JitFrontend.hpp>
+#include <proteus/Frontend/Builtins.hpp>
+#include <proteus/JitInterface.hpp>
+
+using namespace proteus;
+using namespace builtins::gpu;
+
+#define TARGET "hip"
+
+float* attention_host(float* key, float* value, float* query,
+  const int n, const int d)
+{
+// intermediate
+float* dot_product = (float*) malloc (n * sizeof(float));
+float* score = (float*) malloc (n * sizeof(float));
+// result
+float* output = (float*) malloc (d * sizeof(float));
+
+for (int i = 0; i < n; i++) {
+float sum = 0;
+for (int j = 0; j < d; j++)
+sum += key[i * d + j] * query[j];
+dot_product[i] = sum;
+}
+
+float sum = 0;
+for (int i = 0; i < n; i++)
+sum += expf(dot_product[i]);
+
+for (int i = 0; i < n; i++)
+score[i] = expf(dot_product[i]) / sum;
+
+for (int j = 0; j < d; j++) {
+float sum = 0;
+for (int i = 0; i < n; i++)
+sum += score[i] * value[i * d + j];
+output[j] = sum;
+}
+
+free(dot_product);
+free(score);
+return output;
+}
+
+// Kernel 1: Compute dot products and accumulate exp_sum
+static auto getAttentionKernel1(int n, int d) {
+  auto JitMod = std::make_unique<JitModule>(TARGET);
+  auto KernelHandle = JitMod->addKernel<void(float*, float*, float*, float*)>("attention_kernel1");
+  auto &F = KernelHandle.F;
+  {
+    auto Args = F.getArgs();
+    auto &key = std::get<0>(Args);
+    auto &query = std::get<1>(Args);
+    auto &dot_product = std::get<2>(Args);
+    auto &exp_sum = std::get<3>(Args);
+
+    F.beginFunction();
+    {
+      auto &Tidx = F.callBuiltin(getThreadIdX);
+      auto &Bidx = F.callBuiltin(getBlockIdX);
+      auto &Bdimx = F.callBuiltin(getBlockDimX);
+
+      auto &i = F.defVar(Bidx * Bdimx + Tidx);
+      auto &Nvar = F.defRuntimeConst(n);
+      auto &Dvar = F.defRuntimeConst(d);
+      auto &Zero = F.defRuntimeConst(0);
+      auto &One = F.defRuntimeConst(1);
+
+      F.beginIf(i < Nvar);
+      {
+        auto &sum = F.defVar(0.0f);
+        auto &j = F.declVar<int>("j");
+
+        F.forLoop({j, Zero, Dvar, One}, [&]() {
+          auto &keyIdx = i * Dvar + j;
+          sum = sum + (key[keyIdx] * query[j]);
+        }).emit();
+
+        dot_product[i] = sum;
+
+        auto &expVal = expf(sum);
+        F.atomicAdd( exp_sum, expVal);
+      }
+      F.endIf();
+
+      F.ret();
+    }
+    F.endFunction();
+  }
+  return std::make_pair(std::move(JitMod), KernelHandle);
+}
+
+// Kernel 2: Compute scores (normalized with exp_sum)
+static auto getAttentionKernel2(int n) {
+  auto JitMod = std::make_unique<JitModule>(TARGET);
+  auto KernelHandle = JitMod->addKernel<void(float*, float*, float*)>("attention_kernel2");
+  auto &F = KernelHandle.F;
+  {
+    auto Args = F.getArgs();
+    auto &exp_sum = std::get<0>(Args);
+    auto &dot_product = std::get<1>(Args);
+    auto &score = std::get<2>(Args);
+
+    F.beginFunction();
+    {
+      auto &Tidx = F.callBuiltin(getThreadIdX);
+      auto &Bidx = F.callBuiltin(getBlockIdX);
+      auto &Bdimx = F.callBuiltin(getBlockDimX);
+
+      auto &i = F.defVar(Bidx * Bdimx + Tidx);
+      auto &Nvar = F.defRuntimeConst(n);
+
+      F.beginIf(i < Nvar);
+      {
+        auto &expVal = expf(dot_product[i]);
+        auto &expSumVal = F.defVar(exp_sum[0]);
+        score[i] = expVal / expSumVal;
+      }
+      F.endIf();
+
+      F.ret();
+    }
+    F.endFunction();
+  }
+  return std::make_pair(std::move(JitMod), KernelHandle);
+}
+
+// Kernel 3: Compute output using scores and values
+static auto getAttentionKernel3(int n, int d) {
+  auto JitMod = std::make_unique<JitModule>(TARGET);
+  auto KernelHandle = JitMod->addKernel<void(float*, float*, float*)>("attention_kernel3");
+  auto &F = KernelHandle.F;
+  {
+    auto Args = F.getArgs();
+    auto &score = std::get<0>(Args);
+    auto &value = std::get<1>(Args);
+    auto &output = std::get<2>(Args);
+
+    F.beginFunction();
+    {
+      auto &Tidx = F.callBuiltin(getThreadIdX);
+      auto &Bidx = F.callBuiltin(getBlockIdX);
+      auto &Bdimx = F.callBuiltin(getBlockDimX);
+
+      auto &j = F.defVar(Bidx * Bdimx + Tidx);
+      auto &Dvar = F.defRuntimeConst(d);
+      auto &Nvar = F.defRuntimeConst(n);
+      auto &Zero = F.defRuntimeConst(0);
+      auto &One = F.defRuntimeConst(1);
+
+      F.beginIf(j < Dvar);
+      {
+        auto &sum = F.defVar(0.0f);
+        auto &i = F.declVar<int>("i");
+
+        F.forLoop({i, Zero, Nvar, One}, [&]() {
+          auto &valueIdx = i * Dvar + j;
+          sum = sum + (score[i] * value[valueIdx]);
+        }).emit();
+
+        output[j] = sum;
+      }
+      F.endIf();
+
+      F.ret();
+    }
+    F.endFunction();
+  }
+  return std::make_pair(std::move(JitMod), KernelHandle);
+}
+
+float* attention_device(float* key, float* value, float* query,
+                        const int n, const int d, const int repeat, const int verify)
+{
+  // input
+  float *d_key;
+  hipMalloc((void**)&d_key, n * d * sizeof(float));
+  hipMemcpy(d_key, key, n * d * sizeof(float), hipMemcpyHostToDevice);
+
+  float *d_value;
+  hipMalloc((void**)&d_value, n * d * sizeof(float));
+  hipMemcpy(d_value, value, n * d * sizeof(float), hipMemcpyHostToDevice);
+
+  float *d_query;
+  hipMalloc((void**)&d_query, d * sizeof(float));
+  hipMemcpy(d_query, query, d * sizeof(float), hipMemcpyHostToDevice);
+
+  // intermediate
+  float *d_dot_product;
+  hipMalloc((void**)&d_dot_product, n * sizeof(float));
+
+  float *d_exp_sum;
+  hipMalloc((void**)&d_exp_sum, sizeof(float));
+
+  // result
+  float *output = (float*) malloc (d * sizeof(float));
+  float *d_output;
+  hipMalloc((void**)&d_output, d * sizeof(float));
+
+  float *d_score;
+  hipMalloc((void**)&d_score, n * sizeof(float));
+
+  hipDeviceSynchronize();
+
+  // Build and compile kernels
+  auto [JitMod1, KernelHandle1] = getAttentionKernel1(n, d);
+  JitMod1->compile();
+
+  auto [JitMod2, KernelHandle2] = getAttentionKernel2(n);
+  JitMod2->compile();
+
+  auto [JitMod3, KernelHandle3] = getAttentionKernel3(n, d);
+  JitMod3->compile();
+
+  hipDeviceSynchronize();
+
+  auto start = std::chrono::steady_clock::now();
+
+  for (int k = 0; k < repeat; k++) {
+    if(verify) {
+      hipMemset(d_exp_sum, 0, 4);
+    }
+
+    KernelHandle1.launch(
+      {static_cast<unsigned int>((n+255)/256), 1u, 1u},
+      {256u, 1u, 1u},
+      0, nullptr,
+      d_key, d_query, d_dot_product, d_exp_sum);
+
+    KernelHandle2.launch(
+      {static_cast<unsigned int>((n+255)/256), 1u, 1u},
+      {256u, 1u, 1u},
+      0, nullptr,
+      d_exp_sum, d_dot_product, d_score);
+
+    KernelHandle3.launch(
+      {static_cast<unsigned int>((d+255)/256), 1u, 1u},
+      {256u, 1u, 1u},
+      0, nullptr,
+      d_score, d_value, d_output);
+  }
+
+  hipDeviceSynchronize();
+  auto end = std::chrono::steady_clock::now();
+  auto time = std::chrono::duration_cast<std::chrono::nanoseconds>(end - start).count();
+  printf("Average execution time of kernels %f (ms)\n", time * 1e-6f / repeat);
+
+  hipMemcpy(output, d_output, d * sizeof(float), hipMemcpyDeviceToHost);
+  hipFree(d_score);
+  hipFree(d_value);
+  hipFree(d_output);
+  hipFree(d_key);
+  hipFree(d_dot_product);
+  hipFree(d_exp_sum);
+  hipFree(d_query);
+  return output;
+}
+
+int main(int argc, char* argv[]) {
+  proteus::init();
+
+  if (argc != 4 && argc != 5) {
+    printf("Usage: %s <rows> <columns> <repeat> [verify]\n", argv[0]);
+    return 1;
+  }
+  const int n = atoi(argv[1]);
+  const int d = atoi(argv[2]);
+  const int r = atoi(argv[3]);
+  const int verify = (argc == 5) ? atoi(argv[4]) : 0;
+
+  // input
+  float* key = (float*) malloc (n * d * sizeof(float));
+  float* value = (float*) malloc (n * d * sizeof(float));
+  float* query = (float*) malloc (d * sizeof(float));
+
+  std::mt19937 gen(19937);
+  std::uniform_real_distribution<float> dist(-0.01f, 0.01f);
+
+  if (verify) {
+  for (int i = 0; i < n * d; i++) {
+    key[i] = dist(gen);
+    value[i] = dist(gen);
+      query[i % d] = dist(gen);
+    }
+  }
+
+  float* dout = attention_device(key, value, query, n, d, r, verify);
+
+  if (verify) {
+    float* hout = attention_host(key, value, query, n, d);
+
+    float rmse = 0;
+    for (int i = 0; i < d; i++) {
+      rmse += (hout[i] - dout[i]) * (hout[i] - dout[i]);
+    }
+    printf("RMSE = %f\n", sqrtf(rmse / d));
+
+    free(hout);
+  }
+
+  free(key);
+  free(value);
+  free(query);
+  free(dout);
+
+  proteus::finalize();
+  return 0;
+}
diff --git a/benchmarks/hecbench/hip/feynman-kac/Makefile b/benchmarks/hecbench/hip/feynman-kac/Makefile
index 6d082e8..7918cf7 100644
--- a/benchmarks/hecbench/hip/feynman-kac/Makefile
+++ b/benchmarks/hecbench/hip/feynman-kac/Makefile
@@ -3,7 +3,7 @@
 #===============================================================================
 
 # Compiler can be set below, or via environment variable
-CC        ?= ${PROTEUS_CC}
+CC        = ${PROTEUS_CC}
 OPTIMIZE  = yes
 DEBUG     = no
 LAUNCHER  =
diff --git a/driver.py b/driver.py
index 7ea5d4f..095d8d9 100644
--- a/driver.py
+++ b/driver.py
@@ -197,10 +197,8 @@ def run(self):
         results = pd.DataFrame()
         caching = pd.DataFrame()
         assert (
-            self.exemode == "aot"
-            or self.exemode == "proteus"
-            or self.exemode == "jitify"
-        ), "Expected aot or proteus or jitify for exemode"
+            self.exemode in ("aot", "proteus", "jitify", "dsl")
+        ), "Expected aot or proteus or jitify or dsl for exemode"
 
         ctime = self.builder.ctime
         exe_size = (
@@ -343,7 +341,7 @@ def run(self):
             results = pd.concat((results, df), ignore_index=True)
 
             # Skip parsing caching stats when running AOT.
-            if self.exemode != "proteus":
+            if self.exemode not in ("proteus", "dsl"):
                 continue
 
             # Parse Proteus caching info.
@@ -470,7 +468,7 @@ def main():
         "-x",
         "--exemode",
         help="execution mode",
-        choices=("aot", "proteus", "jitify"),
+        choices=("aot", "proteus", "jitify", "dsl"),
     )
     parser.add_argument(
         "-m",
diff --git a/hecbench.toml b/hecbench.toml
index 66a7710..6acb11a 100644
--- a/hecbench.toml
+++ b/hecbench.toml
@@ -3,10 +3,14 @@
 command = "make -j"
 [hecbench.config.build.nvidia.clean]
 command = "make -j clean"
+[hecbench.config.build.nvidia.dsl]
+command = "make -j"
 [hecbench.config.build.amd]
 command = "make -j"
 [hecbench.config.build.amd.clean]
 command = "make -j clean"
+[hecbench.config.build.amd.dsl]
+command = "make -j"
 
 [hecbench.adam]
 [hecbench.adam.nvidia]
@@ -127,3 +131,13 @@ path = "benchmarks/hecbench/hip/wsm5"
 exe = "wsm5-proteus.x"
 [hecbench.wsm5.inputs]
 default = "10"
+
+[hecbench.attention]
+[hecbench.attention.amd.dsl]
+path = "benchmarks/hecbench/dsl/attention"
+exe = "attention-proteus.x"
+[hecbench.attention.nvidia.dsl]
+path = "benchmarks/hecbench/dsl/attention"
+exe = "attention-proteus.x"
+[hecbench.attention.inputs]
+default = "65536 2048 100"
diff --git a/presets/dsl.toml b/presets/dsl.toml
new file mode 100644
index 0000000..206b156
--- /dev/null
+++ b/presets/dsl.toml
@@ -0,0 +1,3 @@
+[[dsl]]
+profmode = "direct"
+[[dsl.env]]
diff --git a/presets/proteus.toml b/presets/proteus.toml
index 6297cb2..3d3ee84 100644
--- a/presets/proteus.toml
+++ b/presets/proteus.toml
@@ -1,51 +1,51 @@
 [[proteus]]
 profmode = "direct"
-[[proteus.env]]
-PROTEUS_USE_STORED_CACHE = "0"
-PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "1"
-PROTEUS_SPECIALIZE_ARGS = "1"
-PROTEUS_SPECIALIZE_DIMS = "1"
+#[[proteus.env]]
+#PROTEUS_USE_STORED_CACHE = "0"
+#PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "1"
+#PROTEUS_SPECIALIZE_ARGS = "1"
+#PROTEUS_SPECIALIZE_DIMS = "1"
 [[proteus.env]]
 PROTEUS_USE_STORED_CACHE = "1"
 PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "1"
 PROTEUS_SPECIALIZE_ARGS = "1"
 PROTEUS_SPECIALIZE_DIMS = "1"
-[[proteus.env]]
-PROTEUS_USE_STORED_CACHE = "0"
-PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "0"
-PROTEUS_SPECIALIZE_ARGS = "0"
-PROTEUS_SPECIALIZE_DIMS = "0"
-[[proteus.env]]
-PROTEUS_USE_STORED_CACHE = "1"
-PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "0"
-PROTEUS_SPECIALIZE_ARGS = "0"
-PROTEUS_SPECIALIZE_DIMS = "0"
+#[[proteus.env]]
+#PROTEUS_USE_STORED_CACHE = "0"
+#PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "0"
+#PROTEUS_SPECIALIZE_ARGS = "0"
+#PROTEUS_SPECIALIZE_DIMS = "0"
+#[[proteus.env]]
+#PROTEUS_USE_STORED_CACHE = "1"
+#PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "0"
+#PROTEUS_SPECIALIZE_ARGS = "0"
+#PROTEUS_SPECIALIZE_DIMS = "0"
 
 
-[[proteus]]
-profmode = "profiler"
-[[proteus.env]]
-PROTEUS_USE_STORED_CACHE = "0"
-PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "1"
-PROTEUS_SPECIALIZE_ARGS = "1"
-PROTEUS_SPECIALIZE_DIMS = "1"
-[[proteus.env]]
-PROTEUS_USE_STORED_CACHE = "0"
-PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "0"
-PROTEUS_SPECIALIZE_ARGS = "0"
-PROTEUS_SPECIALIZE_DIMS = "0"
-[[proteus.env]]
-PROTEUS_USE_STORED_CACHE = "0"
-PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "1"
-PROTEUS_SPECIALIZE_ARGS = "0"
-PROTEUS_SPECIALIZE_DIMS = "0"
-[[proteus.env]]
-PROTEUS_USE_STORED_CACHE = "0"
-PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "0"
-PROTEUS_SPECIALIZE_ARGS = "1"
-PROTEUS_SPECIALIZE_DIMS = "0"
-[[proteus.env]]
-PROTEUS_USE_STORED_CACHE = "0"
-PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "0"
-PROTEUS_SPECIALIZE_ARGS = "0"
-PROTEUS_SPECIALIZE_DIMS = "1"
+#[[proteus]]
+#profmode = "profiler"
+#[[proteus.env]]
+#PROTEUS_USE_STORED_CACHE = "0"
+#PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "1"
+#PROTEUS_SPECIALIZE_ARGS = "1"
+#PROTEUS_SPECIALIZE_DIMS = "1"
+#[[proteus.env]]
+#PROTEUS_USE_STORED_CACHE = "0"
+#PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "0"
+#PROTEUS_SPECIALIZE_ARGS = "0"
+#PROTEUS_SPECIALIZE_DIMS = "0"
+#[[proteus.env]]
+#PROTEUS_USE_STORED_CACHE = "0"
+#PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "1"
+#PROTEUS_SPECIALIZE_ARGS = "0"
+#PROTEUS_SPECIALIZE_DIMS = "0"
+#[[proteus.env]]
+#PROTEUS_USE_STORED_CACHE = "0"
+#PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "0"
+#PROTEUS_SPECIALIZE_ARGS = "1"
+#PROTEUS_SPECIALIZE_DIMS = "0"
+#[[proteus.env]]
+#PROTEUS_USE_STORED_CACHE = "0"
+#PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "0"
+#PROTEUS_SPECIALIZE_ARGS = "0"
+#PROTEUS_SPECIALIZE_DIMS = "1"

From e3cb8631bfe3e12f8a5f09f935989afb514e3ae3 Mon Sep 17 00:00:00 2001
From: Zane Fink <finkzane@gmail.com>
Date: Thu, 2 Oct 2025 17:47:40 -0700
Subject: [PATCH 02/12] oops

---
 presets/proteus.toml | 84 ++++++++++++++++++++++----------------------
 1 file changed, 42 insertions(+), 42 deletions(-)

diff --git a/presets/proteus.toml b/presets/proteus.toml
index 3d3ee84..6297cb2 100644
--- a/presets/proteus.toml
+++ b/presets/proteus.toml
@@ -1,51 +1,51 @@
 [[proteus]]
 profmode = "direct"
-#[[proteus.env]]
-#PROTEUS_USE_STORED_CACHE = "0"
-#PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "1"
-#PROTEUS_SPECIALIZE_ARGS = "1"
-#PROTEUS_SPECIALIZE_DIMS = "1"
+[[proteus.env]]
+PROTEUS_USE_STORED_CACHE = "0"
+PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "1"
+PROTEUS_SPECIALIZE_ARGS = "1"
+PROTEUS_SPECIALIZE_DIMS = "1"
 [[proteus.env]]
 PROTEUS_USE_STORED_CACHE = "1"
 PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "1"
 PROTEUS_SPECIALIZE_ARGS = "1"
 PROTEUS_SPECIALIZE_DIMS = "1"
-#[[proteus.env]]
-#PROTEUS_USE_STORED_CACHE = "0"
-#PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "0"
-#PROTEUS_SPECIALIZE_ARGS = "0"
-#PROTEUS_SPECIALIZE_DIMS = "0"
-#[[proteus.env]]
-#PROTEUS_USE_STORED_CACHE = "1"
-#PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "0"
-#PROTEUS_SPECIALIZE_ARGS = "0"
-#PROTEUS_SPECIALIZE_DIMS = "0"
+[[proteus.env]]
+PROTEUS_USE_STORED_CACHE = "0"
+PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "0"
+PROTEUS_SPECIALIZE_ARGS = "0"
+PROTEUS_SPECIALIZE_DIMS = "0"
+[[proteus.env]]
+PROTEUS_USE_STORED_CACHE = "1"
+PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "0"
+PROTEUS_SPECIALIZE_ARGS = "0"
+PROTEUS_SPECIALIZE_DIMS = "0"
 
 
-#[[proteus]]
-#profmode = "profiler"
-#[[proteus.env]]
-#PROTEUS_USE_STORED_CACHE = "0"
-#PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "1"
-#PROTEUS_SPECIALIZE_ARGS = "1"
-#PROTEUS_SPECIALIZE_DIMS = "1"
-#[[proteus.env]]
-#PROTEUS_USE_STORED_CACHE = "0"
-#PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "0"
-#PROTEUS_SPECIALIZE_ARGS = "0"
-#PROTEUS_SPECIALIZE_DIMS = "0"
-#[[proteus.env]]
-#PROTEUS_USE_STORED_CACHE = "0"
-#PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "1"
-#PROTEUS_SPECIALIZE_ARGS = "0"
-#PROTEUS_SPECIALIZE_DIMS = "0"
-#[[proteus.env]]
-#PROTEUS_USE_STORED_CACHE = "0"
-#PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "0"
-#PROTEUS_SPECIALIZE_ARGS = "1"
-#PROTEUS_SPECIALIZE_DIMS = "0"
-#[[proteus.env]]
-#PROTEUS_USE_STORED_CACHE = "0"
-#PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "0"
-#PROTEUS_SPECIALIZE_ARGS = "0"
-#PROTEUS_SPECIALIZE_DIMS = "1"
+[[proteus]]
+profmode = "profiler"
+[[proteus.env]]
+PROTEUS_USE_STORED_CACHE = "0"
+PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "1"
+PROTEUS_SPECIALIZE_ARGS = "1"
+PROTEUS_SPECIALIZE_DIMS = "1"
+[[proteus.env]]
+PROTEUS_USE_STORED_CACHE = "0"
+PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "0"
+PROTEUS_SPECIALIZE_ARGS = "0"
+PROTEUS_SPECIALIZE_DIMS = "0"
+[[proteus.env]]
+PROTEUS_USE_STORED_CACHE = "0"
+PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "1"
+PROTEUS_SPECIALIZE_ARGS = "0"
+PROTEUS_SPECIALIZE_DIMS = "0"
+[[proteus.env]]
+PROTEUS_USE_STORED_CACHE = "0"
+PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "0"
+PROTEUS_SPECIALIZE_ARGS = "1"
+PROTEUS_SPECIALIZE_DIMS = "0"
+[[proteus.env]]
+PROTEUS_USE_STORED_CACHE = "0"
+PROTEUS_SPECIALIZE_LAUNCH_BOUNDS = "0"
+PROTEUS_SPECIALIZE_ARGS = "0"
+PROTEUS_SPECIALIZE_DIMS = "1"

From 4a8711cf619f7aabc02041fe888faba0ad5324a0 Mon Sep 17 00:00:00 2001
From: Zane Fink <finkzane@gmail.com>
Date: Thu, 2 Oct 2025 17:49:27 -0700
Subject: [PATCH 03/12] undo

---
 benchmarks/hecbench/hip/feynman-kac/Makefile | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/benchmarks/hecbench/hip/feynman-kac/Makefile b/benchmarks/hecbench/hip/feynman-kac/Makefile
index 7918cf7..6d082e8 100644
--- a/benchmarks/hecbench/hip/feynman-kac/Makefile
+++ b/benchmarks/hecbench/hip/feynman-kac/Makefile
@@ -3,7 +3,7 @@
 #===============================================================================
 
 # Compiler can be set below, or via environment variable
-CC        = ${PROTEUS_CC}
+CC        ?= ${PROTEUS_CC}
 OPTIMIZE  = yes
 DEBUG     = no
 LAUNCHER  =

From ebb05ce935cd5e1314130c796ad9724005c733c1 Mon Sep 17 00:00:00 2001
From: Zane Fink <finkzane@gmail.com>
Date: Mon, 6 Oct 2025 11:39:44 -0700
Subject: [PATCH 04/12] conv3d

---
 benchmarks/hecbench/dsl/conv3d/Makefile |  66 ++++
 benchmarks/hecbench/dsl/conv3d/main.cpp | 409 ++++++++++++++++++++++++
 hecbench.toml                           |   7 +
 3 files changed, 482 insertions(+)
 create mode 100644 benchmarks/hecbench/dsl/conv3d/Makefile
 create mode 100644 benchmarks/hecbench/dsl/conv3d/main.cpp

diff --git a/benchmarks/hecbench/dsl/conv3d/Makefile b/benchmarks/hecbench/dsl/conv3d/Makefile
new file mode 100644
index 0000000..d553473
--- /dev/null
+++ b/benchmarks/hecbench/dsl/conv3d/Makefile
@@ -0,0 +1,66 @@
+#===============================================================================
+# User Options
+#===============================================================================
+
+# Compiler can be set below, or via environment variable
+CC        = ${PROTEUS_CC}
+OPTIMIZE  = yes
+DEBUG     = no
+PROTEUS_PATH  ?= /path/to/proteus/install
+
+#===============================================================================
+# Program name & source code list
+#===============================================================================
+
+SUFFIX = -proteus
+
+program = conv3d$(SUFFIX).x
+
+source = main.cpp
+obj = $(source:.cpp=$(SUFFIX).o)
+
+#===============================================================================
+# Sets Flags
+#===============================================================================
+
+# Standard Flags
+OFFLOAD_ARCH ?= gfx942
+HIPFLAGS := --offload-arch=${OFFLOAD_ARCH}
+CFLAGS := $(EXTRA_CFLAGS) -std=c++17 -Wall -DPROTEUS_ENABLE_HIP -I${PROTEUS_PATH}/include -I${ROCM_PATH}/llvm/include/ -I${ROCM_PATH}/include $(shell ${ROCM_PATH}/llvm/bin/llvm-config --cxxflags) -fexceptions $(HIPFLAGS)
+
+# Linker Flags
+LDFLAGS = -L${ROCM_PATH}/lib -L${ROCM_PATH}/llvm/lib \
+	$(shell ${ROCM_PATH}/llvm/bin/llvm-config --libs) \
+	$(shell ${ROCM_PATH}/llvm/bin/llvm-config --system-libs) \
+	-llldCommon -llldELF -lamdhip64 -lhiprtc -lhiprtc-builtins -Wl,-rpath,${ROCM_PATH}/lib
+
+# Debug Flags
+ifeq ($(DEBUG),yes)
+	CFLAGS += -g
+	LDFLAGS  += -g
+endif
+
+# Optimization Flags
+ifeq ($(OPTIMIZE),yes)
+	CFLAGS += -O3
+endif
+
+
+# Always link against Proteus runtime for DSL builds
+LDFLAGS += -Wl,-rpath,${PROTEUS_PATH}/lib64 -L${PROTEUS_PATH}/lib64/ -lproteus
+
+#===============================================================================
+# Targets to Build
+#===============================================================================
+
+$(program): $(obj) Makefile
+	$(CC) $(CFLAGS) $(obj) -o $@ $(LDFLAGS)
+
+%$(SUFFIX).o: %.cpp Makefile
+	$(CC) $(CFLAGS) -x hip -c $< -o $@
+
+clean:
+	rm -rf *.x *.o .proteus
+
+run: $(program)
+	./$(program) 32 96 256 26 26 5 100
diff --git a/benchmarks/hecbench/dsl/conv3d/main.cpp b/benchmarks/hecbench/dsl/conv3d/main.cpp
new file mode 100644
index 0000000..f7578c4
--- /dev/null
+++ b/benchmarks/hecbench/dsl/conv3d/main.cpp
@@ -0,0 +1,409 @@
+/*
+  Reference
+  Chapter 16 in Programming massively parallel processors,
+  A hands-on approach (D. Kirk and W. Hwu)
+ */
+
+#include <assert.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <math.h>
+#include <chrono>
+#include <iostream>
+#include <memory>
+#include <utility>
+#include <hip/hip_runtime.h>
+#define TARGET "hip"
+
+#include <proteus/JitFrontend.hpp>
+#include <proteus/Frontend/Builtins.hpp>
+#include <proteus/JitInterface.hpp>
+
+using namespace proteus;
+using namespace builtins::gpu;
+
+#define TILE_WIDTH 16
+
+#define II(n,c,h,w) ((n)*C*Hin*Win+(c)*Hin*Win+(h)*Win+w)
+#define WI(n,c,h,w) ((n)*C*K*K+(c)*K*K+(h)*K+w)
+#define OI(n,c,h,w) ((n)*M*Hout*Wout+(c)*Hout*Wout+(h)*Wout+w)
+
+void verify (const float* Y, float* Y_ref, size_t Y_size)
+{
+  bool ok = true;
+  for (size_t i = 0; i < Y_size; i++) {
+    if (fabs(Y[i] - Y_ref[i]) > 1e-3f) {
+      printf("%f (device) != %f (reference)\n", Y[i], Y_ref[i]);
+      ok = false;
+      break;
+    }
+  }
+  printf("%s\n", ok ? "PASS" : "FAIL");
+}
+
+// JIT kernel builder for conv3d_s1: grid(N, M, Z)
+static auto getConv3dS1Kernel(int C_, int M_, int K_, int Hin_, int Win_, int Hout_, int Wout_, int W_grid_) {
+  auto JitMod = std::make_unique<JitModule>(TARGET);
+  auto KernelHandle = JitMod->addKernel<void(float*, float*, float*)>("conv3d_s1");
+  auto &F = KernelHandle.F;
+  {
+    auto Args = F.getArgs();
+    auto &X = std::get<0>(Args);
+    auto &W = std::get<1>(Args);
+    auto &Y = std::get<2>(Args);
+
+    F.beginFunction();
+    {
+      auto &C = F.defRuntimeConst(C_);
+      auto &M = F.defRuntimeConst(M_);
+      auto &K = F.defRuntimeConst(K_);
+      auto &Hin = F.defRuntimeConst(Hin_);
+      auto &Win = F.defRuntimeConst(Win_);
+      auto &Hout = F.defRuntimeConst(Hout_);
+      auto &Wout = F.defRuntimeConst(Wout_);
+      auto &WGrid = F.defRuntimeConst(W_grid_);
+
+      // Constants
+      auto &TileWidth = F.defRuntimeConst(TILE_WIDTH);
+      auto &Zero = F.defRuntimeConst(0);
+      auto &One = F.defRuntimeConst(1);
+
+      auto &n = F.callBuiltin(getBlockIdX);
+      auto &m = F.callBuiltin(getBlockIdY);
+      auto &h = F.callBuiltin(getBlockIdZ) / WGrid * TileWidth + F.callBuiltin(getThreadIdY);
+      auto &w = F.callBuiltin(getBlockIdZ) % WGrid * TileWidth + F.callBuiltin(getThreadIdX);
+
+      F.beginIf(h >= Hout);
+      { F.ret();}
+      F.endIf();
+      F.beginIf(w >= Wout);
+      { F.ret();}
+      F.endIf();
+
+      auto &s = F.defVar(0.0f);
+      auto &c = F.declVar<int>("c");
+      auto &p = F.declVar<int>("p");
+      auto &q = F.declVar<int>("q");
+      F.buildLoopNest(
+        F.forLoop({c, Zero, C, One}),
+        F.forLoop({p, Zero, K, One}),
+        F.forLoop({q, Zero, K, One}, [&]() {
+          s += X[II(n, c, h+p, w+q)] * W[WI(m, c, p, q)];
+        })
+      ).emit();
+
+      Y[OI(n, m, h, w)] = s;
+
+      F.ret();
+    }
+    F.endFunction();
+  }
+  return std::make_pair(std::move(JitMod), KernelHandle);
+}
+
+// JIT kernel builder for conv3d_s2: grid(M, Z, N)
+static auto getConv3dS2Kernel(int C_, int M_, int K_, int Hin_, int Win_, int Hout_, int Wout_, int W_grid_) {
+  auto JitMod = std::make_unique<JitModule>(TARGET);
+  auto KernelHandle = JitMod->addKernel<void(float*, float*, float*)>("conv3d_s2");
+  auto &F = KernelHandle.F;
+  {
+    auto Args = F.getArgs();
+    auto &X = std::get<0>(Args);
+    auto &W = std::get<1>(Args);
+    auto &Y = std::get<2>(Args);
+
+    F.beginFunction();
+    {
+      auto &C = F.defRuntimeConst(C_);
+      auto &M = F.defRuntimeConst(M_);
+      auto &K = F.defRuntimeConst(K_);
+      auto &Hin = F.defRuntimeConst(Hin_);
+      auto &Win = F.defRuntimeConst(Win_);
+      auto &Hout = F.defRuntimeConst(Hout_);
+      auto &Wout = F.defRuntimeConst(Wout_);
+      auto &WGrid = F.defRuntimeConst(W_grid_);
+
+      // Constants
+      auto &TileWidth = F.defRuntimeConst(TILE_WIDTH);
+      auto &Zero = F.defRuntimeConst(0);
+      auto &One = F.defRuntimeConst(1);
+
+      auto &m = F.callBuiltin(getBlockIdX);
+      auto &h = F.callBuiltin(getBlockIdY) / WGrid * TileWidth + F.callBuiltin(getThreadIdY);
+      auto &w = F.callBuiltin(getBlockIdY) % WGrid * TileWidth + F.callBuiltin(getThreadIdX);
+      auto &n = F.callBuiltin(getBlockIdZ);
+
+      F.beginIf(h >= Hout);
+      { F.ret();}
+      F.endIf();
+      F.beginIf(w >= Wout);
+      { F.ret();}
+      F.endIf();
+
+      auto &s = F.defVar(0.0f);
+      auto &c = F.declVar<int>("c");
+      auto &p = F.declVar<int>("p");
+      auto &q = F.declVar<int>("q");
+      F.buildLoopNest(
+        F.forLoop({c, Zero, C, One}),
+        F.forLoop({p, Zero, K, One}),
+        F.forLoop({q, Zero, K, One}, [&]() {
+          s += X[II(n, c, h+p, w+q)] * W[WI(m, c, p, q)];
+        })
+      ).emit();
+
+      Y[OI(n, m, h, w)] = s;
+
+      F.ret();
+    }
+    F.endFunction();
+  }
+  return std::make_pair(std::move(JitMod), KernelHandle);
+}
+
+// JIT kernel builder for conv3d_s3: grid(Z, N, M)
+static auto getConv3dS3Kernel(int C_, int M_, int K_, int Hin_, int Win_, int Hout_, int Wout_, int W_grid_) {
+  auto JitMod = std::make_unique<JitModule>(TARGET);
+  auto KernelHandle = JitMod->addKernel<void(float*, float*, float*)>("conv3d_s3");
+  auto &F = KernelHandle.F;
+  {
+    auto Args = F.getArgs();
+    auto &X = std::get<0>(Args);
+    auto &W = std::get<1>(Args);
+    auto &Y = std::get<2>(Args);
+
+    F.beginFunction();
+    {
+      auto &C = F.defRuntimeConst(C_);
+      auto &M = F.defRuntimeConst(M_);
+      auto &K = F.defRuntimeConst(K_);
+      auto &Hin = F.defRuntimeConst(Hin_);
+      auto &Win = F.defRuntimeConst(Win_);
+      auto &Hout = F.defRuntimeConst(Hout_);
+      auto &Wout = F.defRuntimeConst(Wout_);
+      auto &WGrid = F.defRuntimeConst(W_grid_);
+
+      // Constants
+      auto &TileWidth = F.defRuntimeConst(TILE_WIDTH);
+      auto &Zero = F.defRuntimeConst(0);
+      auto &One = F.defRuntimeConst(1);
+
+      auto &h = F.callBuiltin(getBlockIdX) / WGrid * TileWidth + F.callBuiltin(getThreadIdY);
+      auto &w = F.callBuiltin(getBlockIdX) % WGrid * TileWidth + F.callBuiltin(getThreadIdX);
+      auto &n = F.callBuiltin(getBlockIdY);
+      auto &m = F.callBuiltin(getBlockIdZ);
+
+      F.beginIf(h >= Hout);
+      { F.ret();}
+      F.endIf();
+      F.beginIf(w >= Wout);
+      { F.ret();}
+      F.endIf();
+
+      auto &s = F.defVar(0.0f);
+      auto &c = F.declVar<int>("c");
+      auto &p = F.declVar<int>("p");
+      auto &q = F.declVar<int>("q");
+      F.buildLoopNest(
+        F.forLoop({c, Zero, C, One}),
+        F.forLoop({p, Zero, K, One}),
+        F.forLoop({q, Zero, K, One}, [&]() {
+          s += X[II(n, c, h+p, w+q)] * W[WI(m, c, p, q)];
+        })
+      ).emit();
+
+      Y[OI(n, m, h, w)] = s;
+
+      F.ret();
+    }
+    F.endFunction();
+  }
+  return std::make_pair(std::move(JitMod), KernelHandle);
+}
+
+// Hin = Hout-1+K; max(h+p) is Hin - 1 as max(h) = Hout-1 and max(p) = K-1
+void reference(const float * __restrict__ X,
+               const float * __restrict__ W,
+                     float * __restrict__ Y,
+               const int N,
+               const int M,
+               const int C,
+               const int K,
+               const int Hin,
+               const int Win,
+               const int Hout,
+               const int Wout)
+{
+  for(int n = 0; n < N; n++)
+    for(int m = 0; m < M; m++)
+      for(int h = 0; h < Hout; h++)
+        for(int w = 0; w < Wout; w++) {
+          Y[OI(n, m, h, w)] = 0;
+          for(int c = 0; c < C; c++)
+            for(int p = 0; p < K; p++)
+              for(int q = 0; q < K; q++)
+                Y[OI(n, m, h, w)] += X[II(n, c, h+p, w+q)] * W[WI(m, c, p, q)];
+        }
+}
+
+void conv3D(const int N, const int C, const int M, const int Win, const int Hin, const int K, const int repeat, const int do_verify)
+{
+  const int Hout = Hin-K+1;
+  const int Wout = Win-K+1;
+
+  size_t X_size = N * C * Hin * Win;
+  size_t W_size = M * C * K * K;
+  size_t Y_size = N * M * Hout * Wout;
+  size_t X_bytes = X_size * sizeof(float);
+  size_t W_bytes = W_size * sizeof(float);
+  size_t Y_bytes = Y_size * sizeof(float);
+
+  float *X, *W, *Y, *Y_ref;
+  X = (float *)malloc(X_bytes); // input
+  W = (float *)malloc(W_bytes); // filter
+  Y = (float *)malloc(Y_bytes); // output
+
+  srand(123);
+
+
+  if (do_verify) {
+    for (size_t i = 0; i < W_size; i++) W[i] = rand() % 31;
+    for (size_t i = 0; i < X_size; i++) X[i] = rand() % 13;
+
+    for (size_t i = 0; i < Y_size; i++) {
+      Y[i] = -1;
+    }
+    Y_ref = (float *)malloc(Y_bytes);
+    for (size_t i = 0; i < Y_size; i++) {
+      Y_ref[i] = -1;
+    }
+    reference(X, W, Y_ref, N, M, C, K, Hin, Win, Hout, Wout);
+  }
+
+  float *dX, *dW, *dY;
+  hipMalloc((void **)&dX, X_bytes);
+  hipMalloc((void **)&dW, W_bytes);
+  hipMalloc((void **)&dY, Y_bytes);
+
+  hipMemcpy(dX, X, X_bytes, hipMemcpyHostToDevice);
+  hipMemcpy(dW, W, W_bytes, hipMemcpyHostToDevice);
+  hipMemcpy(dY, Y, Y_bytes, hipMemcpyHostToDevice);
+
+  int W_grid = (Wout + TILE_WIDTH - 1) / TILE_WIDTH;
+  int H_grid = (Hout + TILE_WIDTH - 1) / TILE_WIDTH;
+  int Z = H_grid * W_grid;
+
+  printf("input dimensions: C=%d Win=%d Hin=%d\n", C, Win, Hin);
+  printf("output dimensions: M=%d Wout=%d Hout=%d\n", M, Wout, Hout);
+  printf("3D grid dimensions: N=%d M=%d Z=%d\n", N, M, Z);
+
+  // Build and compile kernels
+  auto [JitMod1, KernelHandle1] = getConv3dS1Kernel(C, M, K, Hin, Win, Hout, Wout, W_grid);
+  JitMod1->compile();
+
+  auto [JitMod2, KernelHandle2] = getConv3dS2Kernel(C, M, K, Hin, Win, Hout, Wout, W_grid);
+  JitMod2->compile();
+
+  auto [JitMod3, KernelHandle3] = getConv3dS3Kernel(C, M, K, Hin, Win, Hout, Wout, W_grid);
+  JitMod3->compile();
+
+  hipDeviceSynchronize();
+
+  // Test conv3d_s1 with grid(N, M, Z)
+  auto start = std::chrono::steady_clock::now();
+  for (int i = 0; i < repeat; i++) {
+    KernelHandle1.launch(
+      {static_cast<unsigned int>(N), static_cast<unsigned int>(M), static_cast<unsigned int>(Z)},
+      {TILE_WIDTH, TILE_WIDTH, 1u},
+      0, nullptr,
+      dX, dW, dY);
+  }
+
+  hipDeviceSynchronize();
+  auto end = std::chrono::steady_clock::now();
+  auto time = std::chrono::duration_cast<std::chrono::nanoseconds>(end - start).count();
+  printf("Average kernel execution time of conv3d_s1 kernel: %f (us)\n",
+         (time * 1e-3f) / repeat);
+  if (do_verify) {
+    hipMemcpy(Y, dY, Y_bytes, hipMemcpyDeviceToHost);
+    verify(Y, Y_ref, Y_size);
+  }
+
+  // Test conv3d_s2 with grid(M, Z, N)
+  start = std::chrono::steady_clock::now();
+  for (int i = 0; i < repeat; i++) {
+    KernelHandle2.launch(
+      {static_cast<unsigned int>(M), static_cast<unsigned int>(Z), static_cast<unsigned int>(N)},
+      {TILE_WIDTH, TILE_WIDTH, 1u},
+      0, nullptr,
+      dX, dW, dY);
+  }
+
+  hipDeviceSynchronize();
+  end = std::chrono::steady_clock::now();
+  time = std::chrono::duration_cast<std::chrono::nanoseconds>(end - start).count();
+  printf("Average kernel execution time of conv3d_s2 kernel: %f (us)\n",
+         (time * 1e-3f) / repeat);
+  if (do_verify) {
+    hipMemcpy(Y, dY, Y_bytes, hipMemcpyDeviceToHost);
+    verify(Y, Y_ref, Y_size);
+  }
+
+  // Test conv3d_s3 with grid(Z, N, M)
+  start = std::chrono::steady_clock::now();
+  for (int i = 0; i < repeat; i++) {
+    KernelHandle3.launch(
+      {static_cast<unsigned int>(Z), static_cast<unsigned int>(N), static_cast<unsigned int>(M)},
+      {TILE_WIDTH, TILE_WIDTH, 1u},
+      0, nullptr,
+      dX, dW, dY);
+  }
+
+  hipDeviceSynchronize();
+  end = std::chrono::steady_clock::now();
+  time = std::chrono::duration_cast<std::chrono::nanoseconds>(end - start).count();
+  printf("Average kernel execution time of conv3d_s3 kernel: %f (us)\n",
+         (time * 1e-3f) / repeat);
+  if (do_verify) {
+    hipMemcpy(Y, dY, Y_bytes, hipMemcpyDeviceToHost);
+    verify(Y, Y_ref, Y_size);
+  }
+
+  free(X);
+  free(W);
+  free(Y);
+  if (do_verify) {
+    free(Y_ref);
+  }
+  hipFree(dX);
+  hipFree(dW);
+  hipFree(dY);
+}
+
+int main(int argc, char* argv[]) {
+  proteus::init();
+
+  if (argc != 8 && argc != 9) {
+    printf("Usage: %s <batch size:N> <input channels:C> <output feature maps:M>", argv[0]);
+    printf(" <input width:Win> <input height:Hin> <kernel size:K> <repeat> [verify (0 or 1, default 0)]\n");
+    return 1;
+  }
+
+  int N = atoi(argv[1]);
+  int C = atoi(argv[2]);
+  int M = atoi(argv[3]);
+  int W = atoi(argv[4]);
+  int H = atoi(argv[5]);
+  int K = atoi(argv[6]);
+  int repeat = atoi(argv[7]);
+  int verify = (argc == 9) ? atoi(argv[8]) : 0;
+
+  printf("3D convolution (FP32)\n");
+  printf("\n========== Warmup start ==========\n");
+  conv3D(N, C, M, W, H, K, 1000, verify);
+  printf("\n========== Warmup done ==========\n");
+  conv3D(N, C, M, W, H, K, repeat, verify);
+
+  proteus::finalize();
+  return 0;
+}
diff --git a/hecbench.toml b/hecbench.toml
index 6acb11a..430776d 100644
--- a/hecbench.toml
+++ b/hecbench.toml
@@ -141,3 +141,10 @@ path = "benchmarks/hecbench/dsl/attention"
 exe = "attention-proteus.x"
 [hecbench.attention.inputs]
 default = "65536 2048 100"
+
+[hecbench.conv3d]
+[hecbench.conv3d.amd.dsl]
+path = "benchmarks/hecbench/dsl/conv3d"
+exe = "conv3d-proteus.x"
+[hecbench.conv3d.inputs]
+default = "32 96 256 26 26 5 100"

From 78e2530843bdf6aeee8e643dbb1d68b36710d1f7 Mon Sep 17 00:00:00 2001
From: Zane Fink <finkzane@gmail.com>
Date: Mon, 6 Oct 2025 11:57:33 -0700
Subject: [PATCH 05/12] remove unnecessary proteus path

---
 .../hecbench/data/bezier-surface/control.txt     | 16 ++++++++++++++++
 benchmarks/hecbench/dsl/attention/Makefile       |  2 +-
 2 files changed, 17 insertions(+), 1 deletion(-)
 create mode 100644 benchmarks/hecbench/data/bezier-surface/control.txt

diff --git a/benchmarks/hecbench/data/bezier-surface/control.txt b/benchmarks/hecbench/data/bezier-surface/control.txt
new file mode 100644
index 0000000..68e868c
--- /dev/null
+++ b/benchmarks/hecbench/data/bezier-surface/control.txt
@@ -0,0 +1,16 @@
+-0.508359,-0.280288,-0.160131
+-0.199345,-0.515802,-0.0100263
+0.133988,-0.515802,-0.0100263
+0.418629,-0.382526,-0.407321
+-0.532679,-0.182468,-0.0100263
+0.0505979,-0.195094,1.08832
+0.383931,-0.195094,1.08832
+0.467321,-0.182468,-0.0100263
+-0.532679,0.150865,-0.0100263
+0.0505979,0.13824,1.08832
+0.383931,0.13824,1.08832
+0.467321,0.150865,-0.0100263
+-0.483195,0.349859,-0.0568787
+-0.199345,0.484198,-0.0100263
+0.133988,0.484198,-0.0100263
+0.30311,0.389596,-0.0811675
diff --git a/benchmarks/hecbench/dsl/attention/Makefile b/benchmarks/hecbench/dsl/attention/Makefile
index 63bba36..5bd253c 100644
--- a/benchmarks/hecbench/dsl/attention/Makefile
+++ b/benchmarks/hecbench/dsl/attention/Makefile
@@ -40,7 +40,7 @@ ifeq ($(OPTIMIZE),yes)
 	CFLAGS += -O3
 endif
 
-CFLAGS += -fpass-plugin=${PROTEUS_PATH}/lib64/libProteusPass.so -I${PROTEUS_PATH}/include
+CFLAGS += -I${PROTEUS_PATH}/include
 LDFLAGS += -Wl,-rpath,${PROTEUS_PATH}/lib64 -L${PROTEUS_PATH}/lib64/ -lproteus \
     -L${ROCM_PATH}/lib -L${ROCM_PATH}/llvm/lib \
     -Wl,--start-group \

From fce115e4eb202f113fba8d07f9ad2808c2c85ff1 Mon Sep 17 00:00:00 2001
From: Zane Fink <finkzane@gmail.com>
Date: Mon, 6 Oct 2025 11:57:54 -0700
Subject: [PATCH 06/12] add the run command for conv3d nvidia

---
 hecbench.toml | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/hecbench.toml b/hecbench.toml
index 430776d..fb9897a 100644
--- a/hecbench.toml
+++ b/hecbench.toml
@@ -146,5 +146,8 @@ default = "65536 2048 100"
 [hecbench.conv3d.amd.dsl]
 path = "benchmarks/hecbench/dsl/conv3d"
 exe = "conv3d-proteus.x"
+[hecbench.conv3d.nvidia.dsl]
+path = "benchmarks/hecbench/dsl/conv3d"
+exe = "conv3d-proteus.x"
 [hecbench.conv3d.inputs]
 default = "32 96 256 26 26 5 100"

From f9cd557bf8594e90d91123c39ed4d452592f8217 Mon Sep 17 00:00:00 2001
From: Zane Fink <finkzane@gmail.com>
Date: Mon, 6 Oct 2025 11:58:24 -0700
Subject: [PATCH 07/12] bezier-surface

---
 .../hecbench/dsl/bezier-surface/Makefile      |  70 +++
 .../hecbench/dsl/bezier-surface/main.cpp      | 453 ++++++++++++++++++
 hecbench.toml                                 |  10 +
 3 files changed, 533 insertions(+)
 create mode 100644 benchmarks/hecbench/dsl/bezier-surface/Makefile
 create mode 100644 benchmarks/hecbench/dsl/bezier-surface/main.cpp

diff --git a/benchmarks/hecbench/dsl/bezier-surface/Makefile b/benchmarks/hecbench/dsl/bezier-surface/Makefile
new file mode 100644
index 0000000..06b118c
--- /dev/null
+++ b/benchmarks/hecbench/dsl/bezier-surface/Makefile
@@ -0,0 +1,70 @@
+#===============================================================================
+# User Options
+#===============================================================================
+
+# Compiler can be set below, or via environment variable
+CC        = ${PROTEUS_CC}
+OPTIMIZE  = yes
+DEBUG     = no
+PROTEUS_PATH  ?=/path/to/proteus/install
+
+#===============================================================================
+# Program name & source code list
+#===============================================================================
+
+SUFFIX = "-proteus"
+
+program = bezier-surface$(SUFFIX).x
+
+source = main.cpp
+obj = $(source:.cpp=$(SUFFIX).o)
+
+#===============================================================================
+# Sets Flags
+#===============================================================================
+
+# Standard Flags
+CFLAGS := $(EXTRA_CFLAGS) -std=c++17 -Wall
+
+# Linker Flags
+LDFLAGS =
+
+# Debug Flags
+ifeq ($(DEBUG),yes)
+	CFLAGS += -g
+	LDFLAGS  += -g
+endif
+
+# Optimization Flags
+ifeq ($(OPTIMIZE),yes)
+	CFLAGS += -O3
+endif
+
+CFLAGS += -I${PROTEUS_PATH}/include -DPROTEUS_ENABLE_HIP
+LDFLAGS += -Wl,-rpath,${PROTEUS_PATH}/lib64 -L${PROTEUS_PATH}/lib64/ -lproteus \
+    -L${ROCM_PATH}/lib -L${ROCM_PATH}/llvm/lib \
+    -Wl,--start-group \
+    $(shell ls ${ROCM_PATH}/llvm/lib/libclang*.a) \
+    $(shell ${ROCM_PATH}/llvm/bin/llvm-config --libs) \
+    -Wl,--end-group \
+    $(shell ${ROCM_PATH}/llvm/bin/llvm-config --system-libs) \
+    -llldCommon -llldELF
+
+#===============================================================================
+# Targets to Build
+#===============================================================================
+
+$(program): $(obj) Makefile
+	$(CC) $(CFLAGS) $(obj) -o $@ $(LDFLAGS)
+
+%$(SUFFIX).o: %.cpp Makefile
+	$(CC) $(CFLAGS) -x hip -c $< -o $@
+
+.PHONY: clean run
+
+clean:
+	rm -rf *.x *.o *.ll *.bc .proteus
+
+# Default run (match other bezier-surface frontends)
+run: $(program)
+	./$(program) -n 8192
diff --git a/benchmarks/hecbench/dsl/bezier-surface/main.cpp b/benchmarks/hecbench/dsl/bezier-surface/main.cpp
new file mode 100644
index 0000000..5b4c5fb
--- /dev/null
+++ b/benchmarks/hecbench/dsl/bezier-surface/main.cpp
@@ -0,0 +1,453 @@
+#include <proteus/Frontend/Builtins.hpp>
+#include <proteus/JitFrontend.hpp>
+
+#include <algorithm>
+#include <chrono>
+#include <iostream>
+#include <math.h>
+#include <iomanip>
+#include <stdio.h>
+#include <stdlib.h>
+#include <vector>
+
+using namespace proteus;
+using namespace builtins::gpu;
+
+#if PROTEUS_ENABLE_HIP
+#define TARGET "hip"
+#include <hip/hip_runtime.h>
+#elif PROTEUS_ENABLE_CUDA
+#define TARGET "cuda"
+#else
+#error "Expected PROTEUS_ENABLE_HIP or PROTEUS_ENABLE_CUDA defined"
+#endif
+
+#define divceil(n, m) (((n)-1) / (m) + 1)
+
+// Params ---------------------------------------------------------------------
+struct Params {
+
+    int         work_group_size;
+    const char *file_name;
+    int         in_size_i;
+    int         in_size_j;
+    int         out_size_i;
+    int         out_size_j;
+    int         num_trials;
+    bool        verify;
+
+    Params(int argc, char **argv) {
+      work_group_size = 256;
+      file_name = "../../data/bezier-surface/control.txt";
+      in_size_i = in_size_j = 3;
+      out_size_i = out_size_j = 300;
+      num_trials = 5;
+      verify = false;
+      int opt;
+      while((opt = getopt(argc, argv, "hp:d:i:g:t:w:r:a:f:m:n:v")) >= 0) {
+        switch(opt) {
+          case 'h':
+            usage();
+            exit(0);
+            break;
+          case 'g': work_group_size = atoi(optarg); break;
+          case 't': num_trials = atoi(optarg); break;
+          case 'f': file_name = optarg; break;
+          case 'm': in_size_i = in_size_j = atoi(optarg); break;
+          case 'n': out_size_i = out_size_j = atoi(optarg); break;
+          case 'v': verify = true; break;
+          default:
+              fprintf(stderr, "\nUnrecognized option!\n");
+              usage();
+              exit(0);
+        }
+      }
+    }
+  
+    void usage() {
+      fprintf(stderr,
+          "\nUsage:  ./main [options]"
+          "\n"
+          "\nGeneral options:"
+          "\n    -h        help"
+          "\n    -g <G>    # device work-group size (default=256)"
+          "\n    -t <T>    # number of trials (default=5)"
+          "\n    -v        verify GPU results against CPU"
+          "\n"
+          "\n"
+          "\nBenchmark-specific options:"
+          "\n    -f <F>    name of input file with control points (default=input/control.txt)"
+          "\n    -m <N>    input size in both dimensions (default=3)"
+          "\n    -n <R>    output resolution in both dimensions (default=300)"
+          "\n");
+    }
+  };
+
+  // Input Data -----------------------------------------------------------------
+void read_input(float *in, const Params &p) {
+
+    // Open input file
+    FILE *f = NULL;
+    f       = fopen(p.file_name, "r");
+    if(f == NULL) {
+      puts("Error opening file");
+      exit(-1);
+    } else {
+      printf("Read data from file %s\n", p.file_name);
+    } 
+  
+  
+    // Store points from input file to array
+    int k = 0, ic = 0;
+    float vx[10000];
+    float vy[10000];
+    float vz[10000];
+    while(fscanf(f, "%f,%f,%f", &vx[ic], &vy[ic], &vz[ic]) == 3) {
+      ic++;
+    }
+    for(int i = 0; i <= p.in_size_i; i++) {
+      for(int j = 0; j <= p.in_size_j; j++) {
+        int idx = (i * (p.in_size_j + 1) + j) * 3;
+        in[idx + 0] = vx[k];
+        in[idx + 1] = vy[k];
+        in[idx + 2] = vz[k];
+        k = (k + 1) % 16;
+      }
+    }
+  }
+
+inline int compare_output(const float *outp, const float *outpCPU,
+                          int NI, int NJ, int RESOLUTIONI, int RESOLUTIONJ) {
+  float sum_delta2, sum_ref2, L1norm2;
+  sum_delta2 = 0;
+  sum_ref2   = 0;
+  L1norm2    = 0;
+  for(int i = 0; i < RESOLUTIONI; i++) {
+    for(int j = 0; j < RESOLUTIONJ; j++) {
+      int base = (i * RESOLUTIONJ + j) * 3;
+      sum_delta2 += fabsf(outp[base + 0] - outpCPU[base + 0]);
+      sum_ref2   += fabsf(outpCPU[base + 0]);
+      sum_delta2 += fabsf(outp[base + 1] - outpCPU[base + 1]);
+      sum_ref2   += fabsf(outpCPU[base + 1]);
+      sum_delta2 += fabsf(outp[base + 2] - outpCPU[base + 2]);
+      sum_ref2   += fabsf(outpCPU[base + 2]);
+    }
+  }
+  L1norm2 = sum_ref2 == 0.0f ? 0.0f : (sum_delta2 / sum_ref2);
+  if(L1norm2 >= 1e-6f){
+    printf("Test failed\n");
+    return 1;
+  }
+  return 0;
+}
+
+float BezierBlend(int k, float mu, int n) {
+  int nn, kn, nkn;
+  float blend = 1.0f;
+  nn        = n;
+  kn        = k;
+  nkn       = n - k;
+  while(nn >= 1) {
+    blend *= static_cast<float>(nn);
+    nn--;
+    if(kn > 1) {
+      blend /= static_cast<float>(kn);
+      kn--;
+    }
+    if(nkn > 1) {
+      blend /= static_cast<float>(nkn);
+      nkn--;
+    }
+  }
+  if(k > 0)
+    blend *= powf(mu, static_cast<float>(k));
+  if(n - k > 0)
+    blend *= powf(1.0f - mu, static_cast<float>(n - k));
+  return blend;
+}
+
+// Sequential implementation for comparison purposes
+void BezierCPU(const float *inp,
+               float *outp,
+               const int NI, const int NJ, const int RESOLUTIONI, const int RESOLUTIONJ) {
+  for(int i = 0; i < RESOLUTIONI; i++) {
+    float mui = i / static_cast<float>(RESOLUTIONI - 1);
+    for(int j = 0; j < RESOLUTIONJ; j++) {
+      float muj = j / static_cast<float>(RESOLUTIONJ - 1);
+      float out_x = 0.0f;
+      float out_y = 0.0f;
+      float out_z = 0.0f;
+      for(int ki = 0; ki <= NI; ki++) {
+        float bi = BezierBlend(ki, mui, NI);
+        for(int kj = 0; kj <= NJ; kj++) {
+          float bj = BezierBlend(kj, muj, NJ);
+          int idx = (ki * (NJ + 1) + kj) * 3;
+          float coeff = bi * bj;
+          out_x += inp[idx + 0] * coeff;
+          out_y += inp[idx + 1] * coeff;
+          out_z += inp[idx + 2] * coeff;
+        }
+      }
+      int out_idx = (i * RESOLUTIONJ + j) * 3;
+      outp[out_idx + 0] = out_x;
+      outp[out_idx + 1] = out_y;
+      outp[out_idx + 2] = out_z;
+    }
+  }
+}
+
+auto createJitModule(int _NI, int _NJ, int _RESOLUTIONI, int _RESOLUTIONJ) {
+  auto J = std::make_unique<JitModule>(TARGET);
+  auto KernelHandle =
+      J->addKernel<void(float *, float*)>("BezierGPU");
+  auto &F = KernelHandle.F;
+
+  {
+    auto [inp, outp] = F.getArgs();
+
+    F.beginFunction();
+    {
+      auto [NI, NJ, RESOLUTIONI, RESOLUTIONJ] =
+          F.defRuntimeConsts(_NI, _NJ, _RESOLUTIONI, _RESOLUTIONJ);
+
+      auto &i = F.declVar<int>();
+      i = F.callBuiltin(getBlockDimX) * F.callBuiltin(getBlockIdX) + F.callBuiltin(getThreadIdX);
+
+      F.beginIf(i >= RESOLUTIONI);
+      {
+        F.ret();
+      }
+      F.endIf();
+
+      auto &mui = F.declVar<float>();
+      mui = F.convert<float>(i) / F.convert<float>(RESOLUTIONI - F.defRuntimeConst<int>(1));
+
+      auto &j = F.declVar<int>();
+      auto &InitJ = F.defRuntimeConst<int>(0);
+      auto &IncJ = F.defRuntimeConst<int>(1);
+      F.beginFor(j, InitJ, RESOLUTIONJ, IncJ);
+      {
+        auto &muj = F.convert<float>(j) / F.convert<float>(RESOLUTIONJ - F.defRuntimeConst<int>(1));
+
+        auto &OutX = F.defVar<float>(0.0f);
+        auto &OutY = F.defVar<float>(0.0f);
+        auto &OutZ = F.defVar<float>(0.0f);
+
+        auto &ki = F.declVar<int>();
+        auto &InitKi = F.defRuntimeConst<int>(0);
+        auto &UpperKi = NI + F.defRuntimeConst<int>(1);
+        auto &IncKi = F.defRuntimeConst<int>(1);
+        F.beginFor(ki, InitKi, UpperKi, IncKi);
+        {
+          // float bi = BezierBlend(ki, mui, NI);
+          auto &bi = F.call<float(int, float, int)>("BezierBlend", ki, mui, NI);
+
+          // for(int kj = 0; kj <= NJ; kj++)
+          auto &kj = F.declVar<int>();
+          auto &InitKj = F.defRuntimeConst<int>(0);
+          auto &UpperKj = NJ + F.defRuntimeConst<int>(1);
+          auto &IncKj = F.defRuntimeConst<int>(1);
+          F.beginFor(kj, InitKj, UpperKj, IncKj);
+          {
+            // float bj = BezierBlend(kj, muj, NJ);
+            auto &bj = F.call<float(int, float, int)>("BezierBlend", kj, muj, NJ);
+
+            // int idx = (ki * (NJ + 1) + kj) * 3;
+            auto &idx = F.declVar<int>();
+            idx = (ki * (NJ + F.defRuntimeConst<int>(1)) + kj) * F.defRuntimeConst<int>(3);
+
+            // float coeff = bi * bj;
+            auto &coeff = F.declVar<float>();
+            coeff = bi * bj;
+
+            // out_x += inp[idx + 0] * coeff;
+            OutX += inp[idx + F.defRuntimeConst<int>(0)] * coeff;
+            // out_y += inp[idx + 1] * coeff;
+            OutY += inp[idx + F.defRuntimeConst<int>(1)] * coeff;
+            // out_z += inp[idx + 2] * coeff;
+            OutZ += inp[idx + F.defRuntimeConst<int>(2)] * coeff;
+          }
+          F.endFor();
+        }
+        F.endFor();
+
+        // int out_idx = (i * RESOLUTIONJ + j) * 3;
+        auto &OutIdx = F.declVar<int>();
+        OutIdx = (i * RESOLUTIONJ + j) * F.defRuntimeConst<int>(3);
+
+        // outp[out_idx + 0] = out_x;
+        outp[OutIdx + F.defRuntimeConst<int>(0)] = OutX;
+        // outp[out_idx + 1] = out_y;
+        outp[OutIdx + F.defRuntimeConst<int>(1)] = OutY;
+        // outp[out_idx + 2] = out_z;
+        outp[OutIdx + F.defRuntimeConst<int>(2)] = OutZ;
+      }
+      F.endFor();
+
+      F.ret();
+    }
+    F.endFunction();
+  }
+
+  {
+    auto &F = J->addFunction<float(int, float, int)>("BezierBlend");
+    F.beginFunction();
+    {
+      auto [k, mu, n] = F.getArgs();
+      auto &blend = F.defVar<float>(1.0f);
+      auto &nn = F.declVar<int>();
+      nn = n;
+      auto &kn = F.declVar<int>();
+      kn = k;
+      auto &nkn = F.declVar<int>();
+      nkn = n - k;
+
+      auto &Cond = nn >= F.defRuntimeConst<int>(1);
+      F.beginWhile(Cond);
+      {
+        blend *= F.convert<float>(nn);
+        nn -= F.defRuntimeConst<int>(1);
+
+        F.beginIf(kn > F.defRuntimeConst<int>(1));
+        {
+          blend /= F.convert<float>(kn);
+          kn -= F.defRuntimeConst<int>(1);
+        }
+        F.endIf();
+
+        F.beginIf(nkn > F.defRuntimeConst<int>(1));
+        {
+          blend /= F.convert<float>(nkn);
+          nkn -= F.defRuntimeConst<int>(1);
+        }
+        F.endIf();
+        
+        Cond = nn >= F.defRuntimeConst<int>(1);
+      }
+      F.endWhile();
+
+      F.beginIf(k > F.defRuntimeConst<int>(0));
+      {
+        blend *= powf(mu, F.convert<float>(k));
+      }
+      F.endIf();
+
+      F.beginIf(n - k > F.defRuntimeConst<int>(0));
+      {
+        blend *= powf(F.defRuntimeConst<float>(1.0f) - mu, F.convert<float>(n - k));
+      }
+      F.endIf();
+      
+
+      F.ret(blend);
+    }
+    F.endFunction();
+  }
+
+  return std::make_pair(std::move(J), KernelHandle);
+}
+
+void run(float *in,
+    int in_size_i, int in_size_j, int out_size_i, int out_size_j, const Params &p) {
+
+size_t out_elems = static_cast<size_t>(out_size_i) * static_cast<size_t>(out_size_j);
+float *gpu_out = (float *)malloc(out_elems * 3 * sizeof(float));
+float *cpu_out = nullptr;
+
+// CPU run for verification if requested
+if (p.verify) {
+  cpu_out = (float *)malloc(out_elems * 3 * sizeof(float));
+  auto start = std::chrono::steady_clock::now();
+  BezierCPU(in, cpu_out,
+            in_size_i, in_size_j, out_size_i, out_size_j);
+  auto end = std::chrono::steady_clock::now();
+  auto time = std::chrono::duration<double, std::milli>(end - start).count();
+  std::cout << "host execution time: " << std::fixed << std::setprecision(4) << time << "ms" << std::endl;
+}
+
+// Device run - kernel creation and compilation (done once)
+size_t in_size   = static_cast<size_t>(in_size_i + 1) * static_cast<size_t>(in_size_j + 1) * 3 * sizeof(float);
+size_t out_size  = out_elems * 3 * sizeof(float);
+
+auto create_jit_start = std::chrono::steady_clock::now();
+auto [J, KernelHandle] = createJitModule(in_size_i, in_size_j, out_size_i, out_size_j);
+auto create_jit_end = std::chrono::steady_clock::now();
+J->compile();
+auto compile_end = std::chrono::steady_clock::now();
+auto compile_time = std::chrono::duration<double, std::milli>(compile_end - create_jit_end).count();
+auto create_jit_time = std::chrono::duration<double, std::milli>(create_jit_end - create_jit_start).count();
+std::cout << "kernel creation time: " << std::fixed << std::setprecision(4) << create_jit_time << "ms" << std::endl;
+std::cout << "kernel compilation time: " << std::fixed << std::setprecision(4) << compile_time << "ms" << std::endl;
+
+// Trial loop
+std::vector<double> trial_times;
+dim3 block(p.work_group_size);
+dim3 grid((out_size_i + p.work_group_size - 1) / p.work_group_size);
+
+for (int trial = 0; trial < p.num_trials; trial++) {
+  float *d_in;
+  float *d_out;
+
+  auto trial_start = std::chrono::steady_clock::now();
+
+  // Allocate device memory
+  hipMalloc((void**)&d_in, in_size);
+  hipMalloc((void**)&d_out, out_size);
+
+  // Transfer data to device
+  hipMemcpy(d_in, in, in_size, hipMemcpyHostToDevice);
+
+  // Launch kernel
+  hipDeviceSynchronize();
+  auto kstart = std::chrono::steady_clock::now();
+
+  KernelHandle.launch({grid.x, grid.y, grid.z}, {block.x, block.y, block.z}, 0, nullptr,
+                  d_in,
+                  d_out);
+
+  hipDeviceSynchronize();
+  auto kend = std::chrono::steady_clock::now();
+  auto ktime = std::chrono::duration<double, std::milli>(kend - kstart).count();
+
+  // Transfer data back to host
+  hipMemcpy(gpu_out, d_out, out_size, hipMemcpyDeviceToHost);
+
+  auto trial_end = std::chrono::steady_clock::now();
+  auto trial_time = std::chrono::duration<double, std::milli>(trial_end - trial_start).count();
+  trial_times.push_back(trial_time);
+
+  std::cout << "Trial " << (trial + 1) << " - kernel execution time: " << std::fixed << std::setprecision(4) << ktime << "ms, total time: " << trial_time << "ms" << std::endl;
+
+  // Free device memory
+  hipFree(d_in);
+  hipFree(d_out);
+}
+
+// Verify
+if (p.verify) {
+  int status = compare_output(gpu_out, cpu_out,
+                             in_size_i, in_size_j, out_size_i, out_size_j);
+  printf("%s\n", (status == 0) ? "PASS" : "FAIL");
+  free(cpu_out);
+}
+
+free(gpu_out);
+}
+
+int main(int argc, char **argv) {
+
+    const Params p(argc, argv);
+    int num_points   = (p.in_size_i + 1) * (p.in_size_j + 1);
+    size_t in_size_bytes   = static_cast<size_t>(num_points) * 3 * sizeof(float);
+    //int out_size  = p.out_size_i * p.out_size_j * sizeof(XYZ);
+  
+    // load data into h_in
+    float* in = (float *)malloc(in_size_bytes);
+    read_input(in, p);
+  
+    // run the app on the cpu and gpu
+    run(in, p.in_size_i, p.in_size_j, p.out_size_i, p.out_size_j, p);
+  
+    free(in);
+    return 0;
+  }
diff --git a/hecbench.toml b/hecbench.toml
index fb9897a..feadfca 100644
--- a/hecbench.toml
+++ b/hecbench.toml
@@ -151,3 +151,13 @@ path = "benchmarks/hecbench/dsl/conv3d"
 exe = "conv3d-proteus.x"
 [hecbench.conv3d.inputs]
 default = "32 96 256 26 26 5 100"
+
+[hecbench.bezier-surface]
+[hecbench.bezier-surface.amd.dsl]
+path = "benchmarks/hecbench/dsl/bezier-surface"
+exe = "bezier-surface-proteus.x"
+[hecbench.bezier-surface.nvidia.dsl]
+path = "benchmarks/hecbench/dsl/bezier-surface"
+exe = "bezier-surface-proteus.x"
+[hecbench.bezier-surface.inputs]
+default = "-n 8192"

From 5649cf2d87e93e1275f89733e38d8aee9ab2fb2b Mon Sep 17 00:00:00 2001
From: Zane Fink <finkzane@gmail.com>
Date: Mon, 6 Oct 2025 12:40:01 -0700
Subject: [PATCH 08/12] add pj-dsl adam

---
 benchmarks/hecbench/dsl/adam/Makefile |  66 +++++++
 benchmarks/hecbench/dsl/adam/adam.cpp | 239 ++++++++++++++++++++++++++
 hecbench.toml                         |   8 +
 3 files changed, 313 insertions(+)
 create mode 100644 benchmarks/hecbench/dsl/adam/Makefile
 create mode 100644 benchmarks/hecbench/dsl/adam/adam.cpp

diff --git a/benchmarks/hecbench/dsl/adam/Makefile b/benchmarks/hecbench/dsl/adam/Makefile
new file mode 100644
index 0000000..771a49d
--- /dev/null
+++ b/benchmarks/hecbench/dsl/adam/Makefile
@@ -0,0 +1,66 @@
+#===============================================================================
+# User Options
+#===============================================================================
+
+# Compiler can be set below, or via environment variable
+CC        = ${PROTEUS_CC}
+OPTIMIZE  = yes
+DEBUG     = no
+PROTEUS_PATH  ?= /path/to/proteus/install
+
+#===============================================================================
+# Program name & source code list
+#===============================================================================
+
+SUFFIX = -proteus
+
+program = adam$(SUFFIX).x
+
+source = adam.cpp
+obj = $(source:.cpp=$(SUFFIX).o)
+
+#===============================================================================
+# Sets Flags
+#===============================================================================
+
+# Standard Flags
+OFFLOAD_ARCH ?= gfx942
+HIPFLAGS := --offload-arch=${OFFLOAD_ARCH}
+CFLAGS := $(EXTRA_CFLAGS) -std=c++17 -Wall -DPROTEUS_ENABLE_HIP -I${PROTEUS_PATH}/include -I${ROCM_PATH}/llvm/include/ -I${ROCM_PATH}/include $(shell ${ROCM_PATH}/llvm/bin/llvm-config --cxxflags) -fexceptions $(HIPFLAGS)
+
+# Linker Flags
+LDFLAGS = -L${ROCM_PATH}/lib -L${ROCM_PATH}/llvm/lib \
+	$(shell ${ROCM_PATH}/llvm/bin/llvm-config --libs) \
+	$(shell ${ROCM_PATH}/llvm/bin/llvm-config --system-libs) \
+	-llldCommon -llldELF -lamdhip64 -lhiprtc -lhiprtc-builtins -Wl,-rpath,${ROCM_PATH}/lib
+
+# Debug Flags
+ifeq ($(DEBUG),yes)
+	CFLAGS += -g
+	LDFLAGS  += -g
+endif
+
+# Optimization Flags
+ifeq ($(OPTIMIZE),yes)
+	CFLAGS += -O3
+endif
+
+
+# Always link against Proteus runtime for DSL builds
+LDFLAGS += -Wl,-rpath,${PROTEUS_PATH}/lib64 -L${PROTEUS_PATH}/lib64/ -lproteus
+
+#===============================================================================
+# Targets to Build
+#===============================================================================
+
+$(program): $(obj) Makefile
+	$(CC) $(CFLAGS) $(obj) -o $@ $(LDFLAGS)
+
+%$(SUFFIX).o: %.cpp Makefile
+	$(CC) $(CFLAGS) -x hip -c $< -o $@
+
+clean:
+	rm -rf *.x *.o .proteus
+
+run: $(program)
+	./$(program) 10000 200 100
diff --git a/benchmarks/hecbench/dsl/adam/adam.cpp b/benchmarks/hecbench/dsl/adam/adam.cpp
new file mode 100644
index 0000000..72c8f4f
--- /dev/null
+++ b/benchmarks/hecbench/dsl/adam/adam.cpp
@@ -0,0 +1,239 @@
+// NOLINTBEGIN
+
+// clang-format off
+// RUN: rm -rf "%t.$$.proteus"
+// RUN: PROTEUS_CACHE_DIR="%t.$$.proteus" %build/adam_runconst.%ext 10000 200 100 1 | %FILECHECK %s --check-prefixes=CHECK,CHECK-FIRST
+// Second run uses the object cache.
+// RUN: PROTEUS_CACHE_DIR="%t.$$.proteus" %build/adam_runconst.%ext 10000 200 100 1 | %FILECHECK %s --check-prefixes=CHECK,CHECK-SECOND
+// RUN: rm -rf "%t.$$.proteus"
+// clang-format on
+
+#include <proteus/Frontend/Builtins.hpp>
+#include <proteus/JitFrontend.hpp>
+
+#include <hip/hip_runtime.h>
+
+#include <chrono>
+#include <iostream>
+#include <math.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+using namespace proteus;
+using namespace builtins::gpu;
+
+#if PROTEUS_ENABLE_HIP
+#define TARGET "hip"
+#elif PROTEUS_ENABLE_CUDA
+#define TARGET "cuda"
+#else
+#error "Expected PROTEUS_ENABLE_HIP or PROTEUS_ENABLE_CUDA defined"
+#endif
+
+
+typedef enum {
+  ADAM_MODE_0 = 0, // eps under square root
+  ADAM_MODE_1 = 1  // eps outside square root
+} adamMode_t;
+
+
+
+auto createJitModuleSpecial(float _b1, float _b2, float _eps, float _grad_scale,
+                            float _step_size, int _time_step,
+                            size_t _vector_size, int _mode, float _decay) {
+  auto J = std::make_unique<JitModule>(TARGET);
+  auto KernelHandle =
+      J->addKernel<void(float *, float *, float *, float *)>("adam");
+  auto &F = KernelHandle.F;
+  auto [p, m, v, g] = F.getArgs();
+
+
+  F.beginFunction();
+  {
+    auto [b1, b2, eps, grad_scale, step_size, time_step, vector_size, mode,
+          decay] = F.defRuntimeConsts(_b1, _b2, _eps, _grad_scale, _step_size,
+                                      _time_step, _vector_size, _mode, _decay);
+
+    
+    auto &i = F.declVar<size_t>("i");
+    auto &totThreads = F.declVar<size_t>("totThreads");
+    auto &j = F.declVar<size_t>("j");
+    auto &t = F.declVar<int>("t");
+    auto &inc1 = F.defRuntimeConst<int>(1);
+
+    i = F.callBuiltin(getBlockIdX) * F.callBuiltin(getBlockDimX) +
+        F.callBuiltin(getThreadIdX);
+    totThreads = F.callBuiltin(getGridDimX) * F.callBuiltin(getBlockDimX);
+
+    F.beginFor(j, i, vector_size, totThreads);
+    {
+      auto &lim = F.declVar<int>("lim");
+      t = 1;
+      lim = time_step;
+      F.beginFor(t, t, lim, inc1);
+      {
+        auto &scaled_grad = F.declVar<float>("scale_grad");
+        scaled_grad = g[j] / grad_scale;
+
+        m[j] = b1 * m[j] + (1.f - b1) * scaled_grad;
+        v[j] = b2 * v[j] + (1.f - b2) * scaled_grad * scaled_grad;
+
+        auto &m_corrected = F.declVar<float>("m_corrected");
+        auto &v_corrected = F.declVar<float>("v_corrected");
+        m_corrected = m[j] / (1.f - powf(b1, t));
+        v_corrected = v[j] / (1.f - powf(b2, t));
+
+        auto &denom = F.declVar<float>("denom");
+        F.beginIf(mode == 0);
+        { denom = sqrtf(v_corrected + eps); }
+        F.endIf();
+
+        F.beginIf(mode == 1);
+        { denom = sqrtf(v_corrected) + eps; }
+        F.endIf();
+
+        auto &update = F.declVar<float>("update");
+        update = (m_corrected / denom) + (decay * p[j]);
+
+        p[j] -= (step_size * update);
+      }
+      F.endFor();
+    }
+    F.endFor();
+    F.ret();
+  }
+  F.endFunction();
+
+  return std::make_pair(std::move(J), KernelHandle);
+}
+
+int main(int argc, char *argv[]) {
+  if (argc < 4 || argc > 5) {
+    printf("Usage: %s <vector size> <number of time steps> <repeat>\n",
+           argv[0]);
+    return 1;
+  }
+
+  const int vector_size = atoi(argv[1]);
+  const int time_step = atoi(argv[2]);
+  const int repeat = atoi(argv[3]);
+
+  size_t size_bytes = vector_size * sizeof(float);
+
+  float *m = (float *)malloc(size_bytes);
+  float *v = (float *)malloc(size_bytes);
+  float *g = (float *)malloc(size_bytes);
+  float *p = (float *)malloc(size_bytes);
+  float *r = (float *)malloc(size_bytes);
+
+  srand(123);
+  for (int i = 0; i < vector_size; i++) {
+    m[i] = rand() / (float)RAND_MAX;
+    v[i] = rand() / (float)RAND_MAX;
+    g[i] = rand() / (float)RAND_MAX;
+    r[i] = p[i] = rand() / (float)RAND_MAX;
+  }
+
+  float *d_m, *d_v, *d_g, *d_p;
+
+  hipMalloc((void **)&d_m, size_bytes);
+  hipMemcpy(d_m, m, size_bytes, hipMemcpyHostToDevice);
+
+  hipMalloc((void **)&d_v, size_bytes);
+  hipMemcpy(d_v, v, size_bytes, hipMemcpyHostToDevice);
+
+  hipMalloc((void **)&d_g, size_bytes);
+  hipMemcpy(d_g, g, size_bytes, hipMemcpyHostToDevice);
+
+  hipMalloc((void **)&d_p, size_bytes);
+  hipMemcpy(d_p, p, size_bytes, hipMemcpyHostToDevice);
+
+  // Arbitrary constants
+  const float step_size = 1e-3f;
+  const float decay = 0.5f;
+  const float beta1 = 0.9f;
+  const float beta2 = 0.999f;
+  const float eps = 1e-8f;
+  const float grad_scale = 256.f;
+
+  const int threadsPerBlock = 256;
+  const dim3 grids((vector_size + threadsPerBlock - 1) / threadsPerBlock);
+  const dim3 blocks(threadsPerBlock);
+
+  adamMode_t mode = ADAM_MODE_0;
+
+  auto [J, KernelHandle] =
+      createJitModuleSpecial(beta1, beta2, eps, grad_scale, step_size,
+                             time_step, vector_size, mode, decay);
+
+  J->compile();
+
+  hipDeviceSynchronize();
+
+
+  auto start = std::chrono::steady_clock::now();
+
+  for (int i = 0; i < repeat; i++) {
+    // adam<float, float><<<grids, blocks>>>(d_p, d_m, d_v, d_g, beta1, beta2,
+    // eps,
+    //                                       grad_scale, step_size, time_step,
+    //                                       vector_size, mode, decay);
+
+    KernelHandle.launch({grids.x, 1, 1}, {blocks.x, 1, 1}, 0, 0,
+                                    d_p, d_m, d_v, d_g);
+  }
+
+  hipDeviceSynchronize();
+  auto end = std::chrono::steady_clock::now();
+  auto time =
+      std::chrono::duration_cast<std::chrono::nanoseconds>(end - start).count();
+  printf("Average kernel execution time %f (ms)\n", time * 1e-6f / repeat);
+
+  hipMemcpy(p, d_p, size_bytes, hipMemcpyDeviceToHost);
+
+
+  hipFree(d_p);
+  hipFree(d_m);
+  hipFree(d_v);
+  hipFree(d_g);
+
+  free(p);
+  free(m);
+  free(v);
+  free(g);
+  free(r);
+  return 0;
+}
+
+// clang-format off
+// We got slight differences in the output for the least significant digits.
+// Could be HW, numeric, or the way we handle signedness.
+// CHECK: init p[0] = 0.348563
+// CHECK-NEXT: init p[1] = 0.259322
+// CHECK-NEXT: init p[2] = 0.377145
+// CHECK-NEXT: init p[3] = 0.486632
+// CHECK-NEXT: init p[4] = 0.352038
+// CHECK-NEXT: init p[5] = 0.0784863
+// CHECK-NEXT: init p[6] = 0.968732
+// CHECK-NEXT: init p[7] = 0.852707
+// CHECK-NEXT: init p[8] = 0.153431
+// CHECK-NEXT: init p[9] = 0.559506
+// CHECK-NEXT: Creating JIT module
+// CHECK-NEXT: Compiling JIT module
+// CHECK-NEXT: Average kernel execution time {{.*}} (ms)
+// CHECK-NEXT: p[0] = -0.57293
+// CHECK-NEXT: p[1] = -0.59603
+// CHECK-NEXT: p[2] = -0.592634
+// CHECK-NEXT: p[3] = -0.588154
+// CHECK-NEXT: p[4] = -0.593454
+// CHECK-NEXT: p[5] = -0.591989
+// CHECK-NEXT: p[6] = -0.573486
+// CHECK-NEXT: p[7] = -0.599872
+// CHECK-NEXT: p[8] = -0.58157
+// CHECK-NEXT: p[9] = -0.59015{{[7|8]}}
+// CHECK: JitCache hits 0 total 1
+// CHECK: HashValue {{[0-9]+}} NumExecs 1 NumHits 0
+// CHECK-FIRST: JitStorageCache hits 0 total 1
+// CHECK-SECOND: JitStorageCache hits 1 total 1
+
+// NOLINTEND
diff --git a/hecbench.toml b/hecbench.toml
index feadfca..5470d59 100644
--- a/hecbench.toml
+++ b/hecbench.toml
@@ -32,6 +32,14 @@ exe = "adam-proteus.x"
 [hecbench.adam.inputs]
 default = "160000 1600 1000"
 
+[hecbench.adam.dsl]
+[hecbench.adam.amd.dsl]
+path = "benchmarks/hecbench/dsl/adam"
+exe = "adam-proteus.x"
+[hecbench.adam.nvidia.dsl]
+path = "benchmarks/hecbench/dsl/adam"
+exe = "adam-proteus.x"
+
 [hecbench.feynman-kac]
 [hecbench.feynman-kac.nvidia]
 [hecbench.feynman-kac.nvidia.aot]

From 0b5968ade75901f560f2fb98c357fbcbd8dcdc7f Mon Sep 17 00:00:00 2001
From: Zane Fink <finkzane@gmail.com>
Date: Mon, 6 Oct 2025 12:41:15 -0700
Subject: [PATCH 09/12] Remove the test stuff

---
 benchmarks/hecbench/dsl/adam/adam.cpp | 43 ---------------------------
 1 file changed, 43 deletions(-)

diff --git a/benchmarks/hecbench/dsl/adam/adam.cpp b/benchmarks/hecbench/dsl/adam/adam.cpp
index 72c8f4f..f0160a5 100644
--- a/benchmarks/hecbench/dsl/adam/adam.cpp
+++ b/benchmarks/hecbench/dsl/adam/adam.cpp
@@ -1,13 +1,3 @@
-// NOLINTBEGIN
-
-// clang-format off
-// RUN: rm -rf "%t.$$.proteus"
-// RUN: PROTEUS_CACHE_DIR="%t.$$.proteus" %build/adam_runconst.%ext 10000 200 100 1 | %FILECHECK %s --check-prefixes=CHECK,CHECK-FIRST
-// Second run uses the object cache.
-// RUN: PROTEUS_CACHE_DIR="%t.$$.proteus" %build/adam_runconst.%ext 10000 200 100 1 | %FILECHECK %s --check-prefixes=CHECK,CHECK-SECOND
-// RUN: rm -rf "%t.$$.proteus"
-// clang-format on
-
 #include <proteus/Frontend/Builtins.hpp>
 #include <proteus/JitFrontend.hpp>
 
@@ -204,36 +194,3 @@ int main(int argc, char *argv[]) {
   free(r);
   return 0;
 }
-
-// clang-format off
-// We got slight differences in the output for the least significant digits.
-// Could be HW, numeric, or the way we handle signedness.
-// CHECK: init p[0] = 0.348563
-// CHECK-NEXT: init p[1] = 0.259322
-// CHECK-NEXT: init p[2] = 0.377145
-// CHECK-NEXT: init p[3] = 0.486632
-// CHECK-NEXT: init p[4] = 0.352038
-// CHECK-NEXT: init p[5] = 0.0784863
-// CHECK-NEXT: init p[6] = 0.968732
-// CHECK-NEXT: init p[7] = 0.852707
-// CHECK-NEXT: init p[8] = 0.153431
-// CHECK-NEXT: init p[9] = 0.559506
-// CHECK-NEXT: Creating JIT module
-// CHECK-NEXT: Compiling JIT module
-// CHECK-NEXT: Average kernel execution time {{.*}} (ms)
-// CHECK-NEXT: p[0] = -0.57293
-// CHECK-NEXT: p[1] = -0.59603
-// CHECK-NEXT: p[2] = -0.592634
-// CHECK-NEXT: p[3] = -0.588154
-// CHECK-NEXT: p[4] = -0.593454
-// CHECK-NEXT: p[5] = -0.591989
-// CHECK-NEXT: p[6] = -0.573486
-// CHECK-NEXT: p[7] = -0.599872
-// CHECK-NEXT: p[8] = -0.58157
-// CHECK-NEXT: p[9] = -0.59015{{[7|8]}}
-// CHECK: JitCache hits 0 total 1
-// CHECK: HashValue {{[0-9]+}} NumExecs 1 NumHits 0
-// CHECK-FIRST: JitStorageCache hits 0 total 1
-// CHECK-SECOND: JitStorageCache hits 1 total 1
-
-// NOLINTEND

From a6827100a4c529bf966954056081e05b26bcd2c8 Mon Sep 17 00:00:00 2001
From: Zane Fink <finkzane@gmail.com>
Date: Mon, 20 Oct 2025 16:39:12 -0700
Subject: [PATCH 10/12] Update dsl benchmarks

---
 benchmarks/hecbench/dsl/adam/adam.cpp         | 22 +++---
 benchmarks/hecbench/dsl/attention/main.cpp    | 74 ++++++++++---------
 .../hecbench/dsl/bezier-surface/main.cpp      | 58 +++++++--------
 3 files changed, 77 insertions(+), 77 deletions(-)

diff --git a/benchmarks/hecbench/dsl/adam/adam.cpp b/benchmarks/hecbench/dsl/adam/adam.cpp
index f0160a5..f9e090b 100644
--- a/benchmarks/hecbench/dsl/adam/adam.cpp
+++ b/benchmarks/hecbench/dsl/adam/adam.cpp
@@ -45,11 +45,11 @@ auto createJitModuleSpecial(float _b1, float _b2, float _eps, float _grad_scale,
                                       _time_step, _vector_size, _mode, _decay);
 
     
-    auto &i = F.declVar<size_t>("i");
-    auto &totThreads = F.declVar<size_t>("totThreads");
-    auto &j = F.declVar<size_t>("j");
-    auto &t = F.declVar<int>("t");
-    auto &inc1 = F.defRuntimeConst<int>(1);
+    auto i = F.declVar<size_t>("i");
+    auto totThreads = F.declVar<size_t>("totThreads");
+    auto j = F.declVar<size_t>("j");
+    auto t = F.declVar<int>("t");
+    auto inc1 = F.defRuntimeConst<int>(1);
 
     i = F.callBuiltin(getBlockIdX) * F.callBuiltin(getBlockDimX) +
         F.callBuiltin(getThreadIdX);
@@ -57,23 +57,23 @@ auto createJitModuleSpecial(float _b1, float _b2, float _eps, float _grad_scale,
 
     F.beginFor(j, i, vector_size, totThreads);
     {
-      auto &lim = F.declVar<int>("lim");
+      auto lim = F.declVar<int>("lim");
       t = 1;
       lim = time_step;
       F.beginFor(t, t, lim, inc1);
       {
-        auto &scaled_grad = F.declVar<float>("scale_grad");
+        auto scaled_grad = F.declVar<float>("scale_grad");
         scaled_grad = g[j] / grad_scale;
 
         m[j] = b1 * m[j] + (1.f - b1) * scaled_grad;
         v[j] = b2 * v[j] + (1.f - b2) * scaled_grad * scaled_grad;
 
-        auto &m_corrected = F.declVar<float>("m_corrected");
-        auto &v_corrected = F.declVar<float>("v_corrected");
+        auto m_corrected = F.declVar<float>("m_corrected");
+        auto v_corrected = F.declVar<float>("v_corrected");
         m_corrected = m[j] / (1.f - powf(b1, t));
         v_corrected = v[j] / (1.f - powf(b2, t));
 
-        auto &denom = F.declVar<float>("denom");
+        auto denom = F.declVar<float>("denom");
         F.beginIf(mode == 0);
         { denom = sqrtf(v_corrected + eps); }
         F.endIf();
@@ -82,7 +82,7 @@ auto createJitModuleSpecial(float _b1, float _b2, float _eps, float _grad_scale,
         { denom = sqrtf(v_corrected) + eps; }
         F.endIf();
 
-        auto &update = F.declVar<float>("update");
+        auto update = F.declVar<float>("update");
         update = (m_corrected / denom) + (decay * p[j]);
 
         p[j] -= (step_size * update);
diff --git a/benchmarks/hecbench/dsl/attention/main.cpp b/benchmarks/hecbench/dsl/attention/main.cpp
index d09b969..dc7d017 100644
--- a/benchmarks/hecbench/dsl/attention/main.cpp
+++ b/benchmarks/hecbench/dsl/attention/main.cpp
@@ -63,30 +63,32 @@ static auto getAttentionKernel1(int n, int d) {
 
     F.beginFunction();
     {
-      auto &Tidx = F.callBuiltin(getThreadIdX);
-      auto &Bidx = F.callBuiltin(getBlockIdX);
-      auto &Bdimx = F.callBuiltin(getBlockDimX);
+      auto Tidx = F.callBuiltin(getThreadIdX);
+      auto Bidx = F.callBuiltin(getBlockIdX);
+      auto Bdimx = F.callBuiltin(getBlockDimX);
 
-      auto &i = F.defVar(Bidx * Bdimx + Tidx);
-      auto &Nvar = F.defRuntimeConst(n);
-      auto &Dvar = F.defRuntimeConst(d);
-      auto &Zero = F.defRuntimeConst(0);
-      auto &One = F.defRuntimeConst(1);
+      auto i = F.defVar<int>(Bidx * Bdimx + Tidx);
+      auto Nvar = F.defRuntimeConst<int>(n);
+      auto Dvar = F.defRuntimeConst<int>(d);
+      auto Zero = F.defRuntimeConst<int>(0);
+      auto One = F.defRuntimeConst<int>(1);
 
       F.beginIf(i < Nvar);
       {
-        auto &sum = F.defVar(0.0f);
-        auto &j = F.declVar<int>("j");
+        auto sum = F.defVar<float>(0.0f);
+        auto j = F.declVar<int>("j");
 
-        F.forLoop({j, Zero, Dvar, One}, [&]() {
-          auto &keyIdx = i * Dvar + j;
+        F.beginFor(j, Zero, Dvar, One);
+        {
+          auto keyIdx = i * Dvar + j;
           sum = sum + (key[keyIdx] * query[j]);
-        }).emit();
+        }
+        F.endFor();
 
         dot_product[i] = sum;
 
-        auto &expVal = expf(sum);
-        F.atomicAdd( exp_sum, expVal);
+        auto expVal = expf(sum);
+        F.atomicAdd(exp_sum, F.convert<float>(expVal));
       }
       F.endIf();
 
@@ -110,17 +112,17 @@ static auto getAttentionKernel2(int n) {
 
     F.beginFunction();
     {
-      auto &Tidx = F.callBuiltin(getThreadIdX);
-      auto &Bidx = F.callBuiltin(getBlockIdX);
-      auto &Bdimx = F.callBuiltin(getBlockDimX);
+      auto Tidx = F.callBuiltin(getThreadIdX);
+      auto Bidx = F.callBuiltin(getBlockIdX);
+      auto Bdimx = F.callBuiltin(getBlockDimX);
 
-      auto &i = F.defVar(Bidx * Bdimx + Tidx);
-      auto &Nvar = F.defRuntimeConst(n);
+      auto i = F.defVar<int>(Bidx * Bdimx + Tidx);
+      auto Nvar = F.defRuntimeConst<int>(n);
 
       F.beginIf(i < Nvar);
       {
-        auto &expVal = expf(dot_product[i]);
-        auto &expSumVal = F.defVar(exp_sum[0]);
+        auto expVal = expf(dot_product[i]);
+        auto expSumVal = F.defVar<float>(exp_sum[0]);
         score[i] = expVal / expSumVal;
       }
       F.endIf();
@@ -145,25 +147,27 @@ static auto getAttentionKernel3(int n, int d) {
 
     F.beginFunction();
     {
-      auto &Tidx = F.callBuiltin(getThreadIdX);
-      auto &Bidx = F.callBuiltin(getBlockIdX);
-      auto &Bdimx = F.callBuiltin(getBlockDimX);
+      auto Tidx = F.callBuiltin(getThreadIdX);
+      auto Bidx = F.callBuiltin(getBlockIdX);
+      auto Bdimx = F.callBuiltin(getBlockDimX);
 
-      auto &j = F.defVar(Bidx * Bdimx + Tidx);
-      auto &Dvar = F.defRuntimeConst(d);
-      auto &Nvar = F.defRuntimeConst(n);
-      auto &Zero = F.defRuntimeConst(0);
-      auto &One = F.defRuntimeConst(1);
+      auto j = F.defVar<int>(Bidx * Bdimx + Tidx);
+      auto Dvar = F.defRuntimeConst<int>(d);
+      auto Nvar = F.defRuntimeConst<int>(n);
+      auto Zero = F.defRuntimeConst<int>(0);
+      auto One = F.defRuntimeConst<int>(1);
 
       F.beginIf(j < Dvar);
       {
-        auto &sum = F.defVar(0.0f);
-        auto &i = F.declVar<int>("i");
+        auto sum = F.defVar<float>(0.0f);
+        auto i = F.declVar<int>("i");
 
-        F.forLoop({i, Zero, Nvar, One}, [&]() {
-          auto &valueIdx = i * Dvar + j;
+        F.beginFor(i, Zero, Nvar, One);
+        {
+          auto valueIdx = i * Dvar + j;
           sum = sum + (score[i] * value[valueIdx]);
-        }).emit();
+        }
+        F.endFor();
 
         output[j] = sum;
       }
diff --git a/benchmarks/hecbench/dsl/bezier-surface/main.cpp b/benchmarks/hecbench/dsl/bezier-surface/main.cpp
index 5b4c5fb..c57a4ab 100644
--- a/benchmarks/hecbench/dsl/bezier-surface/main.cpp
+++ b/benchmarks/hecbench/dsl/bezier-surface/main.cpp
@@ -210,8 +210,7 @@ auto createJitModule(int _NI, int _NJ, int _RESOLUTIONI, int _RESOLUTIONJ) {
       auto [NI, NJ, RESOLUTIONI, RESOLUTIONJ] =
           F.defRuntimeConsts(_NI, _NJ, _RESOLUTIONI, _RESOLUTIONJ);
 
-      auto &i = F.declVar<int>();
-      i = F.callBuiltin(getBlockDimX) * F.callBuiltin(getBlockIdX) + F.callBuiltin(getThreadIdX);
+      auto i = F.callBuiltin(getBlockDimX) * F.callBuiltin(getBlockIdX) + F.callBuiltin(getThreadIdX);
 
       F.beginIf(i >= RESOLUTIONI);
       {
@@ -219,45 +218,44 @@ auto createJitModule(int _NI, int _NJ, int _RESOLUTIONI, int _RESOLUTIONJ) {
       }
       F.endIf();
 
-      auto &mui = F.declVar<float>();
-      mui = F.convert<float>(i) / F.convert<float>(RESOLUTIONI - F.defRuntimeConst<int>(1));
+      auto mui = F.convert<float>(i) / F.convert<float>(RESOLUTIONI - F.defRuntimeConst<int>(1));
 
-      auto &j = F.declVar<int>();
-      auto &InitJ = F.defRuntimeConst<int>(0);
-      auto &IncJ = F.defRuntimeConst<int>(1);
+      auto j = F.declVar<int>();
+      auto InitJ = F.defRuntimeConst<int>(0);
+      auto IncJ = F.defRuntimeConst<int>(1);
       F.beginFor(j, InitJ, RESOLUTIONJ, IncJ);
       {
-        auto &muj = F.convert<float>(j) / F.convert<float>(RESOLUTIONJ - F.defRuntimeConst<int>(1));
+        auto muj = F.convert<float>(j) / F.convert<float>(RESOLUTIONJ - F.defRuntimeConst<int>(1));
 
-        auto &OutX = F.defVar<float>(0.0f);
-        auto &OutY = F.defVar<float>(0.0f);
-        auto &OutZ = F.defVar<float>(0.0f);
+        auto OutX = F.defVar<float>(0.0f);
+        auto OutY = F.defVar<float>(0.0f);
+        auto OutZ = F.defVar<float>(0.0f);
 
-        auto &ki = F.declVar<int>();
-        auto &InitKi = F.defRuntimeConst<int>(0);
-        auto &UpperKi = NI + F.defRuntimeConst<int>(1);
-        auto &IncKi = F.defRuntimeConst<int>(1);
+        auto ki = F.declVar<int>();
+        auto InitKi = F.defRuntimeConst<int>(0);
+        auto UpperKi = NI + F.defRuntimeConst<int>(1);
+        auto IncKi = F.defRuntimeConst<int>(1);
         F.beginFor(ki, InitKi, UpperKi, IncKi);
         {
           // float bi = BezierBlend(ki, mui, NI);
-          auto &bi = F.call<float(int, float, int)>("BezierBlend", ki, mui, NI);
+          auto bi = F.call<float(int, float, int)>("BezierBlend", ki, mui, NI);
 
           // for(int kj = 0; kj <= NJ; kj++)
-          auto &kj = F.declVar<int>();
-          auto &InitKj = F.defRuntimeConst<int>(0);
-          auto &UpperKj = NJ + F.defRuntimeConst<int>(1);
-          auto &IncKj = F.defRuntimeConst<int>(1);
+          auto kj = F.declVar<int>();
+          auto InitKj = F.defRuntimeConst<int>(0);
+          auto UpperKj = NJ + F.defRuntimeConst<int>(1);
+          auto IncKj = F.defRuntimeConst<int>(1);
           F.beginFor(kj, InitKj, UpperKj, IncKj);
           {
             // float bj = BezierBlend(kj, muj, NJ);
-            auto &bj = F.call<float(int, float, int)>("BezierBlend", kj, muj, NJ);
+            auto bj = F.call<float(int, float, int)>("BezierBlend", kj, muj, NJ);
 
             // int idx = (ki * (NJ + 1) + kj) * 3;
-            auto &idx = F.declVar<int>();
+            auto idx = F.declVar<int>();
             idx = (ki * (NJ + F.defRuntimeConst<int>(1)) + kj) * F.defRuntimeConst<int>(3);
 
             // float coeff = bi * bj;
-            auto &coeff = F.declVar<float>();
+            auto coeff = F.declVar<float>();
             coeff = bi * bj;
 
             // out_x += inp[idx + 0] * coeff;
@@ -272,7 +270,7 @@ auto createJitModule(int _NI, int _NJ, int _RESOLUTIONI, int _RESOLUTIONJ) {
         F.endFor();
 
         // int out_idx = (i * RESOLUTIONJ + j) * 3;
-        auto &OutIdx = F.declVar<int>();
+        auto OutIdx = F.declVar<int>();
         OutIdx = (i * RESOLUTIONJ + j) * F.defRuntimeConst<int>(3);
 
         // outp[out_idx + 0] = out_x;
@@ -294,16 +292,15 @@ auto createJitModule(int _NI, int _NJ, int _RESOLUTIONI, int _RESOLUTIONJ) {
     F.beginFunction();
     {
       auto [k, mu, n] = F.getArgs();
-      auto &blend = F.defVar<float>(1.0f);
-      auto &nn = F.declVar<int>();
+      auto blend = F.defVar<float>(1.0f);
+      auto nn = F.declVar<int>();
       nn = n;
-      auto &kn = F.declVar<int>();
+      auto kn = F.declVar<int>();
       kn = k;
-      auto &nkn = F.declVar<int>();
+      auto nkn = F.declVar<int>();
       nkn = n - k;
 
-      auto &Cond = nn >= F.defRuntimeConst<int>(1);
-      F.beginWhile(Cond);
+      F.beginWhile([&]() { return nn >= 1; });
       {
         blend *= F.convert<float>(nn);
         nn -= F.defRuntimeConst<int>(1);
@@ -322,7 +319,6 @@ auto createJitModule(int _NI, int _NJ, int _RESOLUTIONI, int _RESOLUTIONJ) {
         }
         F.endIf();
         
-        Cond = nn >= F.defRuntimeConst<int>(1);
       }
       F.endWhile();
 

From 52a6d52ae87ad7f2b9d17916461b9c730c0402df Mon Sep 17 00:00:00 2001
From: Zane Fink <finkzane@gmail.com>
Date: Mon, 20 Oct 2025 17:03:19 -0700
Subject: [PATCH 11/12] Add minibude, 3mm

---
 benchmarks/hecbench/dsl/3mm/Makefile      |  57 ++
 benchmarks/hecbench/dsl/3mm/main.cpp      | 424 ++++++++++++
 benchmarks/hecbench/dsl/minibude/Makefile |  67 ++
 benchmarks/hecbench/dsl/minibude/bude.h   |  40 ++
 benchmarks/hecbench/dsl/minibude/main.cpp | 783 ++++++++++++++++++++++
 hecbench.toml                             |  20 +
 6 files changed, 1391 insertions(+)
 create mode 100644 benchmarks/hecbench/dsl/3mm/Makefile
 create mode 100644 benchmarks/hecbench/dsl/3mm/main.cpp
 create mode 100644 benchmarks/hecbench/dsl/minibude/Makefile
 create mode 100644 benchmarks/hecbench/dsl/minibude/bude.h
 create mode 100644 benchmarks/hecbench/dsl/minibude/main.cpp

diff --git a/benchmarks/hecbench/dsl/3mm/Makefile b/benchmarks/hecbench/dsl/3mm/Makefile
new file mode 100644
index 0000000..f82ee4a
--- /dev/null
+++ b/benchmarks/hecbench/dsl/3mm/Makefile
@@ -0,0 +1,57 @@
+#===============================================================================
+# User Options
+#===============================================================================
+
+# Compiler can be set below, or via environment variable
+CC        = ${PROTEUS_CC}
+OPTIMIZE  = yes
+DEBUG     = no
+PROTEUS_PATH  ?= /path/to/proteus/install
+
+#===============================================================================
+# Program name & source code list
+#===============================================================================
+
+SUFFIX = -proteus
+
+program = 3mm$(SUFFIX).x
+
+source = main.cpp
+obj = $(source:.cpp=$(SUFFIX).o)
+
+#===============================================================================
+# Sets Flags
+#===============================================================================
+
+OFFLOAD_ARCH ?= gfx942
+HIPFLAGS := --offload-arch=${OFFLOAD_ARCH}
+CFLAGS := $(EXTRA_CFLAGS) -std=c++17 -Wall -DPROTEUS_ENABLE_HIP -I${PROTEUS_PATH}/include -I${ROCM_PATH}/llvm/include/ -I${ROCM_PATH}/include $(shell ${ROCM_PATH}/llvm/bin/llvm-config --cxxflags) -fexceptions $(HIPFLAGS)
+
+LDFLAGS = -L${ROCM_PATH}/lib -L${ROCM_PATH}/llvm/lib 	$(shell ${ROCM_PATH}/llvm/bin/llvm-config --libs) 	$(shell ${ROCM_PATH}/llvm/bin/llvm-config --system-libs) 	-llldCommon -llldELF -lamdhip64 -lhiprtc -lhiprtc-builtins -Wl,-rpath,${ROCM_PATH}/lib
+
+ifeq ($(DEBUG),yes)
+	CFLAGS += -g
+	LDFLAGS  += -g
+endif
+
+ifeq ($(OPTIMIZE),yes)
+	CFLAGS += -O3
+endif
+
+LDFLAGS += -Wl,-rpath,${PROTEUS_PATH}/lib64 -L${PROTEUS_PATH}/lib64/ -lproteus
+
+#===============================================================================
+# Targets to Build
+#===============================================================================
+
+$(program): $(obj) Makefile
+	$(CC) $(CFLAGS) $(obj) -o $@ $(LDFLAGS)
+
+%$(SUFFIX).o: %.cpp Makefile
+	$(CC) $(CFLAGS) -x hip -c $< -o $@
+
+clean:
+	rm -rf *.x *.o .proteus
+
+run: $(program)
+	./$(program) --no-verify --N 8192 --kernel jit_regtiled
diff --git a/benchmarks/hecbench/dsl/3mm/main.cpp b/benchmarks/hecbench/dsl/3mm/main.cpp
new file mode 100644
index 0000000..a1995ba
--- /dev/null
+++ b/benchmarks/hecbench/dsl/3mm/main.cpp
@@ -0,0 +1,424 @@
+#include <proteus/JitFrontend.hpp>
+#include <proteus/Frontend/Builtins.hpp>
+#include <proteus/JitInterface.hpp>
+
+#include <cstdlib>
+#include <cstring>
+#include <cctype>
+#include <string>
+#include <iostream>
+#include <memory>
+#include <utility>
+#include <chrono>
+
+using namespace proteus;
+using namespace builtins::gpu;
+
+#if PROTEUS_ENABLE_HIP
+#define TARGET "hip"
+#include <hip/hip_runtime.h>
+#elif PROTEUS_ENABLE_CUDA
+#define TARGET "cuda"
+#include <cuda_runtime.h>
+#define hipError_t cudaError_t
+#define hipSuccess cudaSuccess
+#define hipMalloc cudaMalloc
+#define hipFree cudaFree
+#define hipMemcpy cudaMemcpy
+#define hipMemcpyHostToDevice cudaMemcpyHostToDevice
+#define hipMemcpyDeviceToHost cudaMemcpyDeviceToHost
+#define hipDeviceSynchronize cudaDeviceSynchronize
+#define hipGetErrorString cudaGetErrorString
+#else
+#error "Expected PROTEUS_ENABLE_HIP or PROTEUS_ENABLE_CUDA defined"
+#endif
+
+
+// Naive tiled JIT kernel (C = A x B) for square N x N, double.
+static auto getMatmulKernel(int N, int TileSize) {
+  auto JitMod = std::make_unique<JitModule>(TARGET);
+  auto KernelHandle = JitMod->addKernel<void(double *, double *, double *)>("tiled_matmul");
+  auto &F = KernelHandle.F;
+  {
+    auto Args = F.getArgs();
+    auto &C = std::get<0>(Args);
+    auto &A = std::get<1>(Args);
+    auto &B = std::get<2>(Args);
+
+    F.beginFunction();
+    {
+      auto Tidx = F.callBuiltin(getThreadIdX);
+      auto Bidx = F.callBuiltin(getBlockIdX);
+      auto Tidy = F.callBuiltin(getThreadIdY);
+      auto Bidy = F.callBuiltin(getBlockIdY);
+
+      auto Row = Bidy * TileSize + Tidy;
+      auto Col = Bidx * TileSize + Tidx;
+
+      auto K = F.declVar<int>("K");
+      auto Zero = F.defRuntimeConst(0);
+      auto One = F.defRuntimeConst(1);
+      auto Nvar = F.defRuntimeConst(N);
+      auto Accum = F.defVar(0.0);
+
+      F.beginFor(K, Zero, Nvar, One);
+      {
+        auto AVal = A[Row * N + K];
+        auto BVal = B[K * N + Col];
+        Accum = Accum + (AVal * BVal);
+      }
+      F.endFor();
+
+      auto CIdx = Row * N + Col;
+      C[CIdx] = Accum;
+
+      F.ret();
+    }
+    F.endFunction();
+  }
+  return std::make_pair(std::move(JitMod), KernelHandle);
+}
+
+// Register + shared-memory tiled JIT kernel (C = A x B) for square N x N, double.
+// Configuration mirrors the HIP kernel: 64x64 block tile, 16x16 threads,
+// 4x4 per-thread micro-tile, K tile = 8 (defaults; overridable via CLI).
+static auto getRegSharedTiledMatmulKernel(int N, int blockTileM, int blockTileN, int kTile, int regTileM, int regTileN) {
+  auto JitMod = std::make_unique<JitModule>(TARGET);
+  auto KernelHandle = JitMod->addKernel<void(double *, double *, double *)>("reg_shared_tiled_matmul");
+  auto &F = KernelHandle.F;
+  {
+    auto Args = F.getArgs();
+    auto &C = std::get<0>(Args);
+    auto &A = std::get<1>(Args);
+    auto &B = std::get<2>(Args);
+
+    auto AsTile = F.declVar<double[]>(blockTileM * kTile, AddressSpace::SHARED);
+    auto BsTile = F.declVar<double[]>(kTile * blockTileN, AddressSpace::SHARED);
+
+    auto Areg = F.declVar<double[]>(regTileM);
+    auto Breg = F.declVar<double[]>(regTileN);
+    auto Creg = F.declVar<double[]>(regTileM * regTileN);
+
+    F.beginFunction();
+    {
+      auto Tidx = F.callBuiltin(getThreadIdX);
+      auto Tidy = F.callBuiltin(getThreadIdY);
+      auto Bidx = F.callBuiltin(getBlockIdX);
+      auto Bidy = F.callBuiltin(getBlockIdY);
+
+      auto Nvar = F.defRuntimeConst(N);
+      auto Two = F.defRuntimeConst(2);
+      auto Zero = F.defRuntimeConst(0);
+      auto One = F.defRuntimeConst(1);
+      auto RegTileM = F.defRuntimeConst(regTileM);
+      auto RegTileN = F.defRuntimeConst(regTileN);
+      auto BlockTileM = F.defRuntimeConst(blockTileM);
+      auto BlockTileN = F.defRuntimeConst(blockTileN);
+      auto KTile = F.defRuntimeConst(kTile);
+      auto ThreadsX = F.defRuntimeConst(blockTileN / regTileN);
+
+      auto BlockRow = Bidy * BlockTileM;
+      auto BlockCol = Bidx * BlockTileN;
+
+      auto Row0 = BlockRow + Tidy * RegTileM;
+      auto Col0 = BlockCol + Tidx * RegTileN;
+
+      {
+        auto I = F.declVar<int>("i");
+        auto J = F.declVar<int>("j");
+        F.forLoop<int>({I, Zero, RegTileM, One}, [&]() {
+          F.forLoop<int>({J, Zero, RegTileN, One}, [&]() {
+            auto Cidx = I * RegTileN + J;
+            Creg[Cidx] = 0.0;
+          }).emit();
+        }).emit();
+      }
+
+      auto KBase = F.declVar<int>("KBase");
+      F.forLoop<int>({KBase, Zero, Nvar, KTile}, [&]() {
+        auto Tid = Tidy * ThreadsX + Tidx;
+
+        auto APlaneSize = BlockTileM * KTile;
+        auto AIdx1 = Tid * Two + One;
+        auto AIdx0 = Tid * Two + Zero;
+        auto ARow0 = AIdx0 / KTile;
+        auto ACol0 = AIdx0 % KTile;
+        auto ARow1 = AIdx1 / KTile;
+        auto ACol1 = AIdx1 % KTile;
+        auto AsIdx0 = ARow0 * KTile + ACol0;
+        auto AsIdx1 = ARow1 * KTile + ACol1;
+        auto AGlobIdx0 = (BlockRow + ARow0) * N + (KBase + ACol0);
+        auto AGlobIdx1 = (BlockRow + ARow1) * N + (KBase + ACol1);
+        F.beginIf(AIdx0 < APlaneSize);
+        {
+          AsTile[AsIdx0] = A[AGlobIdx0];
+        }
+        F.endIf();
+        F.beginIf(AIdx1 < APlaneSize);
+        {
+          AsTile[AsIdx1] = A[AGlobIdx1];
+        }
+        F.endIf();
+
+        auto BPlaneSize = KTile * BlockTileN;
+        auto BIdx0 = Tid * Two + Zero;
+        auto BIdx1 = Tid * Two + One;
+        auto BRow0 = BIdx0 / BlockTileN;
+        auto BCol0 = BIdx0 % BlockTileN;
+        auto BRow1 = BIdx1 / BlockTileN;
+        auto BCol1 = BIdx1 % BlockTileN;
+        auto BsIdx0 = BRow0 * BlockTileN + BCol0;
+        auto BsIdx1 = BRow1 * BlockTileN + BCol1;
+        auto BGlobIdx0 = (KBase + BRow0) * N + (BlockCol + BCol0);
+        auto BGlobIdx1 = (KBase + BRow1) * N + (BlockCol + BCol1);
+        F.beginIf(BIdx0 < BPlaneSize);
+        {
+          BsTile[BsIdx0] = B[BGlobIdx0];
+        }
+        F.endIf();
+        F.beginIf(BIdx1 < BPlaneSize);
+        {
+          BsTile[BsIdx1] = B[BGlobIdx1];
+        }
+        F.endIf();
+
+        F.callBuiltin(syncThreads);
+
+        auto KIt = F.declVar<int>("KIt");
+        F.forLoop<int>({KIt, Zero, KTile, One}, [&]() {
+          auto I = F.declVar<int>("i");
+          auto J = F.declVar<int>("j");
+
+          F.forLoop<int>({I, Zero, RegTileM, One}, [&]() {
+            auto r = Tidy * RegTileM + I;
+            auto asIdx = r * KTile + KIt;
+            Areg[I] = AsTile[asIdx];
+          }).emit();
+
+          F.forLoop<int>({J, Zero, RegTileN, One}, [&]() {
+            auto c = Tidx * RegTileN + J;
+            auto bsIdx = KIt * BlockTileN + c;
+            Breg[J] = BsTile[bsIdx];
+          }).emit();
+
+          auto Ii = F.declVar<int>("ii");
+          auto Jj = F.declVar<int>("jj");
+          F.forLoop<int>({Ii, Zero, RegTileM, One}, [&]() {
+            F.forLoop<int>({Jj, Zero, RegTileN, One}, [&]() {
+              auto Cidx = Ii * RegTileN + Jj;
+              Creg[Cidx] = Creg[Cidx] + (Areg[Ii] * Breg[Jj]);
+            }).emit();
+          }).emit();
+        }).emit();
+
+        F.callBuiltin(syncThreads);
+      }).emit();
+
+      {
+        auto I = F.declVar<int>("i");
+        auto J = F.declVar<int>("j");
+        F.forLoop<int>({I, Zero, RegTileM, One}, [&]() {
+          F.forLoop<int>({J, Zero, RegTileN, One}, [&]() {
+            auto Cidx = (Row0 + I) * N + (Col0 + J);
+            auto Ridx = I * RegTileN + J;
+            C[Cidx] = Creg[Ridx];
+          }).emit();
+        }).emit();
+      }
+
+      F.ret();
+    }
+    F.endFunction();
+  }
+  return std::make_pair(std::move(JitMod), KernelHandle);
+}
+
+static bool verifyG(double *G, int N) {
+  double expected = static_cast<double>(N) * static_cast<double>(N) * static_cast<double>(N);
+  for (int i = 0; i < N * N; ++i) {
+    if (G[i] != expected) return false;
+  }
+  return true;
+}
+
+int main(int argc, char** argv) {
+  proteus::init();
+
+  constexpr int MatmulTileSize = 32;
+  constexpr int RegTileM = 4;
+  constexpr int RegTileN = 4;
+  constexpr int BlockTileM = 64;
+  constexpr int BlockTileN = 64;
+  constexpr int KTile = 8;
+
+  unsigned int N = 8192;
+  int NumTrials = 5;
+  bool DoVerify = true;
+  std::string KernelType = "jit_regtiled";
+  int blockTileMArg = BlockTileM;
+  int blockTileNArg = BlockTileN;
+  int kTileArg = KTile;
+  int posIdx = 0;
+
+  for (int i = 1; i < argc; ++i) {
+    if (!std::strcmp(argv[i], "--N") || !std::strcmp(argv[i], "-n")) {
+      if (i + 1 < argc) N = static_cast<unsigned int>(std::atoi(argv[++i]));
+    } else if (!std::strcmp(argv[i], "--trials") || !std::strcmp(argv[i], "-t")) {
+      if (i + 1 < argc) NumTrials = std::atoi(argv[++i]);
+    } else if (!std::strcmp(argv[i], "--kernel")) {
+      if (i + 1 < argc) {
+        KernelType = argv[++i];
+        if (KernelType != "jit" && KernelType != "jit_regtiled") {
+          std::cerr << "Error: Invalid kernel type '" << KernelType
+                    << "'. Valid options: jit, jit_regtiled\n";
+          return 1;
+        }
+      }
+    } else if (!std::strcmp(argv[i], "--verify")) {
+      DoVerify = true;
+    } else if (!std::strcmp(argv[i], "--no-verify")) {
+      DoVerify = false;
+    } else if (!std::strcmp(argv[i], "--help") || !std::strcmp(argv[i], "-h")) {
+      std::cout << "Usage: " << argv[0]
+                << " [-n|--N N] [-t|--trials T] [--kernel KERNEL] [--verify|--no-verify]"
+                << " [blockTileM blockTileN kTile]\n"
+                << "  KERNEL: jit, jit_regtiled (default: jit_regtiled)\n"
+                << "  Positional tile sizes are used for the JIT reg-tiled kernel;"
+                << " defaults are " << BlockTileM << " " << BlockTileN << " " << KTile << "\n";
+      return 0;
+    } else {
+      if (argv[i] && std::isdigit(static_cast<unsigned char>(argv[i][0]))) {
+        int val = std::atoi(argv[i]);
+        if (posIdx == 0) blockTileMArg = val;
+        else if (posIdx == 1) blockTileNArg = val;
+        else if (posIdx == 2) kTileArg = val;
+        ++posIdx;
+      }
+    }
+  }
+
+  std::cout << "3mm PJ-DSL-GPU: N=" << N
+            << ", trials=" << NumTrials
+            << ", verify=" << DoVerify
+            << ", kernel=" << KernelType
+            << ", tiles=(blkM=" << blockTileMArg
+            << ", blkN=" << blockTileNArg
+            << ", k=" << kTileArg << ")" << std::endl;
+
+  // Host allocations
+  double *AH = (double *)new double[N * N];
+  double *BH = (double *)new double[N * N];
+  double *CH = (double *)new double[N * N];
+  double *DH = (double *)new double[N * N];
+  double *EH = (double *)new double[N * N];
+  double *FH = (double *)new double[N * N];
+  double *GH = (double *)new double[N * N];
+
+  if (DoVerify) {
+    for (unsigned int i = 0; i < N * N; ++i) {
+      AH[i] = 1.0; BH[i] = 1.0; CH[i] = 1.0; DH[i] = 1.0;
+      EH[i] = 0.0; FH[i] = 0.0; GH[i] = 0.0;
+    }
+  }
+
+  // Device allocations
+  double *AD, *BD, *CD, *DD, *ED, *FD, *GD;
+  size_t Bytes = sizeof(double) * N * N;
+  hipMalloc(reinterpret_cast<void **>(&AD), Bytes);
+  hipMalloc(reinterpret_cast<void **>(&BD), Bytes);
+  hipMalloc(reinterpret_cast<void **>(&CD), Bytes);
+  hipMalloc(reinterpret_cast<void **>(&DD), Bytes);
+  hipMalloc(reinterpret_cast<void **>(&ED), Bytes);
+  hipMalloc(reinterpret_cast<void **>(&FD), Bytes);
+  hipMalloc(reinterpret_cast<void **>(&GD), Bytes);
+
+  // Stage inputs
+  hipMemcpy(AD, AH, Bytes, hipMemcpyHostToDevice);
+  hipMemcpy(BD, BH, Bytes, hipMemcpyHostToDevice);
+  hipMemcpy(CD, CH, Bytes, hipMemcpyHostToDevice);
+  hipMemcpy(DD, DH, Bytes, hipMemcpyHostToDevice);
+
+
+  if (KernelType == "jit_regtiled") {
+    auto [JitMod, KernelHandle] = getRegSharedTiledMatmulKernel(N, blockTileMArg, blockTileNArg, kTileArg, RegTileM, RegTileN);
+    JitMod->compile(true);
+
+    unsigned int gridX = static_cast<unsigned int>(N / blockTileNArg);
+    unsigned int gridY = static_cast<unsigned int>(N / blockTileMArg);
+    unsigned int blockX = static_cast<unsigned int>(blockTileNArg / RegTileN);
+    unsigned int blockY = static_cast<unsigned int>(blockTileMArg / RegTileM);
+
+    KernelHandle.launch({gridX, gridY, 1u}, {blockX, blockY, 1u}, 0, nullptr, ED, AD, BD);
+    KernelHandle.launch({gridX, gridY, 1u}, {blockX, blockY, 1u}, 0, nullptr, FD, CD, DD);
+    KernelHandle.launch({gridX, gridY, 1u}, {blockX, blockY, 1u}, 0, nullptr, GD, ED, FD);
+    hipDeviceSynchronize();
+
+    double TotalMs = 0.0;
+    auto Start = std::chrono::high_resolution_clock::now();
+    for (int t = 0; t < NumTrials; ++t) {
+      KernelHandle.launch({gridX, gridY, 1u}, {blockX, blockY, 1u}, 0, nullptr, ED, AD, BD);
+      KernelHandle.launch({gridX, gridY, 1u}, {blockX, blockY, 1u}, 0, nullptr, FD, CD, DD);
+      KernelHandle.launch({gridX, gridY, 1u}, {blockX, blockY, 1u}, 0, nullptr, GD, ED, FD);
+    }
+    hipDeviceSynchronize();
+
+    auto End = std::chrono::high_resolution_clock::now();
+    std::chrono::duration<double, std::milli> Ms = End - Start;
+    TotalMs = Ms.count();
+    double AvgMs = TotalMs / static_cast<double>(NumTrials);
+    std::cerr.setf(std::ios::fixed);
+    std::cerr.precision(3);
+    std::cerr << "Average over " << NumTrials << " trials (3 GEMMs): " << AvgMs << " ms" << '\n';
+
+  } else {
+    auto [JitMod, KernelHandle] = getMatmulKernel(N, MatmulTileSize);
+    JitMod->compile(true);
+
+    unsigned int gridX = static_cast<unsigned int>((N + MatmulTileSize - 1) / MatmulTileSize);
+    unsigned int gridY = static_cast<unsigned int>((N + MatmulTileSize - 1) / MatmulTileSize);
+    unsigned int blockX = static_cast<unsigned int>(MatmulTileSize);
+    unsigned int blockY = static_cast<unsigned int>(MatmulTileSize);
+
+    KernelHandle.launch({gridX, gridY, 1u}, {blockX, blockY, 1u}, 0, nullptr, ED, AD, BD);
+    KernelHandle.launch({gridX, gridY, 1u}, {blockX, blockY, 1u}, 0, nullptr, FD, CD, DD);
+    KernelHandle.launch({gridX, gridY, 1u}, {blockX, blockY, 1u}, 0, nullptr, GD, ED, FD);
+    hipDeviceSynchronize();
+
+    double TotalMs = 0.0;
+    auto Start = std::chrono::high_resolution_clock::now();
+    for (int t = 0; t < NumTrials; ++t) {
+      KernelHandle.launch({gridX, gridY, 1u}, {blockX, blockY, 1u}, 0, nullptr, ED, AD, BD);
+      KernelHandle.launch({gridX, gridY, 1u}, {blockX, blockY, 1u}, 0, nullptr, FD, CD, DD);
+      KernelHandle.launch({gridX, gridY, 1u}, {blockX, blockY, 1u}, 0, nullptr, GD, ED, FD);
+    }
+    hipDeviceSynchronize();
+
+    auto End = std::chrono::high_resolution_clock::now();
+    std::chrono::duration<double, std::milli> Ms = End - Start;
+    TotalMs = Ms.count();
+    double AvgMs = TotalMs / static_cast<double>(NumTrials);
+    std::cerr.setf(std::ios::fixed);
+    std::cerr.precision(3);
+    std::cerr << "Average over " << NumTrials << " trials (3 GEMMs): " << AvgMs << " ms" << '\n';
+  }
+
+  if (DoVerify) {
+    hipMemcpy(GH, GD, Bytes, hipMemcpyDeviceToHost);
+    if (verifyG(GH, N)) std::cout << "Verification passed" << std::endl;
+    else                std::cout << "Verification failed" << std::endl;
+  }
+
+  hipFree(AD);
+  hipFree(BD);
+  hipFree(CD);
+  hipFree(DD);
+  hipFree(ED);
+  hipFree(FD);
+  hipFree(GD);
+
+  delete[] AH; delete[] BH; delete[] CH; delete[] DH;
+  delete[] EH; delete[] FH; delete[] GH;
+
+  proteus::finalize();
+  return 0;
+}
diff --git a/benchmarks/hecbench/dsl/minibude/Makefile b/benchmarks/hecbench/dsl/minibude/Makefile
new file mode 100644
index 0000000..c4b031c
--- /dev/null
+++ b/benchmarks/hecbench/dsl/minibude/Makefile
@@ -0,0 +1,67 @@
+#===============================================================================
+# User Options
+#===============================================================================
+
+# Compiler can be set below, or via environment variable
+CC        = ${PROTEUS_CC}
+OPTIMIZE  = yes
+DEBUG     = no
+PROTEUS_PATH  ?=/path/to/proteus/install
+
+#===============================================================================
+# Program name & source code list
+#===============================================================================
+
+SUFFIX = -proteus
+
+program = minibude$(SUFFIX).x
+
+source = main.cpp
+obj = $(source:.cpp=$(SUFFIX).o)
+
+#===============================================================================
+# Sets Flags
+#===============================================================================
+
+# Standard Flags
+CFLAGS := $(EXTRA_CFLAGS) -std=c++17 -Wall
+
+# Linker Flags
+LDFLAGS =
+
+# Debug Flags
+ifeq ($(DEBUG),yes)
+	CFLAGS += -g
+	LDFLAGS  += -g
+endif
+
+# Optimization Flags
+ifeq ($(OPTIMIZE),yes)
+	CFLAGS += -O3
+endif
+
+CFLAGS += -I${PROTEUS_PATH}/include
+LDFLAGS += -Wl,-rpath,${PROTEUS_PATH}/lib64 -L${PROTEUS_PATH}/lib64/ -lproteus \
+    -L${ROCM_PATH}/lib -L${ROCM_PATH}/llvm/lib \
+    -Wl,--start-group \
+    $(shell ls ${ROCM_PATH}/llvm/lib/libclang*.a) \
+    $(shell ${ROCM_PATH}/llvm/bin/llvm-config --libs) \
+    -Wl,--end-group \
+    $(shell ${ROCM_PATH}/llvm/bin/llvm-config --system-libs) \
+    -llldCommon -llldELF
+
+#===============================================================================
+# Targets to Build
+#===============================================================================
+
+$(program): $(obj) Makefile
+	$(CC) $(CFLAGS) $(obj) -o $@ $(LDFLAGS)
+
+%$(SUFFIX).o: %.cpp Makefile
+	$(CC) $(CFLAGS) -x hip -c $< -o $@
+
+clean:
+	rm -rf *.x *.o *.ll *.bc .proteus
+
+run: $(program)
+	./$(program) --deck ../../data/minibude/bm1 --wgsize 256 --iterations 1
diff --git a/benchmarks/hecbench/dsl/minibude/bude.h b/benchmarks/hecbench/dsl/minibude/bude.h
new file mode 100644
index 0000000..35b3e3f
--- /dev/null
+++ b/benchmarks/hecbench/dsl/minibude/bude.h
@@ -0,0 +1,40 @@
+#pragma once
+
+#include <cstdint>
+#include <string>
+#include <iomanip>
+#include <hip/hip_runtime.h>
+
+
+
+#ifndef DEFAULT_PPWI
+#define DEFAULT_PPWI 1
+#endif
+#ifndef DEFAULT_WGSIZE
+#define DEFAULT_WGSIZE 4
+#endif
+
+#define MAX_PPWI 16
+
+#define DEFAULT_ITERS  8
+#define DEFAULT_NPOSES 65536
+#define REF_NPOSES     65536
+
+#define DATA_DIR          "../../data/minibude/bm1"
+#define FILE_LIGAND       "/ligand.in"
+#define FILE_PROTEIN      "/protein.in"
+#define FILE_FORCEFIELD   "/forcefield.in"
+#define FILE_POSES        "/poses.in"
+#define FILE_REF_ENERGIES "/ref_energies.out"
+
+struct __attribute__((__packed__)) Atom {
+  float x, y, z;
+  int32_t type;
+};
+
+struct __attribute__((__packed__)) FFParams {
+  int32_t hbtype;
+  float radius;
+  float hphb;
+  float elsc;
+};
diff --git a/benchmarks/hecbench/dsl/minibude/main.cpp b/benchmarks/hecbench/dsl/minibude/main.cpp
new file mode 100644
index 0000000..efa5b12
--- /dev/null
+++ b/benchmarks/hecbench/dsl/minibude/main.cpp
@@ -0,0 +1,783 @@
+#include <algorithm>
+#include <array>
+#include <chrono>
+#include <cmath>
+#include <cstdint>
+#include <cstdlib>
+#include <fstream>
+#include <functional>
+#include <iomanip>
+#include <iostream>
+#include <initializer_list>
+#include <limits>
+#include <memory>
+#include <hip/hip_runtime.h>
+#include <stdexcept>
+#include <string>
+#include <tuple>
+#include <utility>
+#include <vector>
+
+#include <proteus/JitFrontend.hpp>
+#include <proteus/Frontend/Builtins.hpp>
+#include <proteus/JitInterface.hpp>
+
+#include "bude.h"
+
+using namespace proteus;
+using namespace builtins::gpu;
+
+#define TARGET "hip"
+
+namespace {
+
+typedef std::chrono::high_resolution_clock::time_point TimePoint;
+
+struct Params {
+  size_t natlig;
+  size_t natpro;
+  size_t ntypes;
+  size_t nposes;
+
+  std::vector<Atom> protein;
+  std::vector<Atom> ligand;
+  std::vector<FFParams> forcefield;
+  std::array<std::vector<float>, 6> poses;
+
+  size_t iterations;
+
+  size_t posesPerWI;
+  size_t wgSize;
+  std::string deckDir;
+
+  friend std::ostream &operator<<(std::ostream &os, const Params &params) {
+    os << "natlig:      " << params.natlig << "\n"
+       << "natpro:      " << params.natpro << "\n"
+       << "ntypes:      " << params.ntypes << "\n"
+       << "nposes:      " << params.nposes << "\n"
+       << "iterations:  " << params.iterations << "\n"
+       << "posesPerWI:  " << params.posesPerWI << "\n"
+       << "wgSize:      " << params.wgSize << "\n";
+    return os;
+  }
+};
+
+static auto buildFastenKernel(size_t posesPerWI_, size_t ntypes_, size_t nposes_, size_t natlig_, size_t natpro_) {
+  auto JitMod = std::make_unique<JitModule>(TARGET);
+  auto KernelHandle = JitMod->addKernel<void(float*,float*,float*,int32_t*,float*,float*,float*,int32_t*,float*,float*,float*,float*,float*,float*,int32_t*, float*, float*, float*, float*)>("fasten_main");
+  auto &F = KernelHandle.F;
+  F.beginFunction();
+  {
+    auto [protein_x, protein_y, protein_z, protein_type,
+          ligand_x, ligand_y, ligand_z, ligand_type,
+          transforms_0, transforms_1, transforms_2, transforms_3, transforms_4, transforms_5,
+          ff_hbtype, ff_radius, ff_hphb, ff_elsc, etotals] = F.getArgs();
+
+    auto lid = F.callBuiltin(getThreadIdX);
+    auto gid = F.callBuiltin(getBlockIdX);
+    auto lrange = F.callBuiltin(getBlockDimX);
+    auto ZERO = F.defRuntimeConst<float>(0.0f);
+    auto QUARTER = F.defRuntimeConst<float>(0.25f);
+    auto HALF = F.defRuntimeConst<float>(0.5f);
+    auto ONE = F.defRuntimeConst<float>(1.0f);
+    auto TWO = F.defRuntimeConst<float>(2.0f);
+    auto FOUR = F.defRuntimeConst<float>(4.0f);
+    auto CNSTNT = F.defRuntimeConst<float>(45.0f);
+
+    // Energy evaluation parameters
+    auto HBTYPE_F = F.defRuntimeConst<int>(70);
+    auto HBTYPE_E = F.defRuntimeConst<int>(69);
+    auto HARDNESS = F.defRuntimeConst<float>(38.0f);
+    auto NPNPDIST = F.defRuntimeConst<float>(5.5f);
+    auto NPPDIST = F.defRuntimeConst<float>(1.0f);
+
+    auto FLT_MAX = F.defRuntimeConst<float>(std::numeric_limits<float>::max());
+
+    auto [posesPerWI, ntypes, nposes, natlig, natpro] = F.defRuntimeConsts(posesPerWI_, ntypes_, nposes_, natlig_, natpro_);
+
+    auto etot = F.declVar<float[]>(posesPerWI_);
+
+    // Positions
+    auto lpos_x = F.declVar<float[]>(posesPerWI_);
+    auto lpos_y = F.declVar<float[]>(posesPerWI_);
+    auto lpos_z = F.declVar<float[]>(posesPerWI_);
+
+    // Transformations
+    auto transform_x = F.declVar<float[]>(posesPerWI_ * 3);
+    auto transform_y = F.declVar<float[]>(posesPerWI_ * 3);
+    auto transform_z = F.declVar<float[]>(posesPerWI_ * 3);
+    auto transform_w = F.declVar<float[]>(posesPerWI_ * 3);
+
+    auto ix = gid * lrange * posesPerWI + lid;
+
+    F.beginIf(ix >= nposes);
+    { ix = nposes - posesPerWI; }
+    F.endIf();
+
+    auto i = F.defVar<size_t>(0);
+    F.beginFor(i, i, posesPerWI, F.defRuntimeConst<size_t>(1));
+    {
+      auto index = ix + i * lrange;
+
+      auto sx = sinf(transforms_0[index]);
+      auto cx = cosf(transforms_0[index]);
+      auto sy = sinf(transforms_1[index]);
+      auto cy = cosf(transforms_1[index]);
+      auto sz = sinf(transforms_2[index]);
+      auto cz = cosf(transforms_2[index]);
+
+      auto base = i * 3;
+      transform_x[base + 0] = cy * cz;
+      transform_y[base + 0] = sx * sy * cz - cx * sz;
+      transform_z[base + 0] = cx * sy * cz + sx * sz;
+      transform_w[base + 0] = transforms_3[index];
+
+      transform_x[base + 1] = cy * sz;
+      transform_y[base + 1] = sx * sy * sz + cx * cz;
+      transform_z[base + 1] = cx * sy * sz - sx * cz;
+      transform_w[base + 1] = transforms_4[index];
+
+      transform_x[base + 2] = -1.0f * sy;
+      transform_y[base + 2] = sx * cy;
+      transform_z[base + 2] = cx * cy;
+      transform_w[base + 2] = transforms_5[index];
+
+      etot[i] = ZERO;
+    }
+    F.endFor();
+
+
+    auto il = F.defVar<size_t>(0);
+    // Loop over ligand atoms
+    F.beginFor(il, il, natlig, F.defRuntimeConst<size_t>(1));
+    {
+      // Load ligand atom data
+      auto l_t = ligand_type[il];
+      auto l_x = ligand_x[il];
+      auto l_y = ligand_y[il];
+      auto l_z = ligand_z[il];
+
+      auto l_hbtype = ff_hbtype[l_t];
+      auto l_radius = ff_radius[l_t];
+      auto l_hphb = ff_hphb[l_t];
+      auto l_elsc = ff_elsc[l_t];
+
+      auto lhphb_ltz = l_hphb < ZERO;
+      auto lhphb_gtz = l_hphb > ZERO;
+
+      i = 0;
+      F.beginFor(i, i, posesPerWI, F.defRuntimeConst<size_t>(1));
+      {
+        auto base = i * 3;
+        lpos_x[i] = transform_w[base + 0] + l_x * transform_x[base + 0] + l_y * transform_y[base + 0] + l_z * transform_z[base + 0];
+        lpos_y[i] = transform_w[base + 1] + l_x * transform_x[base + 1] + l_y * transform_y[base + 1] + l_z * transform_z[base + 1];
+        lpos_z[i] = transform_w[base + 2] + l_x * transform_x[base + 2] + l_y * transform_y[base + 2] + l_z * transform_z[base + 2];
+      }
+      F.endFor();
+      auto ip = F.defVar<size_t>(0);
+      // Loop over protein atoms
+      F.beginFor(ip, ip, natpro, F.defRuntimeConst<size_t>(1));
+      {
+        // Load protein atom data
+        auto p_t = protein_type[ip];
+        auto p_x = protein_x[ip];
+        auto p_y = protein_y[ip];
+        auto p_z = protein_z[ip];
+
+        auto p_hbtype = ff_hbtype[p_t];
+        auto p_radius = ff_radius[p_t];
+        auto p_hphb = ff_hphb[p_t];
+        auto p_elsc = ff_elsc[p_t];
+
+        auto radij = p_radius + l_radius;
+        auto r_radij = 1.f / (radij);
+
+        auto elcdst = F.defVar<float>(TWO);
+        F.beginIf(p_hbtype == HBTYPE_F);
+        {
+          F.beginIf(l_hbtype == HBTYPE_F);
+          { elcdst = FOUR; }
+          F.endIf();
+        }
+        F.endIf();
+        auto elcdst1 = F.defVar<float>(HALF);
+        F.beginIf(p_hbtype == HBTYPE_F);
+        {
+          F.beginIf(l_hbtype == HBTYPE_F);
+          { elcdst1 = QUARTER; }
+          F.endIf();
+        }
+        F.endIf();
+
+        auto type_E = F.defVar<bool>(false);
+        F.beginIf(p_hbtype == HBTYPE_E);
+        { type_E = true; }
+        F.endIf();
+        F.beginIf(l_hbtype == HBTYPE_E);
+        { type_E = true; }
+        F.endIf();
+
+        auto phphb_ltz = p_hphb < ZERO;
+        auto phphb_gtz = p_hphb > ZERO;
+        auto phphb_nz = p_hphb != ZERO;
+
+        auto p_hphb_eff = p_hphb * ONE;
+        F.beginIf(phphb_ltz);
+        {
+          F.beginIf(lhphb_gtz);
+          { p_hphb_eff = p_hphb * (-1.0f * ONE); }
+          F.endIf();
+        }
+        F.endIf();
+        auto l_hphb_eff = l_hphb * ONE;
+        F.beginIf(phphb_gtz);
+        {
+          F.beginIf(lhphb_ltz);
+          { l_hphb_eff = l_hphb * (-1.0f * ONE); }
+          F.endIf();
+        }
+        F.endIf();
+
+        auto distdslv = -1.0f * FLT_MAX;
+        F.beginIf(phphb_ltz);
+        {
+          distdslv = NPPDIST;
+          F.beginIf(lhphb_ltz);
+          { distdslv = NPNPDIST; }
+          F.endIf();
+        }
+        F.endIf();
+        F.beginIf(phphb_ltz == false);
+        {
+          F.beginIf(lhphb_ltz);
+          { distdslv = NPPDIST; }
+          F.endIf();
+        }
+        F.endIf();
+        auto r_distdslv = 1.f / (distdslv);
+
+        auto chrg_init = l_elsc * p_elsc;
+        auto dslv_init = p_hphb_eff + l_hphb_eff;
+
+        i = 0;
+        F.beginFor(i, i, posesPerWI, F.defRuntimeConst<size_t>(1));
+        {
+          auto x = lpos_x[i] - p_x;
+          auto y = lpos_y[i] - p_y;
+          auto z = lpos_z[i] - p_z;
+
+          auto distij = sqrtf(x * x + y * y + z * z);
+
+          // Calculate the sum of the sphere radii
+          auto distbb = distij - radij;
+          auto zone1 = (distbb < ZERO);
+
+          // Calculate steric energy
+          auto steric = ONE - (distij * r_radij);
+          auto multiplier = F.defVar<float>(ZERO);
+          F.beginIf(zone1);
+          {
+            multiplier = TWO * HARDNESS;
+          }
+          F.endIf();
+          etot[i] += steric * multiplier;
+
+          // Calculate formal and dipole charge interactions
+          auto chrg_zone_factor = ONE - distbb * elcdst1;
+          F.beginIf(zone1);
+          { chrg_zone_factor = ONE; }
+          F.endIf();
+          auto chrg_dist_factor = F.defVar<float>(ZERO);
+          F.beginIf(distbb < elcdst);
+          { chrg_dist_factor = ONE; }
+          F.endIf();
+          auto chrg_e = chrg_init * chrg_zone_factor * chrg_dist_factor;
+          auto neg_chrg_e = -1.0f * fabs(chrg_e);
+          F.beginIf(type_E);
+          { chrg_e = neg_chrg_e; }
+          F.endIf();
+          etot[i] += chrg_e * CNSTNT;
+
+          // Calculate the two cases for Nonpolar-Polar repulsive interactions
+          auto coeff = ONE - (distbb * r_distdslv);
+          auto dslv_zone_factor = F.defVar<float>(ZERO);
+          F.beginIf(distbb < distdslv);
+          {
+            F.beginIf(phphb_nz);
+            { dslv_zone_factor = ONE; }
+            F.endIf();
+          }
+          F.endIf();
+          auto dslv_e = dslv_init * dslv_zone_factor;
+          auto dslv_scale = coeff;
+          F.beginIf(zone1);
+          { dslv_scale = ONE; }
+          F.endIf();
+          dslv_e *= dslv_scale;
+          etot[i] += dslv_e;
+
+        }
+        F.endFor();
+    }
+    F.endFor();
+  }
+  F.endFor();
+
+
+  auto td_base = gid * lrange * posesPerWI + lid;
+  F.beginIf(td_base < nposes);
+  {
+      i = 0;
+      F.beginFor(i, i, posesPerWI, F.defRuntimeConst<size_t>(1));
+      {
+      etotals[td_base + i * lrange] = etot[i] * HALF;
+      }
+      F.endFor();
+  }
+  F.endIf();
+  F.ret();
+  }
+  F.endFunction();
+
+  std::cout << "Kernel constructed successfully" << std::endl;
+  return std::make_pair(std::move(JitMod), KernelHandle);
+}
+
+double elapsedMillis(const TimePoint &start, const TimePoint &end) {
+  auto elapsedNs =
+      static_cast<double>(std::chrono::duration_cast<std::chrono::nanoseconds>(
+                              end - start)
+                              .count());
+  return elapsedNs * 1e-6;
+}
+
+void printTimings(const Params &params, double millis) {
+  double ms = (millis / params.iterations);
+  double runtime = ms * 1e-3;
+
+  double ops_per_wg =
+      params.posesPerWI * 27 +
+      params.natlig * (3 + params.posesPerWI * 18 +
+                       params.natpro * (11 + params.posesPerWI * 30)) +
+      params.posesPerWI;
+  double total_ops = ops_per_wg * (static_cast<double>(params.nposes) /
+                                   params.posesPerWI);
+  double flops = total_ops / runtime;
+  double gflops = flops / 1e9;
+
+  double interactions = static_cast<double>(params.nposes) *
+                        static_cast<double>(params.natlig) *
+                        static_cast<double>(params.natpro);
+  double interactions_per_sec = interactions / runtime;
+
+  std::cout.precision(3);
+  std::cout << std::fixed;
+  std::cout << "- Total kernel time:    " << (millis) << " ms\n";
+  std::cout << "- Average kernel time:   " << ms << " ms\n";
+  std::cout << "- Interactions/s: " << (interactions_per_sec / 1e9)
+            << " billion\n";
+  std::cout << "- GFLOP/s:        " << gflops << "\n";
+}
+
+template <typename T> std::vector<T> readNStruct(const std::string &path) {
+  std::fstream s(path, std::ios::binary | std::ios::in);
+  if (!s.good()) {
+    throw std::invalid_argument("Bad file: " + path);
+  }
+  s.ignore(std::numeric_limits<std::streamsize>::max());
+  auto len = s.gcount();
+  s.clear();
+  s.seekg(0, std::ios::beg);
+  std::vector<T> xs(static_cast<size_t>(len) / sizeof(T));
+  s.read(reinterpret_cast<char *>(xs.data()), len);
+  s.close();
+  return xs;
+}
+
+Params loadParameters(const std::vector<std::string> &args) {
+  Params params = {};
+
+  params.iterations = DEFAULT_ITERS;
+  params.nposes = DEFAULT_NPOSES;
+  params.wgSize = DEFAULT_WGSIZE;
+  params.deckDir = DATA_DIR;
+  params.posesPerWI = DEFAULT_PPWI;
+
+  const auto readParam = [&args](size_t &current, const std::string &arg,
+                                 const std::initializer_list<std::string> &matches,
+                                 const std::function<void(std::string)> &handle) {
+    if (matches.size() == 0)
+      return false;
+    if (std::find(matches.begin(), matches.end(), arg) != matches.end()) {
+      if (current + 1 < args.size()) {
+        current++;
+        handle(args[current]);
+      } else {
+        std::cerr << "[";
+        for (const auto &m : matches)
+          std::cerr << m;
+        std::cerr << "] specified but no value was given" << std::endl;
+        std::exit(EXIT_FAILURE);
+      }
+      return true;
+    }
+    return false;
+  };
+
+  const auto bindInt = [](const std::string &param, size_t &dest,
+                          const std::string &name) {
+    try {
+      auto parsed = std::stol(param);
+      if (parsed < 0) {
+        std::cerr << "positive integer required for <" << name
+                  << ">: `" << parsed << "`" << std::endl;
+        std::exit(EXIT_FAILURE);
+      }
+      dest = static_cast<size_t>(parsed);
+    } catch (...) {
+      std::cerr << "malformed value, integer required for <" << name
+                << ">: `" << param << "`" << std::endl;
+      std::exit(EXIT_FAILURE);
+    }
+  };
+
+  for (size_t i = 0; i < args.size(); ++i) {
+    using namespace std::placeholders;
+    const auto arg = args[i];
+    if (readParam(i, arg, {"--iterations", "-i"},
+                  std::bind(bindInt, _1, std::ref(params.iterations),
+                            "iterations")))
+      continue;
+    if (readParam(i, arg, {"--numposes", "-n"},
+                  std::bind(bindInt, _1, std::ref(params.nposes), "numposes")))
+      continue;
+    if (readParam(i, arg, {"--posesperwi", "-p"},
+                  std::bind(bindInt, _1, std::ref(params.posesPerWI),
+                            "posesperwi")))
+      continue;
+    if (readParam(i, arg, {"--wgsize", "-w"},
+                  std::bind(bindInt, _1, std::ref(params.wgSize), "wgsize")))
+      continue;
+    if (readParam(i, arg, {"--deck"},
+                  [&](const std::string &param) { params.deckDir = param; }))
+      continue;
+
+    if (arg == "--help" || arg == "-h") {
+      std::cout << "\n";
+      std::cout << "Usage: ./main [OPTIONS]\n\n"
+                << "Options:\n"
+                << "  -h  --help               Print this message\n"
+                << "  -i  --iterations I       Repeat kernel I times (default: "
+                << DEFAULT_ITERS << ")\n"
+                << "  -n  --numposes   N       Compute energies for N poses (default: "
+                << DEFAULT_NPOSES << ")\n"
+                << "  -p  --posesperwi PPWI    Compute PPWI poses per work-item "
+                   "(default: "
+                << DEFAULT_PPWI << ")\n"
+                << "  -w  --wgsize     WGSIZE  Run with work-group size WGSIZE using "
+                   "nd_range, set to 0 for plain range (default: "
+                << DEFAULT_WGSIZE << ")\n"
+                << "      --deck       DECK    Use the DECK directory as input deck "
+                   "(default: "
+                << DATA_DIR << ")"
+                << std::endl;
+      std::exit(EXIT_SUCCESS);
+    }
+
+    std::cout << "Unrecognized argument '" << arg << "' (try '--help')"
+              << std::endl;
+    std::exit(EXIT_FAILURE);
+  }
+
+  if (params.posesPerWI == 0 || params.posesPerWI > MAX_PPWI) {
+    std::cerr << "posesperwi must be in 1.." << MAX_PPWI << std::endl;
+    std::exit(EXIT_FAILURE);
+  }
+
+  params.ligand = readNStruct<Atom>(params.deckDir + FILE_LIGAND);
+  params.natlig = params.ligand.size();
+
+  params.protein = readNStruct<Atom>(params.deckDir + FILE_PROTEIN);
+  params.natpro = params.protein.size();
+
+  params.forcefield = readNStruct<FFParams>(params.deckDir + FILE_FORCEFIELD);
+  params.ntypes = params.forcefield.size();
+
+  auto poses = readNStruct<float>(params.deckDir + FILE_POSES);
+  if (poses.size() / 6 != params.nposes) {
+    throw std::invalid_argument("Bad poses: " +
+                                std::to_string(poses.size()));
+  }
+
+  for (size_t i = 0; i < 6; ++i) {
+    params.poses[i].resize(params.nposes);
+    std::copy(std::next(poses.cbegin(), i * params.nposes),
+              std::next(poses.cbegin(), i * params.nposes + params.nposes),
+              params.poses[i].begin());
+  }
+
+  return params;
+}
+
+std::vector<float> runKernel(const Params &params) {
+  auto energies = std::vector<float>(params.nposes);
+
+  std::vector<float> h_protein_x(params.natpro), h_protein_y(params.natpro),
+      h_protein_z(params.natpro);
+  std::vector<int32_t> h_protein_type(params.natpro);
+  for (size_t i = 0; i < params.natpro; ++i) {
+    h_protein_x[i] = params.protein[i].x;
+    h_protein_y[i] = params.protein[i].y;
+    h_protein_z[i] = params.protein[i].z;
+    h_protein_type[i] = params.protein[i].type;
+  }
+
+  std::vector<float> h_ligand_x(params.natlig), h_ligand_y(params.natlig),
+      h_ligand_z(params.natlig);
+  std::vector<int32_t> h_ligand_type(params.natlig);
+  for (size_t i = 0; i < params.natlig; ++i) {
+    h_ligand_x[i] = params.ligand[i].x;
+    h_ligand_y[i] = params.ligand[i].y;
+    h_ligand_z[i] = params.ligand[i].z;
+    h_ligand_type[i] = params.ligand[i].type;
+  }
+
+  float *d_protein_x = nullptr;
+  float *d_protein_y = nullptr;
+  float *d_protein_z = nullptr;
+  int32_t *d_protein_type = nullptr;
+  hipMalloc(reinterpret_cast<void **>(&d_protein_x),
+            params.natpro * sizeof(float));
+  hipMalloc(reinterpret_cast<void **>(&d_protein_y),
+            params.natpro * sizeof(float));
+  hipMalloc(reinterpret_cast<void **>(&d_protein_z),
+            params.natpro * sizeof(float));
+  hipMalloc(reinterpret_cast<void **>(&d_protein_type),
+            params.natpro * sizeof(int32_t));
+  hipMemcpy(d_protein_x, h_protein_x.data(),
+            params.natpro * sizeof(float),
+            hipMemcpyHostToDevice);
+  hipMemcpy(d_protein_y, h_protein_y.data(),
+            params.natpro * sizeof(float),
+            hipMemcpyHostToDevice);
+  hipMemcpy(d_protein_z, h_protein_z.data(),
+            params.natpro * sizeof(float),
+            hipMemcpyHostToDevice);
+  hipMemcpy(d_protein_type, h_protein_type.data(),
+            params.natpro * sizeof(int32_t),
+            hipMemcpyHostToDevice);
+
+  float *d_ligand_x = nullptr;
+  float *d_ligand_y = nullptr;
+  float *d_ligand_z = nullptr;
+  int32_t *d_ligand_type = nullptr;
+  hipMalloc(reinterpret_cast<void **>(&d_ligand_x),
+            params.natlig * sizeof(float));
+  hipMalloc(reinterpret_cast<void **>(&d_ligand_y),
+            params.natlig * sizeof(float));
+  hipMalloc(reinterpret_cast<void **>(&d_ligand_z),
+            params.natlig * sizeof(float));
+  hipMalloc(reinterpret_cast<void **>(&d_ligand_type),
+            params.natlig * sizeof(int32_t));
+  hipMemcpy(d_ligand_x, h_ligand_x.data(),
+            params.natlig * sizeof(float),
+            hipMemcpyHostToDevice);
+  hipMemcpy(d_ligand_y, h_ligand_y.data(),
+            params.natlig * sizeof(float),
+            hipMemcpyHostToDevice);
+  hipMemcpy(d_ligand_z, h_ligand_z.data(),
+            params.natlig * sizeof(float),
+            hipMemcpyHostToDevice);
+  hipMemcpy(d_ligand_type, h_ligand_type.data(),
+            params.natlig * sizeof(int32_t),
+            hipMemcpyHostToDevice);
+
+  float *transforms_0 = nullptr;
+  float *transforms_1 = nullptr;
+  float *transforms_2 = nullptr;
+  float *transforms_3 = nullptr;
+  float *transforms_4 = nullptr;
+  float *transforms_5 = nullptr;
+  hipMalloc(reinterpret_cast<void **>(&transforms_0),
+            params.nposes * sizeof(float));
+  hipMalloc(reinterpret_cast<void **>(&transforms_1),
+            params.nposes * sizeof(float));
+  hipMalloc(reinterpret_cast<void **>(&transforms_2),
+            params.nposes * sizeof(float));
+  hipMalloc(reinterpret_cast<void **>(&transforms_3),
+            params.nposes * sizeof(float));
+  hipMalloc(reinterpret_cast<void **>(&transforms_4),
+            params.nposes * sizeof(float));
+  hipMalloc(reinterpret_cast<void **>(&transforms_5),
+            params.nposes * sizeof(float));
+
+  float *transformPtrs[6] = {transforms_0, transforms_1, transforms_2,
+                             transforms_3, transforms_4, transforms_5};
+  for (size_t i = 0; i < 6; ++i) {
+    hipMemcpy(transformPtrs[i], params.poses[i].data(),
+              params.nposes * sizeof(float),
+              hipMemcpyHostToDevice);
+  }
+
+  std::vector<int32_t> h_ff_hbtype(params.ntypes);
+  std::vector<float> h_ff_radius(params.ntypes), h_ff_hphb(params.ntypes),
+      h_ff_elsc(params.ntypes);
+  for (size_t i = 0; i < params.ntypes; ++i) {
+    h_ff_hbtype[i] = params.forcefield[i].hbtype;
+    h_ff_radius[i] = params.forcefield[i].radius;
+    h_ff_hphb[i] = params.forcefield[i].hphb;
+    h_ff_elsc[i] = params.forcefield[i].elsc;
+  }
+
+  int32_t *d_ff_hbtype = nullptr;
+  float *d_ff_radius = nullptr;
+  float *d_ff_hphb = nullptr;
+  float *d_ff_elsc = nullptr;
+  hipMalloc(reinterpret_cast<void **>(&d_ff_hbtype),
+            params.ntypes * sizeof(int32_t));
+  hipMalloc(reinterpret_cast<void **>(&d_ff_radius),
+            params.ntypes * sizeof(float));
+  hipMalloc(reinterpret_cast<void **>(&d_ff_hphb),
+            params.ntypes * sizeof(float));
+  hipMalloc(reinterpret_cast<void **>(&d_ff_elsc),
+            params.ntypes * sizeof(float));
+  hipMemcpy(d_ff_hbtype, h_ff_hbtype.data(),
+            params.ntypes * sizeof(int32_t),
+            hipMemcpyHostToDevice);
+  hipMemcpy(d_ff_radius, h_ff_radius.data(),
+            params.ntypes * sizeof(float),
+            hipMemcpyHostToDevice);
+  hipMemcpy(d_ff_hphb, h_ff_hphb.data(),
+            params.ntypes * sizeof(float),
+            hipMemcpyHostToDevice);
+  hipMemcpy(d_ff_elsc, h_ff_elsc.data(),
+            params.ntypes * sizeof(float),
+            hipMemcpyHostToDevice);
+
+  float *results = nullptr;
+  hipMalloc(reinterpret_cast<void **>(&results),
+            params.nposes * sizeof(float));
+
+  auto [JitMod, KernelHandle] = buildFastenKernel(params.posesPerWI, params.ntypes, params.nposes, params.natlig, params.natpro);
+  JitMod->compile(true);
+  // JitMod->print();
+
+  double global = std::ceil(static_cast<double>(params.nposes) /
+                            static_cast<double>(params.posesPerWI));
+  global = std::ceil(global / static_cast<double>(params.wgSize));
+
+  const unsigned int gridDimX = static_cast<unsigned int>(global);
+  const unsigned int blockDimX = static_cast<unsigned int>(params.wgSize);
+
+
+  KernelHandle.launch({gridDimX, 1u, 1u},
+                      {blockDimX, 1u, 1u},
+                      0, nullptr,
+                      d_protein_x, d_protein_y,
+                      d_protein_z, d_protein_type, d_ligand_x,
+                      d_ligand_y, d_ligand_z, d_ligand_type,
+                      transforms_0, transforms_1, transforms_2,
+                      transforms_3, transforms_4, transforms_5,
+                      d_ff_hbtype, d_ff_radius, d_ff_hphb,
+                      d_ff_elsc, results);
+  hipDeviceSynchronize();
+
+  auto kernelStart = std::chrono::high_resolution_clock::now();
+  for (size_t i = 0; i < params.iterations; ++i) {
+    KernelHandle.launch({gridDimX, 1u, 1u},
+                        {blockDimX, 1u, 1u},
+                        0, nullptr,
+                        d_protein_x, d_protein_y,
+                        d_protein_z, d_protein_type, d_ligand_x,
+                        d_ligand_y, d_ligand_z, d_ligand_type,
+                        transforms_0, transforms_1, transforms_2,
+                        transforms_3, transforms_4, transforms_5,
+                        d_ff_hbtype, d_ff_radius, d_ff_hphb,
+                        d_ff_elsc, results);
+
+  }
+  hipDeviceSynchronize();
+  auto kernelEnd = std::chrono::high_resolution_clock::now();
+
+  hipMemcpy(energies.data(), results,
+            params.nposes * sizeof(float),
+            hipMemcpyDeviceToHost);
+
+  printTimings(params, elapsedMillis(kernelStart, kernelEnd));
+
+  hipFree(d_protein_x);
+  hipFree(d_protein_y);
+  hipFree(d_protein_z);
+  hipFree(d_protein_type);
+  hipFree(d_ligand_x);
+  hipFree(d_ligand_y);
+  hipFree(d_ligand_z);
+  hipFree(d_ligand_type);
+  hipFree(transforms_0);
+  hipFree(transforms_1);
+  hipFree(transforms_2);
+  hipFree(transforms_3);
+  hipFree(transforms_4);
+  hipFree(transforms_5);
+  hipFree(d_ff_hbtype);
+  hipFree(d_ff_radius);
+  hipFree(d_ff_hphb);
+  hipFree(d_ff_elsc);
+  hipFree(results);
+
+  return energies;
+}
+
+} // namespace
+
+int main(int argc, char *argv[]) {
+  proteus::init();
+
+  auto args = std::vector<std::string>(argv + 1, argv + argc);
+  auto params = loadParameters(args);
+
+  std::cout << "Poses     : " << params.nposes << std::endl;
+  std::cout << "Iterations: " << params.iterations << std::endl;
+  std::cout << "Ligands   : " << params.natlig << std::endl;
+  std::cout << "Proteins  : " << params.natpro << std::endl;
+  std::cout << "Deck      : " << params.deckDir << std::endl;
+  std::cout << "WG        : " << params.wgSize << std::endl;
+  auto energies = runKernel(params);
+
+#ifdef DUMP
+  FILE *output = fopen("result.out", "w+");
+  printf("\nEnergies\n");
+  for (size_t i = 0; i < params.nposes; i++) {
+    fprintf(output, "%7.2f\n", energies[i]);
+    if (i < 16)
+      printf("%7.2f\n", energies[i]);
+  }
+  fclose(output);
+#endif
+
+  std::ifstream refEnergies(params.deckDir + FILE_REF_ENERGIES);
+  size_t nRefPoses = params.nposes;
+  if (params.nposes > REF_NPOSES) {
+    std::cout << "Only validating the first " << REF_NPOSES << " poses.\n";
+    nRefPoses = REF_NPOSES;
+  }
+
+  std::string line;
+  float maxdiff = 0.0f;
+  for (size_t i = 0; i < nRefPoses; i++) {
+    if (!std::getline(refEnergies, line)) {
+      throw std::logic_error("ran out of ref energies lines to verify");
+    }
+    float e = std::stof(line);
+    if (std::fabs(e) < 1.f && std::fabs(energies[i]) < 1.f)
+      continue;
+
+    float diff = std::fabs(e - energies[i]) / e;
+    if (diff > maxdiff)
+      maxdiff = diff;
+  }
+  std::cout << "Largest difference was "
+            << std::setprecision(3) << (100 * maxdiff) << "%.\n\n";
+
+  return 0;
+}
diff --git a/hecbench.toml b/hecbench.toml
index 5470d59..e1a392c 100644
--- a/hecbench.toml
+++ b/hecbench.toml
@@ -169,3 +169,23 @@ path = "benchmarks/hecbench/dsl/bezier-surface"
 exe = "bezier-surface-proteus.x"
 [hecbench.bezier-surface.inputs]
 default = "-n 8192"
+
+[hecbench.minibude]
+[hecbench.minibude.amd.dsl]
+path = "benchmarks/hecbench/dsl/minibude"
+exe = "minibude-proteus.x"
+[hecbench.minibude.nvidia.dsl]
+path = "benchmarks/hecbench/dsl/minibude"
+exe = "minibude-proteus.x"
+[hecbench.minibude.inputs]
+default = "--deck ../../data/minibude/bm2 --wgsize 256 --iterations 5"
+
+[hecbench.3mm]
+[hecbench.3mm.amd.dsl]
+path = "benchmarks/hecbench/dsl/3mm"
+exe = "3mm-proteus.x"
+[hecbench.3mm.nvidia.dsl]
+path = "benchmarks/hecbench/dsl/3mm"
+exe = "3mm-proteus.x"
+[hecbench.3mm.inputs]
+default = "--no-verify --N 8192 --kernel jit_regtiled"

From f7292528406576d757369d6990fe8d70e4135d0b Mon Sep 17 00:00:00 2001
From: Zane Fink <finkzane@gmail.com>
Date: Mon, 20 Oct 2025 17:39:31 -0700
Subject: [PATCH 12/12] gemm/floyd-warshal

---
 .../hecbench/dsl/floyd-warshall/Makefile      |  63 +++
 .../hecbench/dsl/floyd-warshall/main.cpp      | 243 +++++++++++
 benchmarks/hecbench/dsl/gemm/Makefile         |  60 +++
 benchmarks/hecbench/dsl/gemm/main.cpp         | 411 ++++++++++++++++++
 hecbench.toml                                 |  20 +
 5 files changed, 797 insertions(+)
 create mode 100644 benchmarks/hecbench/dsl/floyd-warshall/Makefile
 create mode 100644 benchmarks/hecbench/dsl/floyd-warshall/main.cpp
 create mode 100644 benchmarks/hecbench/dsl/gemm/Makefile
 create mode 100644 benchmarks/hecbench/dsl/gemm/main.cpp

diff --git a/benchmarks/hecbench/dsl/floyd-warshall/Makefile b/benchmarks/hecbench/dsl/floyd-warshall/Makefile
new file mode 100644
index 0000000..104960b
--- /dev/null
+++ b/benchmarks/hecbench/dsl/floyd-warshall/Makefile
@@ -0,0 +1,63 @@
+#===============================================================================
+# User Options
+#===============================================================================
+
+CC        = ${PROTEUS_CC}
+OPTIMIZE  = yes
+DEBUG     = no
+PROTEUS_PATH ?= /path/to/proteus/install
+
+#===============================================================================
+# Program name & source code list
+#===============================================================================
+
+SUFFIX = "-proteus"
+
+program = floyd-warshall$(SUFFIX).x
+
+source = main.cpp
+obj = $(source:.cpp=$(SUFFIX).o)
+
+#===============================================================================
+# Sets Flags
+#===============================================================================
+
+CFLAGS := $(EXTRA_CFLAGS) -std=c++17 -Wall
+LDFLAGS =
+
+ifeq ($(DEBUG),yes)
+	CFLAGS += -g
+	LDFLAGS += -g
+endif
+
+ifeq ($(OPTIMIZE),yes)
+	CFLAGS += -O3
+endif
+
+CFLAGS += -I${PROTEUS_PATH}/include
+LDFLAGS += -Wl,-rpath,${PROTEUS_PATH}/lib64 -L${PROTEUS_PATH}/lib64/ -lproteus \
+    -L${ROCM_PATH}/lib -L${ROCM_PATH}/llvm/lib \
+    -Wl,--start-group \
+    $(shell ls ${ROCM_PATH}/llvm/lib/libclang*.a) \
+    $(shell ${ROCM_PATH}/llvm/bin/llvm-config --libs) \
+    -Wl,--end-group \
+    $(shell ${ROCM_PATH}/llvm/bin/llvm-config --system-libs) \
+    -llldCommon -llldELF -lhiprand
+
+#===============================================================================
+# Targets to Build
+#===============================================================================
+
+$(program): $(obj) Makefile
+	$(CC) $(CFLAGS) $(obj) -o $@ $(LDFLAGS)
+
+%$(SUFFIX).o: %.cpp Makefile
+	$(CC) $(CFLAGS) -x hip -c $< -o $@
+
+.PHONY: clean run
+
+clean:
+	rm -rf *.x *.o *.ll *.bc .proteus
+
+run: $(program)
+	./$(program) 1024 100 16 0
diff --git a/benchmarks/hecbench/dsl/floyd-warshall/main.cpp b/benchmarks/hecbench/dsl/floyd-warshall/main.cpp
new file mode 100644
index 0000000..67a0c37
--- /dev/null
+++ b/benchmarks/hecbench/dsl/floyd-warshall/main.cpp
@@ -0,0 +1,243 @@
+#include <proteus/Frontend/Builtins.hpp>
+#include <proteus/JitFrontend.hpp>
+#include <proteus/JitInterface.hpp>
+
+#include <cassert>
+#include <chrono>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <memory>
+
+#include <hip/hip_runtime.h>
+#include <hiprand/hiprand.h>
+
+#if !defined(PROTEUS_ENABLE_HIP) && !defined(PROTEUS_ENABLE_CUDA)
+#define PROTEUS_ENABLE_HIP 1
+#endif
+
+#if PROTEUS_ENABLE_HIP
+#define TARGET "hip"
+#elif PROTEUS_ENABLE_CUDA
+#define TARGET "cuda"
+#else
+#error "Expected PROTEUS_ENABLE_HIP or PROTEUS_ENABLE_CUDA defined"
+#endif
+
+constexpr unsigned int MAXDISTANCE = 200;
+
+using namespace proteus;
+using namespace builtins::gpu;
+
+// Map RNG output to [0, MAXDISTANCE] and zero the diagonal (2D launch)
+extern "C" __global__ void initRandomMatrix2D(unsigned int* __restrict__ buf, const unsigned int numNodes) {
+  unsigned int x = threadIdx.x + blockIdx.x * blockDim.x;
+  unsigned int y = threadIdx.y + blockIdx.y * blockDim.y;
+  if (x >= numNodes || y >= numNodes) return;
+  unsigned int idx = y * numNodes + x;
+  unsigned int v = buf[idx] % (MAXDISTANCE + 1);
+  if (x == y) v = 0u;
+  buf[idx] = v;
+}
+
+// Reference CPU implementation for verification
+static void floydWarshallCPUReference(unsigned int * pathDistanceMatrix,
+                                      unsigned int * pathMatrix,
+                                      unsigned int numNodes)
+{
+  unsigned int width = numNodes;
+  for (unsigned int k = 0; k < numNodes; ++k) {
+    for (unsigned int y = 0; y < numNodes; ++y) {
+      unsigned int yXwidth = y * numNodes;
+      for (unsigned int x = 0; x < numNodes; ++x) {
+        unsigned int distanceYtoX = pathDistanceMatrix[yXwidth + x];
+        unsigned int distanceYtoK = pathDistanceMatrix[yXwidth + k];
+        unsigned int distanceKtoX = pathDistanceMatrix[k * width + x];
+        unsigned int indirectDistance = distanceYtoK + distanceKtoX;
+        if (indirectDistance < distanceYtoX) {
+          pathDistanceMatrix[yXwidth + x] = indirectDistance;
+          pathMatrix[yXwidth + x] = k;
+        }
+      }
+    }
+  }
+}
+
+static auto createJitModuleSpecial(unsigned int _numNodes) {
+  auto J = std::make_unique<JitModule>(TARGET);
+  auto KernelHandle = J->addKernel<void(unsigned int*, unsigned int*, unsigned int, unsigned int)>("floydWarshallPass");
+  auto &F = KernelHandle.F;
+  auto [pathDistanceBuffer, pathBuffer, numNodes, pass] = F.getArgs();
+
+  F.beginFunction();
+  {
+    // Bake only numNodes as a runtime constant; pass is a dynamic kernel arg
+    auto rcNumNodes = F.defRuntimeConst(_numNodes);
+
+    auto xValue = F.declVar<unsigned int>("xValue");
+    auto yValue = F.declVar<unsigned int>("yValue");
+    auto k = F.declVar<unsigned int>("k");
+    auto oldWeight = F.declVar<unsigned int>("oldWeight");
+    auto tempWeight = F.declVar<unsigned int>("tempWeight");
+
+    // 2D thread/block indices
+    auto tidx = F.callBuiltin(getThreadIdX);
+    auto tidy = F.callBuiltin(getThreadIdY);
+    auto bidx = F.callBuiltin(getBlockIdX);
+    auto bidy = F.callBuiltin(getBlockIdY);
+    auto bdimx = F.callBuiltin(getBlockDimX);
+    auto bdimy = F.callBuiltin(getBlockDimY);
+
+    xValue = bidx * bdimx + tidx;
+    yValue = bidy * bdimy + tidy;
+    k = pass;
+
+    oldWeight = pathDistanceBuffer[yValue * rcNumNodes + xValue];
+    tempWeight = pathDistanceBuffer[yValue * rcNumNodes + k] +
+                 pathDistanceBuffer[k * rcNumNodes + xValue];
+
+    F.beginIf(tempWeight < oldWeight);
+    {
+      pathDistanceBuffer[yValue * rcNumNodes + xValue] = tempWeight;
+      pathBuffer[yValue * rcNumNodes + xValue] = k;
+    }
+    F.endIf();
+
+    F.ret();
+  }
+  F.endFunction();
+
+  return std::make_pair(std::move(J), KernelHandle);
+}
+
+int main(int argc, char **argv) {
+  proteus::init();
+  if (argc != 4 && argc != 5) {
+    std::printf("Usage: %s <number of nodes> <iterations> <block size> [verify (0 or 1, default 0)]\n", argv[0]);
+    return 1;
+  }
+
+  unsigned int numNodes = static_cast<unsigned int>(std::atoi(argv[1]));
+  unsigned int numIterations = static_cast<unsigned int>(std::atoi(argv[2]));
+  unsigned int blockSize = static_cast<unsigned int>(std::atoi(argv[3]));
+  int do_verify = (argc == 5) ? std::atoi(argv[4]) : 0;
+
+  const size_t matrixSizeBytes = static_cast<size_t>(numNodes) * static_cast<size_t>(numNodes) * sizeof(unsigned int);
+
+  unsigned int *pathDistanceMatrix = nullptr;
+  unsigned int *pathMatrix = nullptr;
+  unsigned int *verificationPathDistanceMatrix = nullptr;
+  unsigned int *verificationPathMatrix = nullptr;
+
+  if (do_verify) {
+    pathDistanceMatrix = (unsigned int *)std::malloc(matrixSizeBytes);
+    pathMatrix = (unsigned int *)std::malloc(matrixSizeBytes);
+    assert(pathDistanceMatrix && pathMatrix);
+
+    // Initialize path matrix
+    for (unsigned int i = 0; i < numNodes; ++i) {
+      for (unsigned int j = 0; j < i; ++j) {
+        pathMatrix[i * numNodes + j] = i;
+        pathMatrix[j * numNodes + i] = j;
+      }
+      pathMatrix[i * numNodes + i] = i;
+    }
+
+    verificationPathDistanceMatrix = (unsigned int *)std::malloc(matrixSizeBytes);
+    verificationPathMatrix = (unsigned int *)std::malloc(matrixSizeBytes);
+    assert(verificationPathDistanceMatrix && verificationPathMatrix);
+    std::memcpy(verificationPathMatrix, pathMatrix, matrixSizeBytes);
+  }
+
+  if (blockSize * blockSize > 256u) {
+    blockSize = 16u;
+  }
+
+  // 2D launch configuration
+  const unsigned int gridX = (numNodes + blockSize - 1u) / blockSize;
+  const unsigned int gridY = (numNodes + blockSize - 1u) / blockSize;
+
+  unsigned int *pathDistanceBuffer = nullptr;
+  unsigned int *pathBuffer = nullptr;
+  (void)hipMalloc(reinterpret_cast<void **>(&pathDistanceBuffer), matrixSizeBytes);
+  (void)hipMalloc(reinterpret_cast<void **>(&pathBuffer), matrixSizeBytes);
+
+  // JIT compile kernel specialized for this numNodes
+  auto [J, KernelHandle] = createJitModuleSpecial(static_cast<unsigned int>(numNodes));
+  J->compile();
+
+  // hipRAND generator
+  hiprandGenerator_t gen;
+  (void)hiprandCreateGenerator(&gen, HIPRAND_RNG_PSEUDO_DEFAULT);
+  (void)hiprandSetPseudoRandomGeneratorSeed(gen, 1234ULL);
+
+  float total_time_ns = 0.0f;
+
+  for (unsigned int n = 0; n < numIterations; n++) {
+    // Generate matrix on device using hipRAND
+    (void)hiprandGenerate(gen, pathDistanceBuffer, numNodes * numNodes);
+    // Map to [0, MAXDISTANCE] and zero diagonal
+    initRandomMatrix2D<<<dim3(gridX, gridY, 1), dim3(blockSize, blockSize, 1)>>>(pathDistanceBuffer, numNodes);
+
+    if (do_verify && n == numIterations - 1) {
+      // Save initial matrix for CPU reference on last iteration
+      (void)hipMemcpy(verificationPathDistanceMatrix, pathDistanceBuffer, matrixSizeBytes, hipMemcpyDeviceToHost);
+    }
+
+    (void)hipDeviceSynchronize();
+    auto start = std::chrono::steady_clock::now();
+
+    for (unsigned int i = 0; i < numNodes; i++) {
+      (void)KernelHandle.launch({gridX, gridY, 1U}, {blockSize, blockSize, 1U}, 0, nullptr,
+                                pathDistanceBuffer, pathBuffer, static_cast<unsigned>(numNodes), static_cast<unsigned>(i));
+    }
+
+    (void)hipDeviceSynchronize();
+    auto end = std::chrono::steady_clock::now();
+    auto time_ns = std::chrono::duration_cast<std::chrono::nanoseconds>(end - start).count();
+    total_time_ns += static_cast<float>(time_ns);
+  }
+
+  (void)hiprandDestroyGenerator(gen);
+
+  std::printf("Average kernel execution time %f (s)\n", (total_time_ns * 1e-9f) / static_cast<float>(numIterations));
+
+  if (do_verify) {
+    (void)hipMemcpy(pathDistanceMatrix, pathDistanceBuffer, matrixSizeBytes, hipMemcpyDeviceToHost);
+  }
+
+  (void)hipFree(pathDistanceBuffer);
+  (void)hipFree(pathBuffer);
+
+  // verify
+  if (do_verify) {
+    floydWarshallCPUReference(verificationPathDistanceMatrix, verificationPathMatrix, numNodes);
+    if (std::memcmp(pathDistanceMatrix, verificationPathDistanceMatrix, matrixSizeBytes) == 0) {
+      std::printf("PASS\n");
+    } else {
+      std::printf("FAIL\n");
+      if (numNodes <= 8) {
+        for (unsigned int i = 0; i < numNodes; i++) {
+          for (unsigned int j = 0; j < numNodes; j++)
+            std::printf("host: %u ", verificationPathDistanceMatrix[i*numNodes+j]);
+          std::printf("\n");
+        }
+        for (unsigned int i = 0; i < numNodes; i++) {
+          for (unsigned int j = 0; j < numNodes; j++)
+            std::printf("device: %u ", pathDistanceMatrix[i*numNodes+j]);
+          std::printf("\n");
+        }
+      }
+    }
+  }
+
+  if (do_verify) {
+    std::free(pathDistanceMatrix);
+    std::free(pathMatrix);
+    std::free(verificationPathDistanceMatrix);
+    std::free(verificationPathMatrix);
+  }
+
+  proteus::finalize();
+  return 0;
+}
diff --git a/benchmarks/hecbench/dsl/gemm/Makefile b/benchmarks/hecbench/dsl/gemm/Makefile
new file mode 100644
index 0000000..336bb21
--- /dev/null
+++ b/benchmarks/hecbench/dsl/gemm/Makefile
@@ -0,0 +1,60 @@
+#===============================================================================
+# User Options
+#===============================================================================
+
+# Compiler can be set below, or via environment variable
+CC        = ${PROTEUS_CC}
+OPTIMIZE  = yes
+DEBUG     = no
+PROTEUS_PATH  ?= /path/to/proteus/install
+
+#===============================================================================
+# Program name & source code list
+#===============================================================================
+
+SUFFIX = -proteus
+
+program = gemm$(SUFFIX).x
+
+source = main.cpp
+obj = $(source:.cpp=$(SUFFIX).o)
+
+#===============================================================================
+# Sets Flags
+#===============================================================================
+
+OFFLOAD_ARCH ?= gfx942
+HIPFLAGS := --offload-arch=${OFFLOAD_ARCH}
+CFLAGS := $(EXTRA_CFLAGS) -std=c++17 -Wall -DPROTEUS_ENABLE_HIP -I${PROTEUS_PATH}/include -I${ROCM_PATH}/llvm/include/ -I${ROCM_PATH}/include $(shell ${ROCM_PATH}/llvm/bin/llvm-config --cxxflags) -fexceptions $(HIPFLAGS)
+
+LDFLAGS = -L${ROCM_PATH}/lib -L${ROCM_PATH}/llvm/lib \
+	$(shell ${ROCM_PATH}/llvm/bin/llvm-config --libs) \
+	$(shell ${ROCM_PATH}/llvm/bin/llvm-config --system-libs) \
+	-llldCommon -llldELF -lamdhip64 -lhiprtc -lhiprtc-builtins -Wl,-rpath,${ROCM_PATH}/lib
+
+ifeq ($(DEBUG),yes)
+	CFLAGS += -g
+	LDFLAGS  += -g
+endif
+
+ifeq ($(OPTIMIZE),yes)
+	CFLAGS += -O3
+endif
+
+LDFLAGS += -Wl,-rpath,${PROTEUS_PATH}/lib64 -L${PROTEUS_PATH}/lib64/ -lproteus
+
+#===============================================================================
+# Targets to Build
+#===============================================================================
+
+$(program): $(obj) Makefile
+	$(CC) $(CFLAGS) $(obj) -o $@ $(LDFLAGS)
+
+%$(SUFFIX).o: %.cpp Makefile
+	$(CC) $(CFLAGS) -x hip -c $< -o $@
+
+clean:
+	rm -rf *.x *.o *.ll *.bc .proteus
+
+run: $(program)
+	./$(program) --no-verify --N 8192 --kernel jit_regtiled
diff --git a/benchmarks/hecbench/dsl/gemm/main.cpp b/benchmarks/hecbench/dsl/gemm/main.cpp
new file mode 100644
index 0000000..01110bb
--- /dev/null
+++ b/benchmarks/hecbench/dsl/gemm/main.cpp
@@ -0,0 +1,411 @@
+#include <hip/hip_runtime.h>
+#include <proteus/JitFrontend.hpp>
+#include <proteus/Frontend/Builtins.hpp>
+#include <proteus/JitInterface.hpp>
+
+#include <cstdlib>
+#include <cstring>
+#include <cctype>
+#include <string>
+#include <iostream>
+#include <memory>
+#include <tuple>
+#include <utility>
+#include <chrono>
+
+using namespace proteus;
+using namespace builtins::gpu;
+
+#define TARGET "hip"
+
+
+
+
+// clang-format off
+// No FileCheck for now; program prints the resulting C matrix
+
+
+
+static auto getMatmulKernel(int N, int TileSize) {
+  auto JitMod = std::make_unique<JitModule>(TARGET);
+  auto KernelHandle =
+      JitMod->addKernel<void(double *, double *, double *)>("tiled_matmul");
+    auto &F = KernelHandle.F;
+  {
+
+    auto Args = F.getArgs();
+    auto &C = std::get<0>(Args);
+    auto &A = std::get<1>(Args);
+    auto &B = std::get<2>(Args);
+
+    F.beginFunction();
+    {
+        auto Tidx = F.callBuiltin(getThreadIdX);
+        auto Bidx = F.callBuiltin(getBlockIdX);
+        auto Tidy = F.callBuiltin(getThreadIdY);
+        auto Bidy = F.callBuiltin(getBlockIdY);
+
+        auto Row = Bidy * TileSize + Tidy;
+        auto Col = Bidx * TileSize + Tidx;
+
+        auto K = F.declVar<int>("K");
+        auto Zero = F.defRuntimeConst(0);
+        auto One = F.defRuntimeConst(1);
+        auto Nvar = F.defRuntimeConst(N);
+        auto Accum = F.defVar(0.0);
+
+        F.beginFor(K, Zero, Nvar, One);
+        {
+            auto AVal = A[Row * N + K];
+            auto BVal = B[K * N + Col];
+            Accum = Accum + (AVal * BVal);
+        }
+        F.endFor();
+
+        auto CIdx = Row * N + Col;
+        C[CIdx] = Accum;
+
+        F.ret();
+    }
+    F.endFunction();
+  }
+  return std::make_pair(std::move(JitMod), KernelHandle);
+}
+
+
+// Register + shared-memory tiled JIT kernel (C = A x B) for square N x N, double.
+// Configuration mirrors the HIP kernel: 64x64 block tile, 16x16 threads,
+// 4x4 per-thread micro-tile, K tile = 8 (defaults; overridable via macros).
+
+static auto getRegSharedTiledMatmulKernel(int N, int blockTileM, int blockTileN, int kTile, int regTileM, int regTileN) {
+  auto JitMod = std::make_unique<JitModule>(TARGET);
+  auto KernelHandle =
+      JitMod->addKernel<void(double *, double *, double *)>("reg_shared_tiled_matmul");
+  auto &F = KernelHandle.F;
+  {
+    auto Args = F.getArgs();
+    auto &C = std::get<0>(Args);
+    auto &A = std::get<1>(Args);
+    auto &B = std::get<2>(Args);
+
+    // Shared tiles for current K-slice
+    auto AsTile = F.declVar<double[]>(blockTileM * kTile, AddressSpace::SHARED);
+    auto BsTile = F.declVar<double[]>(kTile * blockTileN, AddressSpace::SHARED);
+
+    auto Areg = F.declVar<double[]>(regTileM);
+    auto Breg = F.declVar<double[]>(regTileN);
+    auto Creg = F.declVar<double[]>(regTileM * regTileN);
+
+    F.beginFunction();
+    {
+      auto Tidx = F.callBuiltin(getThreadIdX);
+      auto Tidy = F.callBuiltin(getThreadIdY);
+      auto Bidx = F.callBuiltin(getBlockIdX);
+      auto Bidy = F.callBuiltin(getBlockIdY);
+
+      // Constants
+      auto Nvar = F.defRuntimeConst(N);
+      auto Two = F.defRuntimeConst(2);
+      auto Zero = F.defRuntimeConst(0);
+      auto One = F.defRuntimeConst(1);
+      auto RegTileM = F.defRuntimeConst(regTileM);
+      auto RegTileN = F.defRuntimeConst(regTileN);
+      auto BlockTileM = F.defRuntimeConst(blockTileM);
+      auto BlockTileN = F.defRuntimeConst(blockTileN);
+      auto KTile = F.defRuntimeConst(kTile);
+      auto ThreadsX = F.defRuntimeConst(blockTileN / regTileN);
+
+      // Block origin in C
+      auto BlockRow = Bidy * BlockTileM;
+      auto BlockCol = Bidx * BlockTileN;
+
+      // Per-thread micro tile origin in C
+      auto Row0 = BlockRow + Tidy * RegTileM;
+      auto Col0 = BlockCol + Tidx * RegTileN;
+
+      // Zero accumulators
+      {
+        auto I = F.declVar<int>("i");
+        auto J = F.declVar<int>("j");
+          F.forLoop<int>({I, Zero, RegTileM, One}, [&]() {
+            F.forLoop<int>({J, Zero, RegTileN, One}, [&]() {
+              auto Cidx = I * RegTileN + J;
+            Creg[Cidx] = 0.0;
+          }).emit();
+        }).emit();
+      }
+
+      // Loop over K dimension in tiles of kTile
+      auto KBase = F.declVar<int>("KBase");
+      F.forLoop<int>({KBase, Zero, Nvar, KTile}, [&]() {
+        // Cooperative load of A and B tiles into shared memory.
+        auto Tid = Tidy * ThreadsX + Tidx; // 0 .. (blockTileM/regTileM*blockTileN/regTileN - 1)
+
+        // Load A tile: size [blockTileM x kTile]
+        auto APlaneSize = BlockTileM * KTile;
+        auto AIdx1 = Tid * Two + One;
+        auto AIdx0 = Tid * Two + Zero;
+        auto ARow0 = AIdx0 / KTile;
+        auto ACol0 = AIdx0 % KTile;
+        auto ARow1 = AIdx1 / KTile;
+        auto ACol1 = AIdx1 % KTile;
+        auto AsIdx0 = ARow0 * KTile + ACol0;
+        auto AsIdx1 = ARow1 * KTile + ACol1;
+        auto AGlobIdx0 = (BlockRow + ARow0) * N + (KBase + ACol0);
+        auto AGlobIdx1 = (BlockRow + ARow1) * N + (KBase + ACol1);
+        F.beginIf(AIdx0 < APlaneSize);
+        {
+          AsTile[AsIdx0] = A[AGlobIdx0];
+        }
+        F.endIf();
+        F.beginIf(AIdx1 < APlaneSize);
+        {
+          AsTile[AsIdx1] = A[AGlobIdx1];
+        }
+        F.endIf();
+
+        // Load B tile: size [kTile x blockTileN]
+        auto BPlaneSize = KTile * BlockTileN;
+        auto BIdx0 = Tid * Two + Zero;
+        auto BIdx1 = Tid * Two + One;
+        auto BRow0 = BIdx0 / BlockTileN;
+        auto BCol0 = BIdx0 % BlockTileN;
+        auto BRow1 = BIdx1 / BlockTileN;
+        auto BCol1 = BIdx1 % BlockTileN;
+        auto BsIdx0 = BRow0 * BlockTileN + BCol0;
+        auto BsIdx1 = BRow1 * BlockTileN + BCol1;
+        auto BGlobIdx0 = (KBase + BRow0) * N + (BlockCol + BCol0);
+        auto BGlobIdx1 = (KBase + BRow1) * N + (BlockCol + BCol1);
+        F.beginIf(BIdx0 < BPlaneSize);
+        {
+          BsTile[BsIdx0] = B[BGlobIdx0];
+        }
+        F.endIf();
+        F.beginIf(BIdx1 < BPlaneSize);
+        {
+          BsTile[BsIdx1] = B[BGlobIdx1];
+        }
+        F.endIf();
+
+        F.callBuiltin(syncThreads);
+
+        // Compute this micro-tile using the shared tiles and register blocking
+        auto KIt = F.declVar<int>("KIt");
+        F.forLoop<int>({KIt, Zero, KTile, One}, [&]() {
+          // Load rows/cols into registers
+          auto I = F.declVar<int>("i");
+          auto J = F.declVar<int>("j");
+
+          F.forLoop<int>({I, Zero, RegTileM, One}, [&]() {
+            auto r = Tidy * RegTileM + I;
+            auto asIdx = r * KTile + KIt;
+            Areg[I] = AsTile[asIdx];
+          }).emit();
+
+          F.forLoop<int>({J, Zero, RegTileN, One}, [&]() {
+            auto c = Tidx * RegTileN + J;
+            auto bsIdx = KIt * BlockTileN + c;
+            Breg[J] = BsTile[bsIdx];
+          }).emit();
+
+          // FMA on the micro-tile
+          auto Ii = F.declVar<int>("ii");
+          auto Jj = F.declVar<int>("jj");
+          F.forLoop<int>({Ii, Zero, RegTileM, One}, [&]() {
+            F.forLoop<int>({Jj, Zero, RegTileN, One}, [&]() {
+              auto Cidx = Ii * RegTileN + Jj;
+              Creg[Cidx] = Creg[Cidx] + (Areg[Ii] * Breg[Jj]);
+            }).emit();
+          }).emit();
+        }).emit();
+
+        F.callBuiltin(syncThreads);
+      }).emit();
+
+      // Write back the per-thread micro-tile to C
+      {
+        auto I = F.declVar<int>("i");
+        auto J = F.declVar<int>("j");
+        F.forLoop<int>({I, Zero, RegTileM, One}, [&]() {
+          F.forLoop<int>({J, Zero, RegTileN, One}, [&]() {
+            auto Cidx = (Row0 + I) * N + (Col0 + J);
+            auto Ridx = I * RegTileN + J;
+            C[Cidx] = Creg[Ridx];
+          }).emit();
+        }).emit();
+      }
+
+      F.ret();
+    }
+    F.endFunction();
+  }
+  return std::make_pair(std::move(JitMod), KernelHandle);
+}
+
+bool verify(double *C, int N) {
+  for (int I = 0; I < N; I++) {
+    for (int J = 0; J < N; J++) {
+      if (C[I*N + J] != N) {
+        return false;
+      }
+    }
+  }
+  return true;
+}
+
+int main(int argc, char** argv) {
+  proteus::init();
+
+  // Default tile size constants
+  constexpr int MatmulTileSize = 32;
+  constexpr int RegTileM = 4;
+  constexpr int RegTileN = 4;
+  constexpr int BlockTileM = 64;
+  constexpr int BlockTileN = 64;
+  constexpr int KTile = 8;
+
+  unsigned int N = 8192;
+  int NumTrials = 5;
+  bool DoVerify = true;
+  std::string KernelType = "jit_regtiled";
+  int blockTileMArg = BlockTileM;
+  int blockTileNArg = BlockTileN;
+  int kTileArg = KTile;
+  int posIdx = 0;
+
+  for (int i = 1; i < argc; ++i) {
+    if (!std::strcmp(argv[i], "--N") || !std::strcmp(argv[i], "-n")) {
+      if (i + 1 < argc) {
+        N = static_cast<unsigned int>(std::atoi(argv[++i]));
+      }
+    } else if (!std::strcmp(argv[i], "--trials") || !std::strcmp(argv[i], "-t")) {
+      if (i + 1 < argc) {
+        NumTrials = std::atoi(argv[++i]);
+      }
+    } else if (!std::strcmp(argv[i], "--kernel")) {
+      if (i + 1 < argc) {
+        KernelType = argv[++i];
+        if (KernelType != "jit" && KernelType != "jit_regtiled") {
+          std::cerr << "Error: Invalid kernel type '" << KernelType
+                    << "'. Valid options: jit, jit_regtiled\n";
+          return 1;
+        }
+      }
+    } else if (!std::strcmp(argv[i], "--verify")) {
+      DoVerify = true;
+    } else if (!std::strcmp(argv[i], "--no-verify")) {
+      DoVerify = false;
+    } else if (!std::strcmp(argv[i], "--help") || !std::strcmp(argv[i], "-h")) {
+      std::cout << "Usage: " << argv[0]
+                << " [-n|--N N] [-t|--trials T] [--kernel KERNEL] [--verify|--no-verify]"
+                << " [blockTileM blockTileN kTile]\n"
+                << "  KERNEL: jit, jit_regtiled (default: jit_regtiled)\n"
+                << "  Positional tile sizes are used for the JIT reg-tiled kernel;"
+                << " defaults are " << BlockTileM << " " << BlockTileN << " " << KTile << "\n";
+      return 0;
+    } else {
+      // Positional ints for blockTileM, blockTileN, kTile (for JIT reg-tiled)
+      if (argv[i] && std::isdigit(static_cast<unsigned char>(argv[i][0]))) {
+        int val = std::atoi(argv[i]);
+        if (posIdx == 0) blockTileMArg = val;
+        else if (posIdx == 1) blockTileNArg = val;
+        else if (posIdx == 2) kTileArg = val;
+        ++posIdx;
+      }
+    }
+  }
+  std::cout << "Configuration: (N, NumTrials, DoVerify, Kernel, Tiles) = ("
+            << N << ", " << NumTrials << ", " << DoVerify << ", " << KernelType
+            << ", blkM=" << blockTileMArg << ", blkN=" << blockTileNArg
+            << ", k=" << kTileArg << ")" << std::endl;
+
+  // Host allocations
+  double *AH = (double *)new double[N * N];
+  double *BH = (double *)new double[N * N];
+  double *CH = (double *)new double[N * N];
+
+  // Device allocations
+  double *AD;
+  double *BD;
+  double *CD;
+  size_t Bytes = sizeof(double) * N * N;
+  hipMalloc(reinterpret_cast<void **>(&AD), Bytes);
+  hipMalloc(reinterpret_cast<void **>(&BD), Bytes);
+  hipMalloc(reinterpret_cast<void **>(&CD), Bytes);
+
+  for (int I = 0; I < N; I++) {
+    for (int J = 0; J < N; J++) {
+      AH[I * N + J] = 1.0;
+      BH[I * N + J] = 1.0;
+      CH[I * N + J] = 0.0;
+    }
+  }
+  // Stage inputs to device
+  hipMemcpy(AD, AH, Bytes, hipMemcpyHostToDevice);
+  hipMemcpy(BD, BH, Bytes, hipMemcpyHostToDevice);
+  hipMemcpy(CD, CH, Bytes, hipMemcpyHostToDevice);
+
+  // Kernel execution based on type
+  if (KernelType == "jit_regtiled") {
+    auto [JitMod, KernelHandle] = getRegSharedTiledMatmulKernel(N, blockTileMArg, blockTileNArg, kTileArg, RegTileM, RegTileN);
+    JitMod->compile(true);
+    KernelHandle.launch({static_cast<unsigned int>(N / blockTileNArg), static_cast<unsigned int>(N / blockTileMArg), 1u}, {static_cast<unsigned int>(blockTileNArg / RegTileN), static_cast<unsigned int>(blockTileMArg / RegTileM), 1u}, 0, nullptr, CD, AD, BD);
+    hipDeviceSynchronize();
+
+    // Timed trials
+    double TotalMs = 0.0;
+    auto Start = std::chrono::high_resolution_clock::now();
+    for (int T = 0; T < NumTrials; ++T) {
+      KernelHandle.launch({static_cast<unsigned int>(N / blockTileNArg), static_cast<unsigned int>(N / blockTileMArg), 1u}, {static_cast<unsigned int>(blockTileNArg / RegTileN), static_cast<unsigned int>(blockTileMArg / RegTileM), 1u}, 0, nullptr, CD, AD, BD);
+
+    }
+    hipDeviceSynchronize();
+    auto End = std::chrono::high_resolution_clock::now();
+    std::chrono::duration<double, std::milli> Ms = End - Start;
+    TotalMs = Ms.count();
+    double AvgMs = TotalMs / static_cast<double>(NumTrials);
+    std::cerr.setf(std::ios::fixed);
+    std::cerr.precision(3);
+    std::cerr << "Average over " << NumTrials << " trials: " << AvgMs << " ms" << '\n';
+
+  } else if (KernelType == "jit") {
+    auto [JitMod, KernelHandle] = getMatmulKernel(N, MatmulTileSize);
+    JitMod->compile(true);
+    KernelHandle.launch({(N + MatmulTileSize - 1) / MatmulTileSize, (N + MatmulTileSize - 1) / MatmulTileSize, 1}, {MatmulTileSize, MatmulTileSize, 1}, 0, nullptr, CD, AD, BD);
+    hipDeviceSynchronize();
+
+    // Timed trials
+    double TotalMs = 0.0;
+    auto Start = std::chrono::high_resolution_clock::now();
+    for (int T = 0; T < NumTrials; ++T) {
+      KernelHandle.launch({(N + MatmulTileSize - 1) / MatmulTileSize, (N + MatmulTileSize - 1) / MatmulTileSize, 1}, {MatmulTileSize, MatmulTileSize, 1}, 0, nullptr, CD, AD, BD);
+    }
+    hipDeviceSynchronize();
+    auto End = std::chrono::high_resolution_clock::now();
+    std::chrono::duration<double, std::milli> Ms = End - Start;
+    TotalMs = Ms.count();
+    double AvgMs = TotalMs / static_cast<double>(NumTrials);
+    std::cerr.setf(std::ios::fixed);
+    std::cerr.precision(3);
+    std::cerr << "Average over " << NumTrials << " trials: " << AvgMs << " ms" << '\n';
+  }
+
+  // Final result back to host after timing loop
+  if (DoVerify) {
+    hipMemcpy(CH, CD, Bytes, hipMemcpyDeviceToHost);
+    verify(CH, N);
+    std::cout << "Verification passed" << std::endl;
+  }
+
+  // Cleanup device and host memory
+  hipFree(AD);
+  hipFree(BD);
+  hipFree(CD);
+  delete[] AH;
+  delete[] BH;
+  delete[] CH;
+
+  proteus::finalize();
+  return 0;
+
+}
diff --git a/hecbench.toml b/hecbench.toml
index e1a392c..86c4f07 100644
--- a/hecbench.toml
+++ b/hecbench.toml
@@ -150,6 +150,16 @@ exe = "attention-proteus.x"
 [hecbench.attention.inputs]
 default = "65536 2048 100"
 
+[hecbench.gemm]
+[hecbench.gemm.amd.dsl]
+path = "benchmarks/hecbench/dsl/gemm"
+exe = "gemm-proteus.x"
+[hecbench.gemm.nvidia.dsl]
+path = "benchmarks/hecbench/dsl/gemm"
+exe = "gemm-proteus.x"
+[hecbench.gemm.inputs]
+default = "--no-verify --N 8192 --kernel jit_regtiled"
+
 [hecbench.conv3d]
 [hecbench.conv3d.amd.dsl]
 path = "benchmarks/hecbench/dsl/conv3d"
@@ -160,6 +170,16 @@ exe = "conv3d-proteus.x"
 [hecbench.conv3d.inputs]
 default = "32 96 256 26 26 5 100"
 
+[hecbench.floyd-warshall]
+[hecbench.floyd-warshall.amd.dsl]
+path = "benchmarks/hecbench/dsl/floyd-warshall"
+exe = "floyd-warshall-proteus.x"
+[hecbench.floyd-warshall.nvidia.dsl]
+path = "benchmarks/hecbench/dsl/floyd-warshall"
+exe = "floyd-warshall-proteus.x"
+[hecbench.floyd-warshall.inputs]
+default = "8192 10 16 0"
+
 [hecbench.bezier-surface]
 [hecbench.bezier-surface.amd.dsl]
 path = "benchmarks/hecbench/dsl/bezier-surface"