fix transfer_read/write (#152)

Author: makslevental

examples/cuda_e2e.ipynb

@@ -29,65 +29,13 @@
    "cell_type": "markdown",
    "source": "# Download mlir-python-bindings with CUDA support"
   },
-  {
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2025-05-05T16:40:48.337497Z",
-     "start_time": "2025-05-05T16:40:48.251771Z"
-    }
-   },
-   "cell_type": "code",
-   "source": [
-    "from pip._internal.cli import cmdoptions\n",
-    "from pip._internal.commands import index\n",
-    "\n",
-    "cmd = index.IndexCommand(\"blah\", \"\")\n",
-    "\n",
-    "options, args = cmd.parse_args(\n",
-    "    [\n",
-    "        \"index\",\n",
-    "        \"versions\",\n",
-    "        \"mlir-python-bindings\",\n",
-    "        \"-f\",\n",
-    "        \"https://makslevental.github.io/wheels\",\n",
-    "    ]\n",
-    ")\n",
-    "\n",
-    "\n",
-    "def get_available_package_versions(self, options, args):\n",
-    "    target_python = cmdoptions.make_target_python(options)\n",
-    "    query = args[0]\n",
-    "\n",
-    "    with self._build_session(options) as session:\n",
-    "        finder = self._build_package_finder(\n",
-    "            options=options,\n",
-    "            session=session,\n",
-    "            target_python=target_python,\n",
-    "            ignore_requires_python=options.ignore_requires_python,\n",
-    "        )\n",
-    "\n",
-    "        versions = set(\n",
-    "            candidate.version for candidate in finder.find_all_candidates(query)\n",
-    "        )\n",
-    "\n",
-    "        return list(versions)\n",
-    "\n",
-    "\n",
-    "def get_latest_cuda_version(all_versions):\n",
-    "    cuda_versions = list(filter(lambda x: \"cuda\" in x.local, all_versions))\n",
-    "    assert len(cuda_versions), \"couldn't find any cuda versions\"\n",
-    "    cuda_versions.sort(key=lambda x: x.release)\n",
-    "    return cuda_versions[0]\n",
-    "\n"
-   ],
-   "outputs": [],
-   "execution_count": 1
-  },
   {
    "cell_type": "code",
    "source": [
-    "all_versions = get_available_package_versions(cmd, options, [\"mlir-python-bindings\"])\n",
-    "latest_cuda_version = get_latest_cuda_version(all_versions)\n",
+    "!BRANCH=\"${BRANCH:-main}\"\n",
+    "!echo \"using BRANCH=$BRANCH\"\n",
+    "!script_address=\"https://raw.githubusercontent.com/makslevental/mlir-python-extras/refs/heads/$BRANCH/scripts/get_latest_gpu_bindings.py\"\n",
+    "!latest_cuda_version=$(curl $script_address | python -)\n",
     "!pip install -q mlir_python_bindings==$latest_cuda_version -f https://makslevental.github.io/wheels\n",
     "!pip install git+https://github.com/makslevental/mlir-python-extras@$BRANCH &> /dev/null"
    ],

examples/mlir_python_extras.ipynb

@@ -23,6 +23,8 @@
    },
    "outputs": [],
    "source": [
+    "!BRANCH=\"${BRANCH:-main}\"\n",
+    "!echo \"using BRANCH=$BRANCH\"\n",
     "!pip install mlir-python-bindings -f https://makslevental.github.io/wheels &> /dev/null\n",
     "!pip install git+https://github.com/makslevental/mlir-python-extras@$BRANCH &> /dev/null"
    ]

examples/vectorization_e2e.ipynb

@@ -31,6 +31,8 @@
    },
    "outputs": [],
    "source": [
+    "!BRANCH=\"${BRANCH:-main}\"\n",
+    "!echo \"using BRANCH=$BRANCH\"\n",
     "!pip install -q mlir-python-bindings -f https://makslevental.github.io/wheels\n",
     "!pip install -q git+https://github.com/makslevental/mlir-python-extras@$BRANCH"
    ]

mlir/extras/dialects/ext/vector.py

@@ -100,7 +100,7 @@ def transfer_write(
         result=None,
         value_to_store=val,
         # no clue why they chose this name...
-        source=dest,
+        base=dest,
         indices=indices,
         permutation_map=permutation_map,
         mask=mask,
@@ -115,7 +115,7 @@ def transfer_write(
 
 def transfer_read(
     vector_t,
-    source,
+    base,
     indices,
     *,
     permutation_map=None,
@@ -128,20 +128,20 @@ def transfer_read(
     if loc is None:
         loc = get_user_code_loc()
     if permutation_map is None:
-        permutation_map = AffineMap.get_minor_identity(source.type.rank, vector_t.rank)
+        permutation_map = AffineMap.get_minor_identity(base.type.rank, vector_t.rank)
     for j, i in enumerate(indices):
         if isinstance(i, int):
             indices[j] = constant(i, index=True)
     if padding is None:
         padding = 0
     if isinstance(padding, int):
-        padding = constant(padding, type=source.type.element_type)
+        padding = constant(padding, type=base.type.element_type)
     if in_bounds is None:
         raise ValueError("in_bounds cannot be None")
 
     return _transfer_read(
         vector=vector_t,
-        source=source,
+        base=base,
         indices=indices,
         permutation_map=permutation_map,
         padding=padding,

scripts/get_latest_gpu_bindings.py

@@ -0,0 +1,63 @@
+import argparse
+
+from pip._internal.cli import cmdoptions
+from pip._internal.commands import index
+
+
+def get_available_package_versions(self, options, args):
+    target_python = cmdoptions.make_target_python(options)
+    query = args[0]
+
+    with self._build_session(options) as session:
+        finder = self._build_package_finder(
+            options=options,
+            session=session,
+            target_python=target_python,
+            ignore_requires_python=options.ignore_requires_python,
+        )
+
+        versions = set(
+            candidate.version for candidate in finder.find_all_candidates(query)
+        )
+
+        return list(versions)
+
+
+def get_latest_gpu_version(all_versions, platform: str):
+    bindings_versions = list(
+        filter(lambda x: x.local and platform in x.local, all_versions)
+    )
+    assert len(bindings_versions), "couldn't find any bindings versions"
+    bindings_versions.sort(key=lambda x: x.release)
+    return bindings_versions[0]
+
+
+cmd = index.IndexCommand("blah", "")
+
+
+def get_latest_gpu_version_name(platform):
+    options, _args = cmd.parse_args(
+        [
+            "index",
+            "versions",
+            "mlir-python-bindings",
+            "--find-links",
+            "https://makslevental.github.io/wheels",
+        ]
+    )
+    options.no_index = True
+    all_versions = get_available_package_versions(
+        cmd, options, ["mlir-python-bindings"]
+    )
+    return get_latest_gpu_version(all_versions, platform)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("platform", nargs="?", default="cuda")
+    platform = parser.parse_args().platform
+    print(get_latest_gpu_version_name(platform))
+
+
+if __name__ == "__main__":
+    main()