diff --git a/include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp b/include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp index 8e8ece4..0819a10 100644 --- a/include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp +++ b/include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp @@ -2079,6 +2079,69 @@ class ReshapeConverter : public OpConversionPattern { } }; +struct GatherConverter : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + Value castIntToIndex(OpBuilder &b, Location loc, Value v) const { + return b.createOrFold(loc, b.getIndexType(), v); + } + + void createGatherPayload(OpBuilder &b, Location loc, Value input, Value index, + int64_t axis, int64_t rank) const { + SmallVector indices; + for (int i = 0; i < rank; i++) { + if (i == axis) { + indices.push_back(castIntToIndex(b, loc, index)); + } else { + indices.push_back(linalg::IndexOp::create(b, loc, i)); + } + } + // Assert index < input.sizes[axis] + auto dim = tensor::DimOp::create(b, loc, input, axis); + auto indexOverflow = arith::CmpIOp::create( + b, loc, arith::CmpIPredicate::slt, castIntToIndex(b, loc, index), dim); + cf::AssertOp::create( + b, loc, indexOverflow, + b.getStringAttr("index must be smaller than axis size")); + + // Assert index >= 0 + auto cst0 = + arith::ConstantOp::create(b, loc, b.getZeroAttr(index.getType())); + auto indexUnderflow = + arith::CmpIOp::create(b, loc, arith::CmpIPredicate::sge, index, cst0); + cf::AssertOp::create(b, loc, indexUnderflow, + b.getStringAttr("index must be larger or equal to 0")); + + Value extract = tensor::ExtractOp::create(b, loc, input, indices); + linalg::YieldOp::create(b, loc, extract); + } + + LogicalResult + matchAndRewrite(triton::GatherOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto loc = op->getLoc(); + auto src = adaptor.getSrc(); + auto indices = adaptor.getIndices(); + auto axis = op.getAxis(); + auto resultType = cast(op.getType()); + int64_t rank = resultType.getRank(); + + Value empty = tensor::EmptyOp::create(rewriter, loc, resultType.getShape(), + resultType.getElementType()); + SmallVector affineMaps(2, + rewriter.getMultiDimIdentityMap(rank)); + SmallVector iteratorTypes( + rank, utils::IteratorType::parallel); + rewriter.replaceOpWithNewOp( + op, resultType, indices, empty, affineMaps, iteratorTypes, + [&](OpBuilder &b, Location loc, ValueRange args) { + auto index = args[0]; + createGatherPayload(b, loc, src, index, axis, rank); + }); + return success(); + } +}; + class ExternElementwiseBinaryOpConverter : public OpConversionPattern { using OpConversionPattern::OpConversionPattern; diff --git a/lib/Conversion/TritonArithToLinalg/TritonArithToLinalg.cpp b/lib/Conversion/TritonArithToLinalg/TritonArithToLinalg.cpp index 2359a27..070b912 100644 --- a/lib/Conversion/TritonArithToLinalg/TritonArithToLinalg.cpp +++ b/lib/Conversion/TritonArithToLinalg/TritonArithToLinalg.cpp @@ -79,6 +79,7 @@ void mlir::triton::populateTritonArithToLinalgConversionPatterns( patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); populateExternElementwiseOpToMLIROps(patterns); diff --git a/python/examples/conftest.py b/python/examples/conftest.py index 5b00194..f7629f9 100644 --- a/python/examples/conftest.py +++ b/python/examples/conftest.py @@ -91,6 +91,7 @@ def with_allocator(): "test_trans_4d", "test_unsplat", "test_arange", + "test_gather", } annotations_tests_supported = { diff --git a/python/examples/test_gather.py b/python/examples/test_gather.py new file mode 100644 index 0000000..535d9ba --- /dev/null +++ b/python/examples/test_gather.py @@ -0,0 +1,66 @@ +import torch +import triton +import pytest + +import triton.language as tl + +@triton.jit +def gather_test_kernel(src_ptr, idx_ptr, out_ptr, axis: tl.constexpr, src_dim0: tl.constexpr, src_dim1: tl.constexpr, + src_stride0: tl.constexpr, src_stride1: tl.constexpr, idx_dim0: tl.constexpr, + idx_dim1: tl.constexpr, idx_stride0: tl.constexpr, idx_stride1: tl.constexpr, + out_dim0: tl.constexpr, out_dim1: tl.constexpr, out_stride0: tl.constexpr, + out_stride1: tl.constexpr): + src_offs = (tl.arange(0, src_dim0)[:, None] * src_stride0 + tl.arange(0, src_dim1)[None, :] * src_stride1) + src = tl.load(src_ptr + src_offs) + + idx_offs = (tl.arange(0, idx_dim0)[:, None] * idx_stride0 + tl.arange(0, idx_dim1)[None, :] * idx_stride1) + idx = tl.load(idx_ptr + idx_offs) + + out = tl.gather(src, idx, axis) + + out_offs = (tl.arange(0, out_dim0)[:, None] * out_stride0 + tl.arange(0, out_dim1)[None, :] * out_stride1) + tl.store(out_ptr + out_offs, out) + + +@triton.jit +def gather_test_kernel_1d(src_ptr, idx_ptr, out_ptr, axis: tl.constexpr, src_dim0: tl.constexpr, idx_dim0: tl.constexpr, + out_dim0: tl.constexpr): + src_offs = tl.arange(0, src_dim0) + src = tl.load(src_ptr + src_offs) + + idx_offs = tl.arange(0, idx_dim0) + idx = tl.load(idx_ptr + idx_offs) + + out = tl.gather(src, idx, axis) + + out_offs = tl.arange(0, out_dim0) + tl.store(out_ptr + out_offs, out) + + +@pytest.mark.interpreter +@pytest.mark.parametrize("src_shape, indices_shape, axis", [ + ([32], [64], 0), + ([4, 4], [8, 4], 0), + ([128, 64], [256, 64], 0), + ([128, 64], [128, 128], 1), +]) +def test_gather(src_shape, indices_shape, axis, device): + + def triton_gather(src: torch.Tensor, axis: int, indices: torch.Tensor): + output = torch.empty(indices.shape, dtype=src.dtype, device=src.device) + + if len(src_shape) == 1: + gather_test_kernel_1d[(1, )](src, indices, output, axis, src.shape[0], indices.shape[0], output.shape[0]) + else: + gather_test_kernel[(1, )](src, indices, output, axis, src.shape[0], src.shape[1], src.stride(0), + src.stride(1), indices.shape[0], indices.shape[1], indices.stride(0), + indices.stride(1), output.shape[0], output.shape[1], output.stride(0), + output.stride(1)) + + return output + + src = torch.randn(src_shape, device=device) + indices = torch.randint(0, src.shape[axis], indices_shape, device=device) + ref = torch.gather(src, axis, indices) + result = triton_gather(src, axis, indices) + torch.testing.assert_close(result, ref, rtol=0, atol=0) diff --git a/test/Conversion/TritonArithToLinalg/gather.mlir b/test/Conversion/TritonArithToLinalg/gather.mlir new file mode 100644 index 0000000..7babbc2 --- /dev/null +++ b/test/Conversion/TritonArithToLinalg/gather.mlir @@ -0,0 +1,138 @@ +// RUN: triton-shared-opt --triton-arith-to-linalg --split-input-file %s | FileCheck %s +module { + tt.func public @gather_test_kernel(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: !tt.ptr {tt.divisibility = 16 : i32}) attributes {noinline = false} { + %cst = arith.constant dense<4> : tensor<8x1xi32> + %cst_0 = arith.constant dense<4> : tensor<4x1xi32> + %0 = tt.make_range {end = 4 : i32, start = 0 : i32} : tensor<4xi32> + %1 = tt.expand_dims %0 {axis = 1 : i32} : tensor<4xi32> -> tensor<4x1xi32> + %2 = arith.muli %1, %cst_0 : tensor<4x1xi32> + %3 = tt.expand_dims %0 {axis = 0 : i32} : tensor<4xi32> -> tensor<1x4xi32> + %4 = tt.broadcast %2 : tensor<4x1xi32> -> tensor<4x4xi32> + %5 = tt.broadcast %3 : tensor<1x4xi32> -> tensor<4x4xi32> + %6 = arith.addi %4, %5 : tensor<4x4xi32> + %7 = tt.splat %arg0 : !tt.ptr -> tensor<4x4x!tt.ptr> + %8 = tt.addptr %7, %6 : tensor<4x4x!tt.ptr>, tensor<4x4xi32> + %9 = tt.load %8 : tensor<4x4x!tt.ptr> + %10 = tt.make_range {end = 8 : i32, start = 0 : i32} : tensor<8xi32> + %11 = tt.expand_dims %10 {axis = 1 : i32} : tensor<8xi32> -> tensor<8x1xi32> + %12 = arith.muli %11, %cst : tensor<8x1xi32> + %13 = tt.broadcast %12 : tensor<8x1xi32> -> tensor<8x4xi32> + %14 = tt.broadcast %3 : tensor<1x4xi32> -> tensor<8x4xi32> + %15 = arith.addi %13, %14 : tensor<8x4xi32> + %16 = tt.splat %arg1 : !tt.ptr -> tensor<8x4x!tt.ptr> + %17 = tt.addptr %16, %15 : tensor<8x4x!tt.ptr>, tensor<8x4xi32> + %18 = tt.load %17 : tensor<8x4x!tt.ptr> + %19 = tt.gather %9[%18] {axis = 0 : i32} : (tensor<4x4xf32>, tensor<8x4xi64>) -> tensor<8x4xf32> + %20 = tt.splat %arg2 : !tt.ptr -> tensor<8x4x!tt.ptr> + %21 = tt.addptr %20, %15 : tensor<8x4x!tt.ptr>, tensor<8x4xi32> + tt.store %21, %19 : tensor<8x4x!tt.ptr> + tt.return + } +} + +// CHECK: #[[$ATTR_0:.+]] = affine_map<(d0) -> (d0)> +// CHECK: #[[$ATTR_1:.+]] = affine_map<(d0, d1) -> (d0, d1)> +// CHECK: #[[$ATTR_2:.+]] = affine_map<(d0, d1) -> (d0, 0)> +// CHECK: #[[$ATTR_3:.+]] = affine_map<(d0, d1) -> (0, d1)> +// CHECK-LABEL: func.func @gather_test_kernel( +// CHECK-SAME: %[[ARG0:.*]]: !tt.ptr {tt.divisibility = 16 : i32}, %[[ARG1:.*]]: !tt.ptr {tt.divisibility = 16 : i32}, %[[ARG2:.*]]: !tt.ptr {tt.divisibility = 16 : i32}, %[[ARG3:.*]]: i32, %[[ARG4:.*]]: i32, %[[ARG5:.*]]: i32, %[[ARG6:.*]]: i32, %[[ARG7:.*]]: i32, %[[ARG8:.*]]: i32) { +// CHECK: %[[CONSTANT_0:.*]] = arith.constant 4 : index +// CHECK: %[[CONSTANT_1:.*]] = arith.constant 0 : i64 +// CHECK: %[[CONSTANT_2:.*]] = arith.constant 4 : i32 +// CHECK: %[[EMPTY_0:.*]] = tensor.empty() : tensor<8x1xi32> +// CHECK: %[[FILL_0:.*]] = linalg.fill ins(%[[CONSTANT_2]] : i32) outs(%[[EMPTY_0]] : tensor<8x1xi32>) -> tensor<8x1xi32> +// CHECK: %[[EMPTY_1:.*]] = tensor.empty() : tensor<4x1xi32> +// CHECK: %[[FILL_1:.*]] = linalg.fill ins(%[[CONSTANT_2]] : i32) outs(%[[EMPTY_1]] : tensor<4x1xi32>) -> tensor<4x1xi32> +// CHECK: %[[EMPTY_2:.*]] = tensor.empty() : tensor<4xi32> +// CHECK: %[[GENERIC_0:.*]] = linalg.generic {indexing_maps = [#[[$ATTR_0]]], iterator_types = ["parallel"]} outs(%[[EMPTY_2]] : tensor<4xi32>) { +// CHECK: ^bb0(%[[VAL_0:.*]]: i32): +// CHECK: %[[INDEX_0:.*]] = linalg.index 0 : index +// CHECK: %[[INDEX_CAST_0:.*]] = arith.index_cast %[[INDEX_0]] : index to i32 +// CHECK: linalg.yield %[[INDEX_CAST_0]] : i32 +// CHECK: } -> tensor<4xi32> +// CHECK: %[[EXPAND_SHAPE_0:.*]] = tensor.expand_shape %[[GENERIC_0]] {{\[\[}}0, 1]] output_shape [4, 1] : tensor<4xi32> into tensor<4x1xi32> +// CHECK: %[[GENERIC_1:.*]] = linalg.generic {indexing_maps = [#[[$ATTR_1]], #[[$ATTR_1]], #[[$ATTR_1]]], iterator_types = ["parallel", "parallel"]} ins(%[[EXPAND_SHAPE_0]], %[[FILL_1]] : tensor<4x1xi32>, tensor<4x1xi32>) outs(%[[EXPAND_SHAPE_0]] : tensor<4x1xi32>) { +// CHECK: ^bb0(%[[VAL_1:.*]]: i32, %[[VAL_2:.*]]: i32, %[[VAL_3:.*]]: i32): +// CHECK: %[[MULI_0:.*]] = arith.muli %[[VAL_1]], %[[VAL_2]] : i32 +// CHECK: linalg.yield %[[MULI_0]] : i32 +// CHECK: } -> tensor<4x1xi32> +// CHECK: %[[EXPAND_SHAPE_1:.*]] = tensor.expand_shape %[[GENERIC_0]] {{\[\[}}0, 1]] output_shape [1, 4] : tensor<4xi32> into tensor<1x4xi32> +// CHECK: %[[EMPTY_3:.*]] = tensor.empty() : tensor<4x4xi32> +// CHECK: %[[GENERIC_2:.*]] = linalg.generic {indexing_maps = [#[[$ATTR_2]], #[[$ATTR_1]]], iterator_types = ["parallel", "parallel"]} ins(%[[GENERIC_1]] : tensor<4x1xi32>) outs(%[[EMPTY_3]] : tensor<4x4xi32>) attrs = {broadcastDims = array} { +// CHECK: ^bb0(%[[VAL_4:.*]]: i32, %[[VAL_5:.*]]: i32): +// CHECK: linalg.yield %[[VAL_4]] : i32 +// CHECK: } -> tensor<4x4xi32> +// CHECK: %[[EMPTY_4:.*]] = tensor.empty() : tensor<4x4xi32> +// CHECK: %[[GENERIC_3:.*]] = linalg.generic {indexing_maps = [#[[$ATTR_3]], #[[$ATTR_1]]], iterator_types = ["parallel", "parallel"]} ins(%[[EXPAND_SHAPE_1]] : tensor<1x4xi32>) outs(%[[EMPTY_4]] : tensor<4x4xi32>) attrs = {broadcastDims = array} { +// CHECK: ^bb0(%[[VAL_6:.*]]: i32, %[[VAL_7:.*]]: i32): +// CHECK: linalg.yield %[[VAL_6]] : i32 +// CHECK: } -> tensor<4x4xi32> +// CHECK: %[[GENERIC_4:.*]] = linalg.generic {indexing_maps = [#[[$ATTR_1]], #[[$ATTR_1]], #[[$ATTR_1]]], iterator_types = ["parallel", "parallel"]} ins(%[[GENERIC_2]], %[[GENERIC_3]] : tensor<4x4xi32>, tensor<4x4xi32>) outs(%[[GENERIC_2]] : tensor<4x4xi32>) { +// CHECK: ^bb0(%[[VAL_8:.*]]: i32, %[[VAL_9:.*]]: i32, %[[VAL_10:.*]]: i32): +// CHECK: %[[ADDI_0:.*]] = arith.addi %[[VAL_8]], %[[VAL_9]] : i32 +// CHECK: linalg.yield %[[ADDI_0]] : i32 +// CHECK: } -> tensor<4x4xi32> +// CHECK: %[[SPLAT_0:.*]] = tensor.splat %[[ARG0]] : tensor<4x4x!tt.ptr> +// CHECK: %[[GENERIC_5:.*]] = linalg.generic {indexing_maps = [#[[$ATTR_1]], #[[$ATTR_1]], #[[$ATTR_1]]], iterator_types = ["parallel", "parallel"]} ins(%[[SPLAT_0]], %[[GENERIC_4]] : tensor<4x4x!tt.ptr>, tensor<4x4xi32>) outs(%[[SPLAT_0]] : tensor<4x4x!tt.ptr>) { +// CHECK: ^bb0(%[[VAL_11:.*]]: !tt.ptr, %[[VAL_12:.*]]: i32, %[[VAL_13:.*]]: !tt.ptr): +// CHECK: %[[ADDPTR_0:.*]] = tt.addptr %[[VAL_11]], %[[VAL_12]] : !tt.ptr, i32 +// CHECK: linalg.yield %[[ADDPTR_0]] : !tt.ptr +// CHECK: } -> tensor<4x4x!tt.ptr> +// CHECK: %[[LOAD_0:.*]] = tt.load %[[GENERIC_5]] : tensor<4x4x!tt.ptr> +// CHECK: %[[EMPTY_5:.*]] = tensor.empty() : tensor<8xi32> +// CHECK: %[[GENERIC_6:.*]] = linalg.generic {indexing_maps = [#[[$ATTR_0]]], iterator_types = ["parallel"]} outs(%[[EMPTY_5]] : tensor<8xi32>) { +// CHECK: ^bb0(%[[VAL_14:.*]]: i32): +// CHECK: %[[INDEX_1:.*]] = linalg.index 0 : index +// CHECK: %[[INDEX_CAST_1:.*]] = arith.index_cast %[[INDEX_1]] : index to i32 +// CHECK: linalg.yield %[[INDEX_CAST_1]] : i32 +// CHECK: } -> tensor<8xi32> +// CHECK: %[[EXPAND_SHAPE_2:.*]] = tensor.expand_shape %[[GENERIC_6]] {{\[\[}}0, 1]] output_shape [8, 1] : tensor<8xi32> into tensor<8x1xi32> +// CHECK: %[[GENERIC_7:.*]] = linalg.generic {indexing_maps = [#[[$ATTR_1]], #[[$ATTR_1]], #[[$ATTR_1]]], iterator_types = ["parallel", "parallel"]} ins(%[[EXPAND_SHAPE_2]], %[[FILL_0]] : tensor<8x1xi32>, tensor<8x1xi32>) outs(%[[EXPAND_SHAPE_2]] : tensor<8x1xi32>) { +// CHECK: ^bb0(%[[VAL_15:.*]]: i32, %[[VAL_16:.*]]: i32, %[[VAL_17:.*]]: i32): +// CHECK: %[[MULI_1:.*]] = arith.muli %[[VAL_15]], %[[VAL_16]] : i32 +// CHECK: linalg.yield %[[MULI_1]] : i32 +// CHECK: } -> tensor<8x1xi32> +// CHECK: %[[EMPTY_6:.*]] = tensor.empty() : tensor<8x4xi32> +// CHECK: %[[GENERIC_8:.*]] = linalg.generic {indexing_maps = [#[[$ATTR_2]], #[[$ATTR_1]]], iterator_types = ["parallel", "parallel"]} ins(%[[GENERIC_7]] : tensor<8x1xi32>) outs(%[[EMPTY_6]] : tensor<8x4xi32>) attrs = {broadcastDims = array} { +// CHECK: ^bb0(%[[VAL_18:.*]]: i32, %[[VAL_19:.*]]: i32): +// CHECK: linalg.yield %[[VAL_18]] : i32 +// CHECK: } -> tensor<8x4xi32> +// CHECK: %[[EMPTY_7:.*]] = tensor.empty() : tensor<8x4xi32> +// CHECK: %[[GENERIC_9:.*]] = linalg.generic {indexing_maps = [#[[$ATTR_3]], #[[$ATTR_1]]], iterator_types = ["parallel", "parallel"]} ins(%[[EXPAND_SHAPE_1]] : tensor<1x4xi32>) outs(%[[EMPTY_7]] : tensor<8x4xi32>) attrs = {broadcastDims = array} { +// CHECK: ^bb0(%[[VAL_20:.*]]: i32, %[[VAL_21:.*]]: i32): +// CHECK: linalg.yield %[[VAL_20]] : i32 +// CHECK: } -> tensor<8x4xi32> +// CHECK: %[[GENERIC_10:.*]] = linalg.generic {indexing_maps = [#[[$ATTR_1]], #[[$ATTR_1]], #[[$ATTR_1]]], iterator_types = ["parallel", "parallel"]} ins(%[[GENERIC_8]], %[[GENERIC_9]] : tensor<8x4xi32>, tensor<8x4xi32>) outs(%[[GENERIC_8]] : tensor<8x4xi32>) { +// CHECK: ^bb0(%[[VAL_22:.*]]: i32, %[[VAL_23:.*]]: i32, %[[VAL_24:.*]]: i32): +// CHECK: %[[ADDI_1:.*]] = arith.addi %[[VAL_22]], %[[VAL_23]] : i32 +// CHECK: linalg.yield %[[ADDI_1]] : i32 +// CHECK: } -> tensor<8x4xi32> +// CHECK: %[[SPLAT_1:.*]] = tensor.splat %[[ARG1]] : tensor<8x4x!tt.ptr> +// CHECK: %[[GENERIC_11:.*]] = linalg.generic {indexing_maps = [#[[$ATTR_1]], #[[$ATTR_1]], #[[$ATTR_1]]], iterator_types = ["parallel", "parallel"]} ins(%[[SPLAT_1]], %[[GENERIC_10]] : tensor<8x4x!tt.ptr>, tensor<8x4xi32>) outs(%[[SPLAT_1]] : tensor<8x4x!tt.ptr>) { +// CHECK: ^bb0(%[[VAL_25:.*]]: !tt.ptr, %[[VAL_26:.*]]: i32, %[[VAL_27:.*]]: !tt.ptr): +// CHECK: %[[ADDPTR_1:.*]] = tt.addptr %[[VAL_25]], %[[VAL_26]] : !tt.ptr, i32 +// CHECK: linalg.yield %[[ADDPTR_1]] : !tt.ptr +// CHECK: } -> tensor<8x4x!tt.ptr> +// CHECK: %[[LOAD_1:.*]] = tt.load %[[GENERIC_11]] : tensor<8x4x!tt.ptr> +// CHECK: %[[EMPTY_8:.*]] = tensor.empty() : tensor<8x4xf32> +// CHECK: %[[GENERIC_12:.*]] = linalg.generic {indexing_maps = [#[[$ATTR_1]], #[[$ATTR_1]]], iterator_types = ["parallel", "parallel"]} ins(%[[LOAD_1]] : tensor<8x4xi64>) outs(%[[EMPTY_8]] : tensor<8x4xf32>) { +// CHECK: ^bb0(%[[VAL_28:.*]]: i64, %[[VAL_29:.*]]: f32): +// CHECK: %[[INDEX_CAST_2:.*]] = arith.index_cast %[[VAL_28]] : i64 to index +// CHECK: %[[INDEX_2:.*]] = linalg.index 1 : index +// CHECK: %[[INDEX_CAST_3:.*]] = arith.index_cast %[[VAL_28]] : i64 to index +// CHECK: %[[CMPI_0:.*]] = arith.cmpi slt, %[[INDEX_CAST_3]], %[[CONSTANT_0]] : index +// CHECK: cf.assert %[[CMPI_0]], "index must be smaller than axis size" +// CHECK: %[[CMPI_1:.*]] = arith.cmpi sge, %[[VAL_28]], %[[CONSTANT_1]] : i64 +// CHECK: cf.assert %[[CMPI_1]], "index must be larger or equal to 0" +// CHECK: %[[EXTRACT_0:.*]] = tensor.extract %[[LOAD_0]]{{\[}}%[[INDEX_CAST_2]], %[[INDEX_2]]] : tensor<4x4xf32> +// CHECK: linalg.yield %[[EXTRACT_0]] : f32 +// CHECK: } -> tensor<8x4xf32> +// CHECK: %[[SPLAT_2:.*]] = tensor.splat %[[ARG2]] : tensor<8x4x!tt.ptr> +// CHECK: %[[GENERIC_13:.*]] = linalg.generic {indexing_maps = [#[[$ATTR_1]], #[[$ATTR_1]], #[[$ATTR_1]]], iterator_types = ["parallel", "parallel"]} ins(%[[SPLAT_2]], %[[GENERIC_10]] : tensor<8x4x!tt.ptr>, tensor<8x4xi32>) outs(%[[SPLAT_2]] : tensor<8x4x!tt.ptr>) { +// CHECK: ^bb0(%[[VAL_30:.*]]: !tt.ptr, %[[VAL_31:.*]]: i32, %[[VAL_32:.*]]: !tt.ptr): +// CHECK: %[[ADDPTR_2:.*]] = tt.addptr %[[VAL_30]], %[[VAL_31]] : !tt.ptr, i32 +// CHECK: linalg.yield %[[ADDPTR_2]] : !tt.ptr +// CHECK: } -> tensor<8x4x!tt.ptr> +// CHECK: tt.store %[[GENERIC_13]], %[[GENERIC_12]] : tensor<8x4x!tt.ptr> +// CHECK: return +// CHECK: } \ No newline at end of file