From 8133aad0a7441565947818233f8ec7987e7fdc5a Mon Sep 17 00:00:00 2001 From: "Li, Tingqian" Date: Fri, 19 Jun 2026 08:29:28 +0000 Subject: [PATCH 1/5] fix slice return type on coordtensor --- lib/Dialect/Fly/IR/FlyOps.cpp | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/lib/Dialect/Fly/IR/FlyOps.cpp b/lib/Dialect/Fly/IR/FlyOps.cpp index 4f6029d276..ad3da63e1a 100644 --- a/lib/Dialect/Fly/IR/FlyOps.cpp +++ b/lib/Dialect/Fly/IR/FlyOps.cpp @@ -903,11 +903,15 @@ FLY_INFER_RETURN_TYPES(SliceOp) { if (auto srcCoordTensorTy = dyn_cast(srcTy)) { Attribute layout = srcCoordTensorTy.getLayout(); Attribute newLayout; - if (auto la = dyn_cast(layout)) + if (auto la = dyn_cast(layout)) { newLayout = sliceLayout(la); - else + IntTupleAttr offsetAttr = layoutCrd2Idx(builder, coordAttr, la.getShape(), la.getStride()); + IntTupleAttr newBase = intTupleAdd(builder, srcCoordTensorTy.getBase(), offsetAttr); + inferredReturnTypes.assign({CoordTensorType::get(newBase, newLayout)}); + } else { newLayout = sliceComposed(cast(layout)); - inferredReturnTypes.assign({CoordTensorType::get(srcCoordTensorTy.getBase(), newLayout)}); + inferredReturnTypes.assign({CoordTensorType::get(srcCoordTensorTy.getBase(), newLayout)}); + } return success(); } From b9bf1b5d1b0ae450e5fa96954040911b233b4b35 Mon Sep 17 00:00:00 2001 From: "Li, Tingqian" Date: Fri, 19 Jun 2026 09:45:30 +0000 Subject: [PATCH 2/5] add test case --- tests/unit/test_coord_tensor.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 tests/unit/test_coord_tensor.py diff --git a/tests/unit/test_coord_tensor.py b/tests/unit/test_coord_tensor.py new file mode 100644 index 0000000000..4b5e0cc795 --- /dev/null +++ b/tests/unit/test_coord_tensor.py @@ -0,0 +1,23 @@ + +import torch + +import flydsl.compiler as flyc +import flydsl.expr as fx + + +def test_coord_tensor_slicing(): + @flyc.jit + def test(i: fx.Int32, output: fx.Tensor): + row_base = fx.make_int_tuple(0) + row_layout = fx.make_layout((128, 64), (1, 0)) + coord_tensor = fx.Tensor(fx.make_view(row_base, row_layout)) + ct0 = coord_tensor[6, None] + ct1 = coord_tensor[i, None] + assert ct0[0].get_static_leaf_int == 6, f"expected static slice to be 6, got {ct0[0].get_static_leaf_int}" + output[0] = fx.get_scalar(ct0[0]) + output[1] = fx.get_scalar(ct1[0]) + + output = torch.zeros(2, dtype=torch.int32, device="cpu", requires_grad=False) + test(6, output) + assert output[0].item() == 6, f"expected 6 for static slice, got {output[0].item()}" + assert output[1].item() == 6, f"expected 6 for dynamic slice, got {output[1].item()}" From 8e39dce18ace45a67884642172b639fce59d1ac8 Mon Sep 17 00:00:00 2001 From: "Li, Tingqian" Date: Fri, 19 Jun 2026 10:38:20 +0000 Subject: [PATCH 3/5] fix python style --- tests/unit/test_coord_tensor.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unit/test_coord_tensor.py b/tests/unit/test_coord_tensor.py index 4b5e0cc795..0ff07c267f 100644 --- a/tests/unit/test_coord_tensor.py +++ b/tests/unit/test_coord_tensor.py @@ -1,4 +1,3 @@ - import torch import flydsl.compiler as flyc From 5ecd81dc3a25c60a6a110667987dbe29c15556a4 Mon Sep 17 00:00:00 2001 From: "Li, Tingqian" Date: Mon, 22 Jun 2026 03:15:36 +0000 Subject: [PATCH 4/5] fix unit test --- tests/unit/test_coord_tensor.py | 31 +++++++++++++++++++++++++++---- 1 file changed, 27 insertions(+), 4 deletions(-) diff --git a/tests/unit/test_coord_tensor.py b/tests/unit/test_coord_tensor.py index 0ff07c267f..ffbce30833 100644 --- a/tests/unit/test_coord_tensor.py +++ b/tests/unit/test_coord_tensor.py @@ -1,12 +1,26 @@ -import torch +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +from __future__ import annotations + +import pytest + +try: + import torch +except ImportError: + torch = None + +pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower] +if torch is None or not torch.cuda.is_available(): + pytest.skip("CUDA/ROCm not available", allow_module_level=True) import flydsl.compiler as flyc import flydsl.expr as fx def test_coord_tensor_slicing(): - @flyc.jit - def test(i: fx.Int32, output: fx.Tensor): + @flyc.kernel + def kernel(i: fx.Int32, output: fx.Tensor): row_base = fx.make_int_tuple(0) row_layout = fx.make_layout((128, 64), (1, 0)) coord_tensor = fx.Tensor(fx.make_view(row_base, row_layout)) @@ -16,7 +30,16 @@ def test(i: fx.Int32, output: fx.Tensor): output[0] = fx.get_scalar(ct0[0]) output[1] = fx.get_scalar(ct1[0]) - output = torch.zeros(2, dtype=torch.int32, device="cpu", requires_grad=False) + @flyc.jit + def test(i: fx.Int32, output: fx.Tensor): + kernel(i, output).launch( + grid=(1, 1, 1), + block=(1, 1, 1), + stream=fx.Stream(None) + ) + + output = torch.zeros(2, dtype=torch.int32, device="cuda", requires_grad=False) test(6, output) + torch.cuda.synchronize() assert output[0].item() == 6, f"expected 6 for static slice, got {output[0].item()}" assert output[1].item() == 6, f"expected 6 for dynamic slice, got {output[1].item()}" From ff825a48e7dd61e132977ecc1a8ea37a0d6f3e9e Mon Sep 17 00:00:00 2001 From: "Li, Tingqian" Date: Mon, 22 Jun 2026 03:20:23 +0000 Subject: [PATCH 5/5] fix coding style --- tests/unit/test_coord_tensor.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) diff --git a/tests/unit/test_coord_tensor.py b/tests/unit/test_coord_tensor.py index ffbce30833..18f22aa0d2 100644 --- a/tests/unit/test_coord_tensor.py +++ b/tests/unit/test_coord_tensor.py @@ -5,6 +5,9 @@ import pytest +import flydsl.compiler as flyc +import flydsl.expr as fx + try: import torch except ImportError: @@ -14,9 +17,6 @@ if torch is None or not torch.cuda.is_available(): pytest.skip("CUDA/ROCm not available", allow_module_level=True) -import flydsl.compiler as flyc -import flydsl.expr as fx - def test_coord_tensor_slicing(): @flyc.kernel @@ -32,11 +32,7 @@ def kernel(i: fx.Int32, output: fx.Tensor): @flyc.jit def test(i: fx.Int32, output: fx.Tensor): - kernel(i, output).launch( - grid=(1, 1, 1), - block=(1, 1, 1), - stream=fx.Stream(None) - ) + kernel(i, output).launch(grid=(1, 1, 1), block=(1, 1, 1), stream=fx.Stream(None)) output = torch.zeros(2, dtype=torch.int32, device="cuda", requires_grad=False) test(6, output)