Sync pr-g-examples with base branch pr-f-test-migration

examples/mha/mha_module.py

@@ -0,0 +1,152 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""
+Multi-Head Attention as an ``ark.Module`` subclass.
+
+Demonstrates:
+  - Subclassing ``ark.Module`` with ``ark.parameter()``.
+  - Using the functional API (``ark.matmul``, ``ark.reshape``,
+    ``ark.transpose``, ``ark.softmax``, ``ark.mul``).
+  - Numerical validation against a manual PyTorch reference implementation.
+
+Run: ``python examples/mha/mha_module.py``
+"""
+
+import math
+import numpy as np
+import torch
+import ark
+
+# ---------- hyperparameters ----------
+batch_size = 1
+seq_len = 64
+d_model = 128
+n_heads = 4
+
+
+# ---------- ARK Module ----------
+class MultiHeadAttention(ark.Module):
+    """Scaled dot-product multi-head attention (no bias, no mask)."""
+
+    def __init__(
+        self, d_model: int, n_heads: int, batch_size: int, seq_len: int
+    ):
+        super().__init__()
+        self.n_heads = n_heads
+        self.d_k = d_model // n_heads
+        self.scale = 1.0 / math.sqrt(self.d_k)
+        self.batch_size = batch_size
+        self.seq_len = seq_len
+
+        # Projection weights: [d_model, d_model]
+        self.wq = ark.parameter([d_model, d_model], ark.fp16)
+        self.wk = ark.parameter([d_model, d_model], ark.fp16)
+        self.wv = ark.parameter([d_model, d_model], ark.fp16)
+        self.wo = ark.parameter([d_model, d_model], ark.fp16)
+
+    def forward(self, x):
+        # x: [batch_size, seq_len, d_model]
+        # Shape is fixed at graph-build time; stored as instance attributes.
+        H = self.n_heads
+        dk = self.d_k
+
+        # Linear projections
+        q = ark.matmul(x, self.wq)  # [batch_size, seq_len, d_model]
+        k = ark.matmul(x, self.wk)
+        v = ark.matmul(x, self.wv)
+
+        # Reshape to [batch_size, seq_len, H, dk] then transpose to [batch_size, H, seq_len, dk]
+        q = ark.transpose(
+            ark.reshape(q, [self.batch_size, self.seq_len, H, dk]), [0, 2, 1, 3]
+        )
+        k = ark.transpose(
+            ark.reshape(k, [self.batch_size, self.seq_len, H, dk]), [0, 2, 1, 3]
+        )
+        v = ark.transpose(
+            ark.reshape(v, [self.batch_size, self.seq_len, H, dk]), [0, 2, 1, 3]
+        )
+
+        # Scaled dot-product attention
+        # scores: [batch_size, H, seq_len, seq_len]
+        scores = ark.matmul(q, k, transpose_other=True)
+        scores = ark.mul(scores, self.scale)
+        attn = ark.softmax(scores)  # along last axis
+        # context: [batch_size, H, seq_len, dk]
+        context = ark.matmul(attn, v)
+
+        # Transpose back and reshape: [batch_size, seq_len, d_model]
+        context = ark.reshape(
+            ark.transpose(context, [0, 2, 1, 3]),
+            [self.batch_size, self.seq_len, H * dk],
+        )
+
+        # Output projection
+        out = ark.matmul(context, self.wo)
+        return out
+
+
+# ---------- PyTorch reference ----------
+def pytorch_mha(x_np, wq, wk, wv, wo):
+    """Manual MHA in PyTorch matching the ARK implementation above."""
+    x = torch.from_numpy(x_np).cuda()
+    B, S, D = x.shape
+    H = n_heads
+    dk = d_model // n_heads
+
+    q = x @ torch.from_numpy(wq).cuda()
+    k = x @ torch.from_numpy(wk).cuda()
+    v = x @ torch.from_numpy(wv).cuda()
+
+    q = q.view(B, S, H, dk).permute(0, 2, 1, 3)
+    k = k.view(B, S, H, dk).permute(0, 2, 1, 3)
+    v = v.view(B, S, H, dk).permute(0, 2, 1, 3)
+
+    scores = q @ k.transpose(-2, -1) / math.sqrt(dk)
+    attn = torch.softmax(scores, dim=-1)
+    context = attn @ v
+
+    context = context.permute(0, 2, 1, 3).contiguous().view(B, S, D)
+    out = context @ torch.from_numpy(wo).cuda()
+    return out.cpu().numpy()
+
+
+# ---------- main ----------
+def main():
+    ark.init()
+
+    # Build graph
+    x_ark = ark.tensor([batch_size, seq_len, d_model], ark.fp16)
+    model = MultiHeadAttention(d_model, n_heads, batch_size, seq_len)
+    y_ark = model(x_ark)
+
+    # Launch runtime
+    runtime = ark.Runtime()
+    runtime.launch()
+
+    # Random inputs and weights (small magnitude for fp16 stability)
+    rng = np.random.RandomState(42)
+    x_np = (rng.randn(batch_size, seq_len, d_model) * 0.02).astype(np.float16)
+    wq_np = (rng.randn(d_model, d_model) * 0.02).astype(np.float16)
+    wk_np = (rng.randn(d_model, d_model) * 0.02).astype(np.float16)
+    wv_np = (rng.randn(d_model, d_model) * 0.02).astype(np.float16)
+    wo_np = (rng.randn(d_model, d_model) * 0.02).astype(np.float16)
+
+    x_ark.from_numpy(x_np)
+    model.load_state_dict({"wq": wq_np, "wk": wk_np, "wv": wv_np, "wo": wo_np})
+
+    # Run ARK
+    runtime.run()
+    y_host = y_ark.to_numpy()
+
+    # Run PyTorch reference
+    y_ref = pytorch_mha(x_np, wq_np, wk_np, wv_np, wo_np)
+
+    # Validate
+    np.testing.assert_allclose(y_host, y_ref, atol=1e-2, rtol=1e-2)
+    max_err = np.max(np.abs(y_host - y_ref))
+    print(f"MHA module test passed  (max error: {max_err:.6f})")
+
+
+if __name__ == "__main__":
+    main()

examples/tutorial/tensor_ops_tutorial.py

@@ -0,0 +1,146 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+"""
+Tensor ops tutorial — covers tensor creation, elementwise ops,
+reductions, and ``Runtime`` usage with the current ARK functional API.
+
+Complements ``quickstart_tutorial.py`` by exercising more ops:
+  - ``ark.tensor``, ``ark.ones``, ``ark.zeros``, ``ark.constant``
+  - ``ark.add``, ``ark.sub``, ``ark.mul``, ``ark.div``
+  - ``ark.exp``, ``ark.sqrt``, ``ark.rsqrt``
+  - ``ark.relu``, ``ark.gelu``, ``ark.sigmoid``
+  - ``ark.reduce_sum``, ``ark.reduce_mean``, ``ark.reduce_max``
+  - ``ark.layernorm``
+  - ``ark.reshape``, ``ark.transpose``
+
+Each op result is validated against a NumPy reference.
+
+Run: ``python examples/tutorial/tensor_ops_tutorial.py``
+"""
+
+import math
+
+import numpy as np
+import ark
+
+
+def tensor_ops_tutorial():
+    ark.init()
+
+    M, N = 64, 128
+
+    # ---- Tensor creation ----
+    x = ark.tensor([M, N], ark.fp32)
+    y = ark.tensor([M, N], ark.fp32)
+    one = ark.ones([M, N], ark.fp32)
+    zero = ark.zeros([M, N], ark.fp32)
+    c5 = ark.constant(5.0, [M, N], ark.fp32)
+
+    # ---- Elementwise arithmetic ----
+    a_add = ark.add(x, y)
+    a_sub = ark.sub(x, y)
+    a_mul = ark.mul(x, y)
+    a_div = ark.div(x, ark.add(y, one))  # avoid div-by-zero
+
+    # ---- Unary math ----
+    a_exp = ark.exp(x)
+    a_sqrt = ark.sqrt(ark.add(x, c5))  # shift to positive
+    a_rsqrt = ark.rsqrt(ark.add(x, c5))
+
+    # ---- Activations ----
+    a_relu = ark.relu(x)
+    a_gelu = ark.gelu(x)
+    a_sig = ark.sigmoid(x)
+
+    # ---- Reductions ----
+    r_sum = ark.reduce_sum(x, axis=-1)  # [M, 1]
+    r_mean = ark.reduce_mean(x, axis=-1)
+    r_max = ark.reduce_max(x, axis=-1)
+
+    # ---- Layernorm ----
+    a_ln = ark.layernorm(x)
+
+    # ---- Reshape / transpose ----
+    a_reshape = ark.reshape(x, [M * N])
+    a_trans = ark.transpose(
+        ark.reshape(x, [M, N // 2, 2]), [0, 2, 1]
+    )  # [M, 2, N//2]
+
+    # ---- Launch runtime ----
+    runtime = ark.Runtime()
+    runtime.launch()
+
+    rng = np.random.RandomState(0)
+    x_np = (rng.randn(M, N) * 0.5).astype(np.float32)
+    y_np = (rng.randn(M, N) * 0.5).astype(np.float32)
+
+    x.from_numpy(x_np)
+    y.from_numpy(y_np)
+
+    runtime.run()
+
+    # ---- Validate ----
+    def check(name, ark_tensor, expected, atol=1e-5):
+        got = ark_tensor.to_numpy()
+        np.testing.assert_allclose(
+            got, expected, atol=atol, rtol=1e-4, err_msg=name
+        )
+
+    # Elementwise
+    check("add", a_add, x_np + y_np)
+    check("sub", a_sub, x_np - y_np)
+    check("mul", a_mul, x_np * y_np)
+    check("div", a_div, x_np / (y_np + 1.0))
+
+    # Unary
+    check("exp", a_exp, np.exp(x_np), atol=1e-4)
+    check("sqrt", a_sqrt, np.sqrt(x_np + 5.0), atol=1e-5)
+    check("rsqrt", a_rsqrt, 1.0 / np.sqrt(x_np + 5.0), atol=1e-5)
+
+    # Activations
+    check("relu", a_relu, np.maximum(x_np, 0))
+
+    # GELU: approximate check (ARK uses erff-based GELU)
+
+    gelu_ref = 0.5 * x_np * (1 + np.vectorize(math.erf)(x_np / np.sqrt(2)))
+    check("gelu", a_gelu, gelu_ref, atol=1e-4)
+
+    sig_ref = 1.0 / (1.0 + np.exp(-x_np))
+    check("sigmoid", a_sig, sig_ref, atol=1e-5)
+
+    # Reductions (keepdims=True by default → [M, 1])
+    check("reduce_sum", r_sum, x_np.sum(axis=-1, keepdims=True))
+    check(
+        "reduce_mean",
+        r_mean,
+        x_np.mean(axis=-1, keepdims=True),
+        atol=1e-4,
+    )
+    check("reduce_max", r_max, x_np.max(axis=-1, keepdims=True))
+
+    # Layernorm
+    mu = x_np.mean(axis=-1, keepdims=True)
+    # eps must match ark.layernorm default (1e-6)
+    var = ((x_np - mu) ** 2).mean(axis=-1, keepdims=True)
+    ln_ref = (x_np - mu) / np.sqrt(var + 1e-6)
+    check("layernorm", a_ln, ln_ref, atol=1e-4)
+
+    # Reshape / transpose
+    check("reshape", a_reshape, x_np.reshape(M * N))
+    check(
+        "transpose",
+        a_trans,
+        x_np.reshape(M, N // 2, 2).transpose(0, 2, 1),
+    )
+
+    # Constant tensors
+    check("ones", one, np.ones((M, N), dtype=np.float32))
+    check("zeros", zero, np.zeros((M, N), dtype=np.float32))
+    check("constant", c5, np.full((M, N), 5.0, dtype=np.float32))
+
+    print("Tensor ops tutorial passed — all ops validated!")
+
+
+if __name__ == "__main__":
+    tensor_ops_tutorial()