leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 26 Packages Projects Releases Wiki Activity

Merge pull request #2919 from pbelevich/patch-1

Browse Source 
Refactor binary_op functions to remove unused result parameter
This commit is contained in:
drazi
2026-02-11 11:48:58 +08:00
committed by GitHub
parent 6b3e607b85 b6d7703e02
commit 71aa7a0abc
1 changed files with 5 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/python/CuTeDSL/notebooks/tensorssa.ipynb
+5 -7
View File
@@ -167,7 +167,7 @@
"outputs": [], "outputs": [],
"source": [ "source": [
"@cute.jit\n", "@cute.jit\n",
"def binary_op_1(res: cute.Tensor, a: cute.Tensor, b: cute.Tensor):\n", "def binary_op_1(a: cute.Tensor, b: cute.Tensor):\n",
" a_vec = a.load()\n", " a_vec = a.load()\n",
" b_vec = b.load()\n", " b_vec = b.load()\n",
"\n", "\n",
@@ -184,7 +184,7 @@
" cute.print_tensor(div_res) # prints [0.500000, 0.500000, 0.500000]\n", " cute.print_tensor(div_res) # prints [0.500000, 0.500000, 0.500000]\n",
"\n", "\n",
" floor_div_res = a_vec // b_vec\n", " floor_div_res = a_vec // b_vec\n",
" cute.print_tensor(res) # prints [0.000000, 0.000000, 0.000000]\n", " cute.print_tensor(floor_div_res) # prints [0.000000, 0.000000, 0.000000]\n",
"\n", "\n",
" mod_res = a_vec % b_vec\n", " mod_res = a_vec % b_vec\n",
" cute.print_tensor(mod_res) # prints [1.000000, 1.000000, 1.000000]\n", " cute.print_tensor(mod_res) # prints [1.000000, 1.000000, 1.000000]\n",
@@ -194,8 +194,7 @@
"a.fill(1.0)\n", "a.fill(1.0)\n",
"b = np.empty((3,), dtype=np.float32)\n", "b = np.empty((3,), dtype=np.float32)\n",
"b.fill(2.0)\n", "b.fill(2.0)\n",
"res = np.empty((3,), dtype=np.float32)\n", "binary_op_1(from_dlpack(a), from_dlpack(b))"
"binary_op_1(from_dlpack(res), from_dlpack(a), from_dlpack(b))"
] ]
}, },
{ {
@@ -205,7 +204,7 @@
"outputs": [], "outputs": [],
"source": [ "source": [
"@cute.jit\n", "@cute.jit\n",
"def binary_op_2(res: cute.Tensor, a: cute.Tensor, c: cutlass.Constexpr):\n", "def binary_op_2(a: cute.Tensor, c: cutlass.Constexpr):\n",
" a_vec = a.load()\n", " a_vec = a.load()\n",
"\n", "\n",
" add_res = a_vec + c\n", " add_res = a_vec + c\n",
@@ -230,8 +229,7 @@
"a = np.empty((3,), dtype=np.float32)\n", "a = np.empty((3,), dtype=np.float32)\n",
"a.fill(1.0)\n", "a.fill(1.0)\n",
"c = 2.0\n", "c = 2.0\n",
"res = np.empty((3,), dtype=np.float32)\n", "binary_op_2(from_dlpack(a), c)"
"binary_op_2(from_dlpack(res), from_dlpack(a), c)"
] ]
}, },
{ {