{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cutlass\n",
    "import cutlass.cute as cute\n",
    "from cutlass.cute.runtime import from_dlpack\n",
    "\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to the TensorSSA in CuTe DSL\n",
    "\n",
    "This tutorial introduces what is the `TensorSSA` and why we need it. We also give some examples to show how to use `TensorSSA`.\n",
    "\n",
    "## What is TensorSSA\n",
    "\n",
    "`TensorSSA` is a Python class that represents a tensor value in Static Single Assignment (SSA) form within the CuTe DSL. You can think of it as a tensor residing in a (simulated) register.\n",
    "\n",
    "## Why TensorSSA\n",
    "\n",
    "`TensorSSA` encapsulates the underlying MLIR tensor value into an object that's easier to manipulate in Python. By overloading numerous Python operators (like `+`, `-`, `*`, `/`, `[]`, etc.), it allows users to express tensor computations (primarily element-wise operations and reductions) in a more Pythonic way. These element-wise operations are then translated into optimized vectorization instructions.\n",
    "\n",
    "It's part of the CuTe DSL, serving as a bridge between the user-described computational logic and the lower-level MLIR IR, particularly for representing and manipulating register-level data.\n",
    "\n",
    "## When to use TensorSSA\n",
    "\n",
    "`TensorSSA` is primarily used in the following scenarios:\n",
    "\n",
    "### Load from memory and store to memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@cute.jit\n",
    "def load_and_store(res: cute.Tensor, a: cute.Tensor, b: cute.Tensor):\n",
    "    \"\"\"\n",
    "    Load data from memory and store the result to memory.\n",
    "\n",
    "    :param res: The destination tensor to store the result.\n",
    "    :param a: The source tensor to be loaded.\n",
    "    :param b: The source tensor to be loaded.\n",
    "    \"\"\"\n",
    "    a_vec = a.load()\n",
    "    print(f\"a_vec: {a_vec}\")  # prints `a_vec: vector<12xf32> o (3, 4)`\n",
    "    b_vec = b.load()\n",
    "    print(f\"b_vec: {b_vec}\")  # prints `b_vec: vector<12xf32> o (3, 4)`\n",
    "    res.store(a_vec + b_vec)\n",
    "    cute.print_tensor(res)\n",
    "\n",
    "\n",
    "a = np.ones(12).reshape((3, 4)).astype(np.float32)\n",
    "b = np.ones(12).reshape((3, 4)).astype(np.float32)\n",
    "c = np.zeros(12).reshape((3, 4)).astype(np.float32)\n",
    "load_and_store(from_dlpack(c), from_dlpack(a), from_dlpack(b))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Register-Level Tensor Operations\n",
    "\n",
    "When writing kernel logic, various computations, transformations, slicing, etc., are performed on data loaded into registers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@cute.jit\n",
    "def apply_slice(src: cute.Tensor, dst: cute.Tensor, indices: cutlass.Constexpr):\n",
    "    \"\"\"\n",
    "    Apply slice operation on the src tensor and store the result to the dst tensor.\n",
    "\n",
    "    :param src: The source tensor to be sliced.\n",
    "    :param dst: The destination tensor to store the result.\n",
    "    :param indices: The indices to slice the source tensor.\n",
    "    \"\"\"\n",
    "    src_vec = src.load()\n",
    "    dst_vec = src_vec[indices]\n",
    "    print(f\"{src_vec} -> {dst_vec}\")\n",
    "    if cutlass.const_expr(isinstance(dst_vec, cute.TensorSSA)):\n",
    "        dst.store(dst_vec)\n",
    "        cute.print_tensor(dst)\n",
    "    else:\n",
    "        dst[0] = dst_vec\n",
    "        cute.print_tensor(dst)\n",
    "\n",
    "\n",
    "def slice_1():\n",
    "    src_shape = (4, 2, 3)\n",
    "    dst_shape = (4, 3)\n",
    "    indices = (None, 1, None)\n",
    "\n",
    "    \"\"\"\n",
    "    a:\n",
    "    [[[ 0.  1.  2.]\n",
    "      [ 3.  4.  5.]]\n",
    "\n",
    "     [[ 6.  7.  8.]\n",
    "      [ 9. 10. 11.]]\n",
    "\n",
    "     [[12. 13. 14.]\n",
    "      [15. 16. 17.]]\n",
    "\n",
    "     [[18. 19. 20.]\n",
    "      [21. 22. 23.]]]\n",
    "    \"\"\"\n",
    "    a = np.arange(np.prod(src_shape)).reshape(*src_shape).astype(np.float32)\n",
    "    dst = np.random.randn(*dst_shape).astype(np.float32)\n",
    "    apply_slice(from_dlpack(a), from_dlpack(dst), indices)\n",
    "\n",
    "\n",
    "slice_1()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def slice_2():\n",
    "    src_shape = (4, 2, 3)\n",
    "    dst_shape = (1,)\n",
    "    indices = 10\n",
    "    a = np.arange(np.prod(src_shape)).reshape(*src_shape).astype(np.float32)\n",
    "    dst = np.random.randn(*dst_shape).astype(np.float32)\n",
    "    apply_slice(from_dlpack(a), from_dlpack(dst), indices)\n",
    "\n",
    "\n",
    "slice_2()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Arithmetic Operations\n",
    "\n",
    "As we mentioned earlier, there're many tensor operations whose operands are `TensorSSA`. And they are all element-wise operations. We give some examples below.\n",
    "\n",
    "### Binary Operations\n",
    "\n",
    "For binary operations, the LHS operand is `TensorSSA` and the RHS operand can be either `TensorSSA` or `Numeric`. When the RHS is `Numeric`, it will be broadcast to a `TensorSSA`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@cute.jit\n",
    "def binary_op_1(res: cute.Tensor, a: cute.Tensor, b: cute.Tensor):\n",
    "    a_vec = a.load()\n",
    "    b_vec = b.load()\n",
    "\n",
    "    add_res = a_vec + b_vec\n",
    "    cute.print_tensor(add_res)  # prints [3.000000, 3.000000, 3.000000]\n",
    "\n",
    "    sub_res = a_vec - b_vec\n",
    "    cute.print_tensor(sub_res)  # prints [-1.000000, -1.000000, -1.000000]\n",
    "\n",
    "    mul_res = a_vec * b_vec\n",
    "    cute.print_tensor(mul_res)  # prints [2.000000, 2.000000, 2.000000]\n",
    "\n",
    "    div_res = a_vec / b_vec\n",
    "    cute.print_tensor(div_res)  # prints [0.500000, 0.500000, 0.500000]\n",
    "\n",
    "    floor_div_res = a_vec // b_vec\n",
    "    cute.print_tensor(res)  # prints [0.000000, 0.000000, 0.000000]\n",
    "\n",
    "    mod_res = a_vec % b_vec\n",
    "    cute.print_tensor(mod_res)  # prints [1.000000, 1.000000, 1.000000]\n",
    "\n",
    "\n",
    "a = np.empty((3,), dtype=np.float32)\n",
    "a.fill(1.0)\n",
    "b = np.empty((3,), dtype=np.float32)\n",
    "b.fill(2.0)\n",
    "res = np.empty((3,), dtype=np.float32)\n",
    "binary_op_1(from_dlpack(res), from_dlpack(a), from_dlpack(b))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@cute.jit\n",
    "def binary_op_2(res: cute.Tensor, a: cute.Tensor, c: cutlass.Constexpr):\n",
    "    a_vec = a.load()\n",
    "\n",
    "    add_res = a_vec + c\n",
    "    cute.print_tensor(add_res)  # prints [3.000000, 3.000000, 3.000000]\n",
    "\n",
    "    sub_res = a_vec - c\n",
    "    cute.print_tensor(sub_res)  # prints [-1.000000, -1.000000, -1.000000]\n",
    "\n",
    "    mul_res = a_vec * c\n",
    "    cute.print_tensor(mul_res)  # prints [2.000000, 2.000000, 2.000000]\n",
    "\n",
    "    div_res = a_vec / c\n",
    "    cute.print_tensor(div_res)  # prints [0.500000, 0.500000, 0.500000]\n",
    "\n",
    "    floor_div_res = a_vec // c\n",
    "    cute.print_tensor(floor_div_res)  # prints [0.000000, 0.000000, 0.000000]\n",
    "\n",
    "    mod_res = a_vec % c\n",
    "    cute.print_tensor(mod_res)  # prints [1.000000, 1.000000, 1.000000]\n",
    "\n",
    "\n",
    "a = np.empty((3,), dtype=np.float32)\n",
    "a.fill(1.0)\n",
    "c = 2.0\n",
    "res = np.empty((3,), dtype=np.float32)\n",
    "binary_op_2(from_dlpack(res), from_dlpack(a), c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@cute.jit\n",
    "def binary_op_3(res: cute.Tensor, a: cute.Tensor, b: cute.Tensor):\n",
    "    a_vec = a.load()\n",
    "    b_vec = b.load()\n",
    "\n",
    "    gt_res = a_vec > b_vec\n",
    "    res.store(gt_res)\n",
    "\n",
    "    \"\"\"\n",
    "    ge_res = a_ >= b_   # [False, True, False]\n",
    "    lt_res = a_ < b_    # [True, False, True]\n",
    "    le_res = a_ <= b_   # [True, False, True]\n",
    "    eq_res = a_ == b_   # [False, False, False]\n",
    "    \"\"\"\n",
    "\n",
    "\n",
    "a = np.array([1, 2, 3], dtype=np.float32)\n",
    "b = np.array([2, 1, 4], dtype=np.float32)\n",
    "res = np.empty((3,), dtype=np.bool_)\n",
    "binary_op_3(from_dlpack(res), from_dlpack(a), from_dlpack(b))\n",
    "print(res)  # prints [False, True, False]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@cute.jit\n",
    "def binary_op_4(res: cute.Tensor, a: cute.Tensor, b: cute.Tensor):\n",
    "    a_vec = a.load()\n",
    "    b_vec = b.load()\n",
    "\n",
    "    xor_res = a_vec ^ b_vec\n",
    "    res.store(xor_res)\n",
    "\n",
    "    # or_res = a_vec | b_vec\n",
    "    # res.store(or_res)     # prints [3, 2, 7]\n",
    "\n",
    "    # and_res = a_vec & b_vec\n",
    "    # res.store(and_res)      # prints [0, 2, 0]\n",
    "\n",
    "\n",
    "a = np.array([1, 2, 3], dtype=np.int32)\n",
    "b = np.array([2, 2, 4], dtype=np.int32)\n",
    "res = np.empty((3,), dtype=np.int32)\n",
    "binary_op_4(from_dlpack(res), from_dlpack(a), from_dlpack(b))\n",
    "print(res)  # prints [3, 0, 7]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Unary Operations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@cute.jit\n",
    "def unary_op_1(res: cute.Tensor, a: cute.Tensor):\n",
    "    a_vec = a.load()\n",
    "\n",
    "    sqrt_res = cute.math.sqrt(a_vec)\n",
    "    cute.print_tensor(sqrt_res)  # prints [2.000000, 2.000000, 2.000000]\n",
    "\n",
    "    sin_res = cute.math.sin(a_vec)\n",
    "    res.store(sin_res)\n",
    "    cute.print_tensor(sin_res)  # prints [-0.756802, -0.756802, -0.756802]\n",
    "\n",
    "    exp2_res = cute.math.exp2(a_vec)\n",
    "    cute.print_tensor(exp2_res)  # prints [16.000000, 16.000000, 16.000000]\n",
    "\n",
    "\n",
    "a = np.array([4.0, 4.0, 4.0], dtype=np.float32)\n",
    "res = np.empty((3,), dtype=np.float32)\n",
    "unary_op_1(from_dlpack(res), from_dlpack(a))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Reduction Operation\n",
    "\n",
    "The `TensorSSA`'s `reduce` method applies a specified reduction operation (`ReductionOp.ADD`, \n",
    "`ReductionOp.MUL`, `ReductionOp.MAX`, `ReductionOp.MIN`) starting with an initial value, and \n",
    "performs this reduction along the dimensions specified by the `reduction_profile`. The result \n",
    "is typically a new `TensorSSA` with reduced dimensions or a scalar value if it reduces across \n",
    "all axes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@cute.jit\n",
    "def reduction_op(a: cute.Tensor):\n",
    "    \"\"\"\n",
    "    Apply reduction operation on the src tensor.\n",
    "\n",
    "    :param src: The source tensor to be reduced.\n",
    "    \"\"\"\n",
    "    a_vec = a.load()\n",
    "    red_res = a_vec.reduce(cute.ReductionOp.ADD, 0.0, reduction_profile=0)\n",
    "    cute.printf(red_res)  # prints 21.000000\n",
    "\n",
    "    red_res = a_vec.reduce(cute.ReductionOp.ADD, 0.0, reduction_profile=(None, 1))\n",
    "    cute.print_tensor(red_res)  # prints [6.000000, 15.000000]\n",
    "\n",
    "    red_res = a_vec.reduce(cute.ReductionOp.ADD, 1.0, reduction_profile=(1, None))\n",
    "    cute.print_tensor(red_res)  # prints [6.000000, 8.000000, 10.000000]\n",
    "\n",
    "\n",
    "a = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)\n",
    "reduction_op(from_dlpack(a))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Broadcast\n",
    "\n",
    "`TensorSSA` supports broadcasting operations following NumPy's broadcasting rules. Broadcasting \n",
    "allows you to perform operations on arrays of different shapes when certain conditions are met. \n",
    "The key rules are:\n",
    "\n",
    "1. Source shape is padded with 1's to match the rank of target shape\n",
    "2. The size in each mode of source shape must either be 1 or equal to target shape\n",
    "3. After broadcasting, all modes should match target shape\n",
    "\n",
    "Let's look at some examples of broadcasting in action:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cutlass\n",
    "import cutlass.cute as cute\n",
    "\n",
    "\n",
    "@cute.jit\n",
    "def broadcast_examples():\n",
    "    a = cute.make_rmem_tensor((1, 3), dtype=cutlass.Float32)\n",
    "    a[0] = 0.0\n",
    "    a[1] = 1.0\n",
    "    a[2] = 2.0\n",
    "    a_val = a.load()\n",
    "    cute.print_tensor(a_val.broadcast_to((4, 3)))\n",
    "    # tensor(raw_ptr(0x00007ffe26625740: f32, rmem, align<32>) o (4,3):(1,4), data=\n",
    "    #    [[ 0.000000,  1.000000,  2.000000, ],\n",
    "    #     [ 0.000000,  1.000000,  2.000000, ],\n",
    "    #     [ 0.000000,  1.000000,  2.000000, ],\n",
    "    #     [ 0.000000,  1.000000,  2.000000, ]])\n",
    "\n",
    "    c = cute.make_rmem_tensor((4, 1), dtype=cutlass.Float32)\n",
    "    c[0] = 0.0\n",
    "    c[1] = 1.0\n",
    "    c[2] = 2.0\n",
    "    c[3] = 3.0\n",
    "    cute.print_tensor(a.load() + c.load())\n",
    "    # tensor(raw_ptr(0x00007ffe26625780: f32, rmem, align<32>) o (4,3):(1,4), data=\n",
    "    #        [[ 0.000000,  1.000000,  2.000000, ],\n",
    "    #         [ 1.000000,  2.000000,  3.000000, ],\n",
    "    #         [ 2.000000,  3.000000,  4.000000, ],\n",
    "    #         [ 3.000000,  4.000000,  5.000000, ]])\n",
    "\n",
    "\n",
    "broadcast_examples()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "raw"
    }
   },
   "source": [
    "The examples above demonstrate two key broadcasting scenarios:\n",
    "\n",
    "1. **Row Vector Broadcasting**: In the first example, we create a row vector `a` with shape \n",
    "   (1, 3) containing values [0.0, 1.0, 2.0]. When we broadcast it to shape (4, 3), the values \n",
    "   are repeated across the first dimension, resulting in:\n",
    "   ```\n",
    "   [[0.0, 1.0, 2.0],\n",
    "    [0.0, 1.0, 2.0],\n",
    "    [0.0, 1.0, 2.0],\n",
    "    [0.0, 1.0, 2.0]]\n",
    "   ```\n",
    "   This demonstrates how a row vector can be broadcast to create multiple identical rows.\n",
    "\n",
    "2. **Column Vector and Row Vector Addition**: In the second example, we have:\n",
    "   - A row vector `a` with shape (1, 3) containing [0.0, 1.0, 2.0]\n",
    "   - A column vector `c` with shape (4, 1) containing [0.0, 1.0, 2.0, 3.0]\n",
    "   \n",
    "   When we add these together, both vectors are broadcast to shape (4, 3):\n",
    "   - The row vector is broadcast vertically (4 times)\n",
    "   - The column vector is broadcast horizontally (3 times)\n",
    "   \n",
    "   The result is:\n",
    "   ```\n",
    "   [[0.0 + 0.0, 1.0 + 0.0, 2.0 + 0.0],\n",
    "    [0.0 + 1.0, 1.0 + 1.0, 2.0 + 1.0],\n",
    "    [0.0 + 2.0, 1.0 + 2.0, 2.0 + 2.0],\n",
    "    [0.0 + 3.0, 1.0 + 3.0, 2.0 + 3.0]]\n",
    "   ```\n",
    "   =\n",
    "   ```\n",
    "   [[0.0, 1.0, 2.0],\n",
    "    [1.0, 2.0, 3.0],\n",
    "    [2.0, 3.0, 4.0],\n",
    "    [3.0, 4.0, 5.0]]\n",
    "   ```\n",
    "\n",
    "This demonstrates how `TensorSSA` can automatically handle broadcasting of both row and column \n",
    "vectors in arithmetic operations, following the broadcasting rules where each dimension must \n",
    "either be 1 or match the target size. The broadcasting is handled implicitly during operations, \n",
    "making it easy to work with tensors of different shapes.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv3_12",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}