diff --git a/ptodsl/lib/a5/README.md b/ptodsl/lib/a5/README.md index ed5b600a..df755ec3 100644 --- a/ptodsl/lib/a5/README.md +++ b/ptodsl/lib/a5/README.md @@ -33,6 +33,8 @@ Entry points: - [`kernels.py`](./kernels.py): translated example kernels, including the no-section `build_hivm_vadd_demo()` flow that lowers through PTOAS VPTO into `llvm.hivm.*` intrinsics +- [`tile_op_kernels.py`](./tile_op_kernels.py): one builder per public tile op + that already has a micro-backed implementation - [`generated`](./generated): emitted `.pto` artifacts from `scripts/generate_a5_pto.py` Recommended usage: @@ -56,6 +58,14 @@ PYTHONPATH=/Users/zhoubot/github/.llvm-19.1.7/build-mlir-py312/tools/mlir/python /Users/zhoubot/github/.venv-ptoas-src312/bin/python scripts/generate_a5_pto.py --emit-hivm-llvm ``` +To emit one `.pto` file per public tile op that already has a micro-backed +implementation: + +```bash +PYTHONPATH=/Users/zhoubot/github/.llvm-19.1.7/build-mlir-py312/tools/mlir/python_packages/mlir_core:/Users/zhoubot/github/pto-org/PTOAS/install-src312:/Users/zhoubot/github/pto-org/PTOAS/build-src312/python \ +/Users/zhoubot/github/.venv-ptoas-src312/bin/python scripts/generate_tile_op_pto.py +``` + `--emit-cpp` and `--emit-hivm-llvm` are intentionally asymmetric: - pure micro kernels such as `a5_hivm_vadd_demo` now lower end-to-end through PTOAS VPTO into `llvm.hivm.*` diff --git a/ptodsl/lib/a5/__init__.py b/ptodsl/lib/a5/__init__.py index 87a72335..e3fab8a3 100644 --- a/ptodsl/lib/a5/__init__.py +++ b/ptodsl/lib/a5/__init__.py @@ -10,6 +10,11 @@ build_templated_elementwise_add, build_vector_copy, ) +from .tile_op_kernels import ( + TILE_OP_KERNEL_BUILDERS, + TILE_OP_KERNEL_SPECS, + tile_op_generation_index_markdown, +) from .ops import * from .tile_micro_coverage import ( TILE_MICRO_COVERAGE, @@ -21,6 +26,8 @@ "A5_HEADER_COVERAGE", "HIVM_LLVM_KERNELS", "KERNEL_BUILDERS", + "TILE_OP_KERNEL_BUILDERS", + "TILE_OP_KERNEL_SPECS", "TILE_MICRO_COVERAGE", "a5_header_coverage_markdown", "build_cube_matmul", @@ -37,5 +44,6 @@ "treduce", "tscalar", "tsort", + "tile_op_generation_index_markdown", "tunary", ] diff --git a/ptodsl/lib/a5/generated/tile_ops/TILE_OP_GENERATION_INDEX.md b/ptodsl/lib/a5/generated/tile_ops/TILE_OP_GENERATION_INDEX.md new file mode 100644 index 00000000..263e3a05 --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/TILE_OP_GENERATION_INDEX.md @@ -0,0 +1,38 @@ +# Tile Op PTO Generation + +| tile op | status | artifact | note | +| --- | --- | --- | --- | +| `mov` | `generated` | `tile_ops/mov.pto` | UB stage + vlds/vsts copy loop. | +| `add` | `generated` | `tile_ops/add.pto` | UB stage + constexpr-specialized TBinOp-style vlds/vadd/vsts lowering. | +| `sub` | `generated` | `tile_ops/sub.pto` | UB stage + constexpr-specialized TBinOp-style vlds/vsub/vsts lowering. | +| `div` | `generated` | `tile_ops/div.pto` | UB stage + constexpr-specialized TBinOp-style vlds/vdiv/vsts lowering. | +| `mul` | `generated` | `tile_ops/mul.pto` | UB stage + constexpr-specialized TBinOp-style vlds/vmul/vsts lowering. | +| `or_` | `generated` | `tile_ops/or_.pto` | UB stage + constexpr-specialized TBinOp-style vlds/vor/vsts lowering. | +| `gather` | `generated` | `tile_ops/gather.pto` | Indexed gather is implemented via vgather2 for same-width source/index pairs; mask-pattern gather still needs unsupported vsqz-style micro support. | +| `exp` | `generated` | `tile_ops/exp.pto` | UB stage + vlds/vexp/vsts loop. | +| `log` | `generated` | `tile_ops/log.pto` | UB stage + vlds/vln/vsts loop. | +| `relu` | `generated` | `tile_ops/relu.pto` | UB stage + vlds/vrelu/vsts loop. | +| `abs` | `generated` | `tile_ops/abs.pto` | UB stage + vlds/vabs/vsts loop. | +| `sqrt` | `generated` | `tile_ops/sqrt.pto` | UB stage + vlds/vsqrt/vsts loop. | +| `rsqrt` | `generated` | `tile_ops/rsqrt.pto` | UB stage + vsqrt/vrec sequence. | +| `reciprocal` | `generated` | `tile_ops/reciprocal.pto` | UB stage + vlds/vrec/vsts loop. | +| `matmul` | `blocked` | - | Cube/L0 path is not a pure vector-micro rewrite target. | +| `matmul_bias` | `blocked` | - | Cube/L0 path is not a pure vector-micro rewrite target. | +| `matmul_acc` | `blocked` | - | Cube/L0 path is not a pure vector-micro rewrite target. | +| `extract` | `blocked` | - | Layout/L0 extraction op, not a vector-micro compute rewrite. | +| `row_sum` | `generated` | `tile_ops/row_sum.pto` | Static-shape row reduction via vcadd + point-store. | +| `row_min` | `generated` | `tile_ops/row_min.pto` | Static-shape row reduction via vcmin + point-store. | +| `row_max` | `generated` | `tile_ops/row_max.pto` | Static-shape row reduction via vcmax + point-store. | +| `row_prod` | `blocked` | - | No row-product micro lowering is wired yet. | +| `row_expand` | `generated` | `tile_ops/row_expand.pto` | Static-shape canonical broadcast via vldas/vldus/vdup/vsts. | +| `row_expand_sub` | `generated` | `tile_ops/row_expand_sub.pto` | Static-shape canonical broadcast via vldas/vldus/vdup/vsub/vsts. | +| `row_expand_div` | `generated` | `tile_ops/row_expand_div.pto` | Static-shape canonical broadcast via vldas/vldus/vdup/vdiv/vsts. | +| `row_expand_mul` | `generated` | `tile_ops/row_expand_mul.pto` | Static-shape canonical broadcast via vldas/vldus/vdup/vmul/vsts. | +| `col_sum` | `generated` | `tile_ops/col_sum.pto` | Static-shape TColReduceOps-style column reduction via vadd. | +| `col_min` | `generated` | `tile_ops/col_min.pto` | Static-shape TColReduceOps-style column reduction via vmin. | +| `col_max` | `generated` | `tile_ops/col_max.pto` | Static-shape TColReduceOps-style column reduction via vmax. | +| `col_prod` | `blocked` | - | No column-product micro lowering is wired yet. | +| `col_expand` | `generated` | `tile_ops/col_expand.pto` | Static-shape canonical broadcast via vlds/vsts replication. | +| `mrgsort` | `generated` | `tile_ops/mrgsort.pto` | Single-list row-major merge sort via vmrgsort4. | +| `sort32` | `generated` | `tile_ops/sort32.pto` | Static-shape block sort via vbitsort. | +| `subset` | `not_applicable` | - | View helper only, not a tile compute op. | diff --git a/ptodsl/lib/a5/generated/tile_ops/abs.pto b/ptodsl/lib/a5/generated/tile_ops/abs.pto new file mode 100644 index 00000000..ef021091 --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/abs.pto @@ -0,0 +1,75 @@ +module { + func.func @tile_op_abs(%arg0: !pto.ptr, %arg1: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c64_1 = arith.constant 64 : index + %c1_2 = arith.constant 1 : index + %c8_3 = arith.constant 8 : index + %c64_4 = arith.constant 64 : index + %1 = pto.make_tensor_view %arg1, shape = [%c8_3, %c64_4], strides = [%c64_1, %c1_2] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_5 = arith.constant 0 : index + %c8_6 = arith.constant 8 : index + %c64_7 = arith.constant 64 : index + %2 = pto.partition_view %0, offsets = [%c0, %c0_5], sizes = [%c8_6, %c64_7] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_8 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c8_10 = arith.constant 8 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %1, offsets = [%c0_8, %c0_9], sizes = [%c8_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %4 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %5 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + pto.tload ins(%2 : !pto.partition_tensor_view<8x64xf32>) outs(%4 : !pto.tile_buf) + %6 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %7 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %8 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c0_12 = arith.constant 0 : index + %9 = pto.vlds %6[%c0_12] : !pto.ptr -> !pto.vreg<64xf32> + %10 = pto.vabs %9, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %10, %7[%c0_12], %8 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %11 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c64_13 = arith.constant 64 : index + %12 = pto.vlds %6[%c64_13] : !pto.ptr -> !pto.vreg<64xf32> + %13 = pto.vabs %12, %11 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %13, %7[%c64_13], %11 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %14 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c128 = arith.constant 128 : index + %15 = pto.vlds %6[%c128] : !pto.ptr -> !pto.vreg<64xf32> + %16 = pto.vabs %15, %14 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %16, %7[%c128], %14 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %17 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c192 = arith.constant 192 : index + %18 = pto.vlds %6[%c192] : !pto.ptr -> !pto.vreg<64xf32> + %19 = pto.vabs %18, %17 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %19, %7[%c192], %17 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %20 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c256 = arith.constant 256 : index + %21 = pto.vlds %6[%c256] : !pto.ptr -> !pto.vreg<64xf32> + %22 = pto.vabs %21, %20 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %22, %7[%c256], %20 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %23 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c320 = arith.constant 320 : index + %24 = pto.vlds %6[%c320] : !pto.ptr -> !pto.vreg<64xf32> + %25 = pto.vabs %24, %23 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %25, %7[%c320], %23 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %26 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c384 = arith.constant 384 : index + %27 = pto.vlds %6[%c384] : !pto.ptr -> !pto.vreg<64xf32> + %28 = pto.vabs %27, %26 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %28, %7[%c384], %26 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %29 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c448 = arith.constant 448 : index + %30 = pto.vlds %6[%c448] : !pto.ptr -> !pto.vreg<64xf32> + %31 = pto.vabs %30, %29 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %31, %7[%c448], %29 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + pto.tstore ins(%5 : !pto.tile_buf) outs(%3 : !pto.partition_tensor_view<8x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/add.pto b/ptodsl/lib/a5/generated/tile_ops/add.pto new file mode 100644 index 00000000..5b4d9507 --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/add.pto @@ -0,0 +1,90 @@ +module { + func.func @tile_op_add(%arg0: !pto.ptr, %arg1: !pto.ptr, %arg2: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c64_1 = arith.constant 64 : index + %c1_2 = arith.constant 1 : index + %c8_3 = arith.constant 8 : index + %c64_4 = arith.constant 64 : index + %1 = pto.make_tensor_view %arg1, shape = [%c8_3, %c64_4], strides = [%c64_1, %c1_2] : !pto.tensor_view + %c64_5 = arith.constant 64 : index + %c1_6 = arith.constant 1 : index + %c8_7 = arith.constant 8 : index + %c64_8 = arith.constant 64 : index + %2 = pto.make_tensor_view %arg2, shape = [%c8_7, %c64_8], strides = [%c64_5, %c1_6] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c8_10 = arith.constant 8 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %0, offsets = [%c0, %c0_9], sizes = [%c8_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_12 = arith.constant 0 : index + %c0_13 = arith.constant 0 : index + %c8_14 = arith.constant 8 : index + %c64_15 = arith.constant 64 : index + %4 = pto.partition_view %1, offsets = [%c0_12, %c0_13], sizes = [%c8_14, %c64_15] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_16 = arith.constant 0 : index + %c0_17 = arith.constant 0 : index + %c8_18 = arith.constant 8 : index + %c64_19 = arith.constant 64 : index + %5 = pto.partition_view %2, offsets = [%c0_16, %c0_17], sizes = [%c8_18, %c64_19] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %c4096_i64 = arith.constant 4096 : i64 + %6 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %7 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + %8 = pto.alloc_tile addr = %c4096_i64 : !pto.tile_buf + pto.tload ins(%3 : !pto.partition_tensor_view<8x64xf32>) outs(%6 : !pto.tile_buf) + pto.tload ins(%4 : !pto.partition_tensor_view<8x64xf32>) outs(%7 : !pto.tile_buf) + %9 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %10 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %11 = pto.castptr %c4096_i64 : i64 -> !pto.ptr + %c64_20 = arith.constant 64 : index + %12 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result, %updated_source = pto.vlds_post %9[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_21, %updated_source_22 = pto.vlds_post %10[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %13 = pto.vadd %result, %result_21, %12 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %14 = pto.vsts_post %13, %11[%c64_20], %12 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %15 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_23, %updated_source_24 = pto.vlds_post %updated_source[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_25, %updated_source_26 = pto.vlds_post %updated_source_22[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %16 = pto.vadd %result_23, %result_25, %15 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %17 = pto.vsts_post %16, %14[%c64_20], %15 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %18 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_27, %updated_source_28 = pto.vlds_post %updated_source_24[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_29, %updated_source_30 = pto.vlds_post %updated_source_26[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %19 = pto.vadd %result_27, %result_29, %18 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %20 = pto.vsts_post %19, %17[%c64_20], %18 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %21 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_31, %updated_source_32 = pto.vlds_post %updated_source_28[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_33, %updated_source_34 = pto.vlds_post %updated_source_30[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %22 = pto.vadd %result_31, %result_33, %21 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %23 = pto.vsts_post %22, %20[%c64_20], %21 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %24 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_35, %updated_source_36 = pto.vlds_post %updated_source_32[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_37, %updated_source_38 = pto.vlds_post %updated_source_34[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %25 = pto.vadd %result_35, %result_37, %24 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %26 = pto.vsts_post %25, %23[%c64_20], %24 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %27 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_39, %updated_source_40 = pto.vlds_post %updated_source_36[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_41, %updated_source_42 = pto.vlds_post %updated_source_38[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %28 = pto.vadd %result_39, %result_41, %27 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %29 = pto.vsts_post %28, %26[%c64_20], %27 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %30 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_43, %updated_source_44 = pto.vlds_post %updated_source_40[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_45, %updated_source_46 = pto.vlds_post %updated_source_42[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %31 = pto.vadd %result_43, %result_45, %30 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %32 = pto.vsts_post %31, %29[%c64_20], %30 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %33 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_47, %updated_source_48 = pto.vlds_post %updated_source_44[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_49, %updated_source_50 = pto.vlds_post %updated_source_46[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %34 = pto.vadd %result_47, %result_49, %33 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %35 = pto.vsts_post %34, %32[%c64_20], %33 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + pto.tstore ins(%8 : !pto.tile_buf) outs(%5 : !pto.partition_tensor_view<8x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/col_expand.pto b/ptodsl/lib/a5/generated/tile_ops/col_expand.pto new file mode 100644 index 00000000..7b987b15 --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/col_expand.pto @@ -0,0 +1,54 @@ +module { + func.func @tile_op_col_expand(%arg0: !pto.ptr, %arg1: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c1_0 = arith.constant 1 : index + %c64_1 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c1_0, %c64_1], strides = [%c64, %c1] : !pto.tensor_view + %c64_2 = arith.constant 64 : index + %c1_3 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_4 = arith.constant 64 : index + %1 = pto.make_tensor_view %arg1, shape = [%c8, %c64_4], strides = [%c64_2, %c1_3] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_5 = arith.constant 0 : index + %c1_6 = arith.constant 1 : index + %c64_7 = arith.constant 64 : index + %2 = pto.partition_view %0, offsets = [%c0, %c0_5], sizes = [%c1_6, %c64_7] : !pto.tensor_view -> !pto.partition_tensor_view<1x64xf32> + %c0_8 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c8_10 = arith.constant 8 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %1, offsets = [%c0_8, %c0_9], sizes = [%c8_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %4 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %5 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + pto.tload ins(%2 : !pto.partition_tensor_view<1x64xf32>) outs(%4 : !pto.tile_buf) + %6 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %7 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %8 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c0_12 = arith.constant 0 : index + %9 = pto.vlds %6[%c0_12] : !pto.ptr -> !pto.vreg<64xf32> + %c0_13 = arith.constant 0 : index + pto.vsts %9, %7[%c0_13], %8 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c64_14 = arith.constant 64 : index + pto.vsts %9, %7[%c64_14], %8 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c128 = arith.constant 128 : index + pto.vsts %9, %7[%c128], %8 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c192 = arith.constant 192 : index + pto.vsts %9, %7[%c192], %8 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c256 = arith.constant 256 : index + pto.vsts %9, %7[%c256], %8 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c320 = arith.constant 320 : index + pto.vsts %9, %7[%c320], %8 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c384 = arith.constant 384 : index + pto.vsts %9, %7[%c384], %8 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c448 = arith.constant 448 : index + pto.vsts %9, %7[%c448], %8 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + pto.tstore ins(%5 : !pto.tile_buf) outs(%3 : !pto.partition_tensor_view<8x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/col_max.pto b/ptodsl/lib/a5/generated/tile_ops/col_max.pto new file mode 100644 index 00000000..55459b11 --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/col_max.pto @@ -0,0 +1,93 @@ +module { + func.func @tile_op_col_max(%arg0: !pto.ptr, %arg1: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c64_1 = arith.constant 64 : index + %c1_2 = arith.constant 1 : index + %c1_3 = arith.constant 1 : index + %c64_4 = arith.constant 64 : index + %1 = pto.make_tensor_view %arg1, shape = [%c1_3, %c64_4], strides = [%c64_1, %c1_2] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_5 = arith.constant 0 : index + %c8_6 = arith.constant 8 : index + %c64_7 = arith.constant 64 : index + %2 = pto.partition_view %0, offsets = [%c0, %c0_5], sizes = [%c8_6, %c64_7] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_8 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c1_10 = arith.constant 1 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %1, offsets = [%c0_8, %c0_9], sizes = [%c1_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<1x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %4 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %5 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + pto.tload ins(%2 : !pto.partition_tensor_view<8x64xf32>) outs(%4 : !pto.tile_buf) + %6 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %7 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %cst = arith.constant 0xFF800000 : f32 + %8 = pto.vbr %cst : f32 -> !pto.vreg<64xf32> + %9 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c64_12 = arith.constant 64 : index + %10 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c0_13 = arith.constant 0 : index + %11 = pto.addptr %6, %c0_13 : -> + %c64_14 = arith.constant 64 : index + %result, %updated_source = pto.vlds_post %11[%c64_14] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %12 = pto.pset_b32 "PAT_ALL" : !pto.mask + %13 = pto.vsel %result, %8, %12 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %14 = pto.vsel %13, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %15 = pto.vmax %8, %14, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_15 = arith.constant 64 : index + %result_16, %updated_source_17 = pto.vlds_post %updated_source[%c64_15] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %16 = pto.pset_b32 "PAT_ALL" : !pto.mask + %17 = pto.vsel %result_16, %8, %16 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %18 = pto.vsel %17, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %19 = pto.vmax %15, %18, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_18 = arith.constant 64 : index + %result_19, %updated_source_20 = pto.vlds_post %updated_source_17[%c64_18] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %20 = pto.pset_b32 "PAT_ALL" : !pto.mask + %21 = pto.vsel %result_19, %8, %20 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %22 = pto.vsel %21, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %23 = pto.vmax %19, %22, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_21 = arith.constant 64 : index + %result_22, %updated_source_23 = pto.vlds_post %updated_source_20[%c64_21] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %24 = pto.pset_b32 "PAT_ALL" : !pto.mask + %25 = pto.vsel %result_22, %8, %24 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %26 = pto.vsel %25, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %27 = pto.vmax %23, %26, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_24 = arith.constant 64 : index + %result_25, %updated_source_26 = pto.vlds_post %updated_source_23[%c64_24] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %28 = pto.pset_b32 "PAT_ALL" : !pto.mask + %29 = pto.vsel %result_25, %8, %28 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %30 = pto.vsel %29, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %31 = pto.vmax %27, %30, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_27 = arith.constant 64 : index + %result_28, %updated_source_29 = pto.vlds_post %updated_source_26[%c64_27] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %32 = pto.pset_b32 "PAT_ALL" : !pto.mask + %33 = pto.vsel %result_28, %8, %32 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %34 = pto.vsel %33, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %35 = pto.vmax %31, %34, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_30 = arith.constant 64 : index + %result_31, %updated_source_32 = pto.vlds_post %updated_source_29[%c64_30] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %36 = pto.pset_b32 "PAT_ALL" : !pto.mask + %37 = pto.vsel %result_31, %8, %36 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %38 = pto.vsel %37, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %39 = pto.vmax %35, %38, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_33 = arith.constant 64 : index + %result_34, %updated_source_35 = pto.vlds_post %updated_source_32[%c64_33] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %40 = pto.pset_b32 "PAT_ALL" : !pto.mask + %41 = pto.vsel %result_34, %8, %40 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %42 = pto.vsel %41, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %43 = pto.vmax %39, %42, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c0_36 = arith.constant 0 : index + %44 = pto.addptr %7, %c0_36 : -> + %45 = pto.vsts_post %43, %44[%c64_12], %10 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + pto.tstore ins(%5 : !pto.tile_buf) outs(%3 : !pto.partition_tensor_view<1x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/col_min.pto b/ptodsl/lib/a5/generated/tile_ops/col_min.pto new file mode 100644 index 00000000..665e0a4f --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/col_min.pto @@ -0,0 +1,93 @@ +module { + func.func @tile_op_col_min(%arg0: !pto.ptr, %arg1: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c64_1 = arith.constant 64 : index + %c1_2 = arith.constant 1 : index + %c1_3 = arith.constant 1 : index + %c64_4 = arith.constant 64 : index + %1 = pto.make_tensor_view %arg1, shape = [%c1_3, %c64_4], strides = [%c64_1, %c1_2] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_5 = arith.constant 0 : index + %c8_6 = arith.constant 8 : index + %c64_7 = arith.constant 64 : index + %2 = pto.partition_view %0, offsets = [%c0, %c0_5], sizes = [%c8_6, %c64_7] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_8 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c1_10 = arith.constant 1 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %1, offsets = [%c0_8, %c0_9], sizes = [%c1_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<1x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %4 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %5 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + pto.tload ins(%2 : !pto.partition_tensor_view<8x64xf32>) outs(%4 : !pto.tile_buf) + %6 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %7 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %cst = arith.constant 0x7F800000 : f32 + %8 = pto.vbr %cst : f32 -> !pto.vreg<64xf32> + %9 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c64_12 = arith.constant 64 : index + %10 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c0_13 = arith.constant 0 : index + %11 = pto.addptr %6, %c0_13 : -> + %c64_14 = arith.constant 64 : index + %result, %updated_source = pto.vlds_post %11[%c64_14] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %12 = pto.pset_b32 "PAT_ALL" : !pto.mask + %13 = pto.vsel %result, %8, %12 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %14 = pto.vsel %13, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %15 = pto.vmin %8, %14, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_15 = arith.constant 64 : index + %result_16, %updated_source_17 = pto.vlds_post %updated_source[%c64_15] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %16 = pto.pset_b32 "PAT_ALL" : !pto.mask + %17 = pto.vsel %result_16, %8, %16 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %18 = pto.vsel %17, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %19 = pto.vmin %15, %18, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_18 = arith.constant 64 : index + %result_19, %updated_source_20 = pto.vlds_post %updated_source_17[%c64_18] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %20 = pto.pset_b32 "PAT_ALL" : !pto.mask + %21 = pto.vsel %result_19, %8, %20 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %22 = pto.vsel %21, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %23 = pto.vmin %19, %22, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_21 = arith.constant 64 : index + %result_22, %updated_source_23 = pto.vlds_post %updated_source_20[%c64_21] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %24 = pto.pset_b32 "PAT_ALL" : !pto.mask + %25 = pto.vsel %result_22, %8, %24 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %26 = pto.vsel %25, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %27 = pto.vmin %23, %26, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_24 = arith.constant 64 : index + %result_25, %updated_source_26 = pto.vlds_post %updated_source_23[%c64_24] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %28 = pto.pset_b32 "PAT_ALL" : !pto.mask + %29 = pto.vsel %result_25, %8, %28 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %30 = pto.vsel %29, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %31 = pto.vmin %27, %30, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_27 = arith.constant 64 : index + %result_28, %updated_source_29 = pto.vlds_post %updated_source_26[%c64_27] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %32 = pto.pset_b32 "PAT_ALL" : !pto.mask + %33 = pto.vsel %result_28, %8, %32 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %34 = pto.vsel %33, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %35 = pto.vmin %31, %34, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_30 = arith.constant 64 : index + %result_31, %updated_source_32 = pto.vlds_post %updated_source_29[%c64_30] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %36 = pto.pset_b32 "PAT_ALL" : !pto.mask + %37 = pto.vsel %result_31, %8, %36 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %38 = pto.vsel %37, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %39 = pto.vmin %35, %38, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_33 = arith.constant 64 : index + %result_34, %updated_source_35 = pto.vlds_post %updated_source_32[%c64_33] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %40 = pto.pset_b32 "PAT_ALL" : !pto.mask + %41 = pto.vsel %result_34, %8, %40 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %42 = pto.vsel %41, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %43 = pto.vmin %39, %42, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c0_36 = arith.constant 0 : index + %44 = pto.addptr %7, %c0_36 : -> + %45 = pto.vsts_post %43, %44[%c64_12], %10 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + pto.tstore ins(%5 : !pto.tile_buf) outs(%3 : !pto.partition_tensor_view<1x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/col_sum.pto b/ptodsl/lib/a5/generated/tile_ops/col_sum.pto new file mode 100644 index 00000000..78b09c22 --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/col_sum.pto @@ -0,0 +1,93 @@ +module { + func.func @tile_op_col_sum(%arg0: !pto.ptr, %arg1: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c64_1 = arith.constant 64 : index + %c1_2 = arith.constant 1 : index + %c1_3 = arith.constant 1 : index + %c64_4 = arith.constant 64 : index + %1 = pto.make_tensor_view %arg1, shape = [%c1_3, %c64_4], strides = [%c64_1, %c1_2] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_5 = arith.constant 0 : index + %c8_6 = arith.constant 8 : index + %c64_7 = arith.constant 64 : index + %2 = pto.partition_view %0, offsets = [%c0, %c0_5], sizes = [%c8_6, %c64_7] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_8 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c1_10 = arith.constant 1 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %1, offsets = [%c0_8, %c0_9], sizes = [%c1_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<1x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %4 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %5 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + pto.tload ins(%2 : !pto.partition_tensor_view<8x64xf32>) outs(%4 : !pto.tile_buf) + %6 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %7 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %cst = arith.constant 0.000000e+00 : f32 + %8 = pto.vbr %cst : f32 -> !pto.vreg<64xf32> + %9 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c64_12 = arith.constant 64 : index + %10 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c0_13 = arith.constant 0 : index + %11 = pto.addptr %6, %c0_13 : -> + %c64_14 = arith.constant 64 : index + %result, %updated_source = pto.vlds_post %11[%c64_14] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %12 = pto.pset_b32 "PAT_ALL" : !pto.mask + %13 = pto.vsel %result, %8, %12 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %14 = pto.vsel %13, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %15 = pto.vadd %8, %14, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_15 = arith.constant 64 : index + %result_16, %updated_source_17 = pto.vlds_post %updated_source[%c64_15] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %16 = pto.pset_b32 "PAT_ALL" : !pto.mask + %17 = pto.vsel %result_16, %8, %16 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %18 = pto.vsel %17, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %19 = pto.vadd %15, %18, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_18 = arith.constant 64 : index + %result_19, %updated_source_20 = pto.vlds_post %updated_source_17[%c64_18] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %20 = pto.pset_b32 "PAT_ALL" : !pto.mask + %21 = pto.vsel %result_19, %8, %20 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %22 = pto.vsel %21, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %23 = pto.vadd %19, %22, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_21 = arith.constant 64 : index + %result_22, %updated_source_23 = pto.vlds_post %updated_source_20[%c64_21] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %24 = pto.pset_b32 "PAT_ALL" : !pto.mask + %25 = pto.vsel %result_22, %8, %24 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %26 = pto.vsel %25, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %27 = pto.vadd %23, %26, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_24 = arith.constant 64 : index + %result_25, %updated_source_26 = pto.vlds_post %updated_source_23[%c64_24] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %28 = pto.pset_b32 "PAT_ALL" : !pto.mask + %29 = pto.vsel %result_25, %8, %28 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %30 = pto.vsel %29, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %31 = pto.vadd %27, %30, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_27 = arith.constant 64 : index + %result_28, %updated_source_29 = pto.vlds_post %updated_source_26[%c64_27] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %32 = pto.pset_b32 "PAT_ALL" : !pto.mask + %33 = pto.vsel %result_28, %8, %32 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %34 = pto.vsel %33, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %35 = pto.vadd %31, %34, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_30 = arith.constant 64 : index + %result_31, %updated_source_32 = pto.vlds_post %updated_source_29[%c64_30] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %36 = pto.pset_b32 "PAT_ALL" : !pto.mask + %37 = pto.vsel %result_31, %8, %36 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %38 = pto.vsel %37, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %39 = pto.vadd %35, %38, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_33 = arith.constant 64 : index + %result_34, %updated_source_35 = pto.vlds_post %updated_source_32[%c64_33] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %40 = pto.pset_b32 "PAT_ALL" : !pto.mask + %41 = pto.vsel %result_34, %8, %40 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %42 = pto.vsel %41, %8, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %43 = pto.vadd %39, %42, %9 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c0_36 = arith.constant 0 : index + %44 = pto.addptr %7, %c0_36 : -> + %45 = pto.vsts_post %43, %44[%c64_12], %10 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + pto.tstore ins(%5 : !pto.tile_buf) outs(%3 : !pto.partition_tensor_view<1x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/div.pto b/ptodsl/lib/a5/generated/tile_ops/div.pto new file mode 100644 index 00000000..ac2fb1e4 --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/div.pto @@ -0,0 +1,90 @@ +module { + func.func @tile_op_div(%arg0: !pto.ptr, %arg1: !pto.ptr, %arg2: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c64_1 = arith.constant 64 : index + %c1_2 = arith.constant 1 : index + %c8_3 = arith.constant 8 : index + %c64_4 = arith.constant 64 : index + %1 = pto.make_tensor_view %arg1, shape = [%c8_3, %c64_4], strides = [%c64_1, %c1_2] : !pto.tensor_view + %c64_5 = arith.constant 64 : index + %c1_6 = arith.constant 1 : index + %c8_7 = arith.constant 8 : index + %c64_8 = arith.constant 64 : index + %2 = pto.make_tensor_view %arg2, shape = [%c8_7, %c64_8], strides = [%c64_5, %c1_6] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c8_10 = arith.constant 8 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %0, offsets = [%c0, %c0_9], sizes = [%c8_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_12 = arith.constant 0 : index + %c0_13 = arith.constant 0 : index + %c8_14 = arith.constant 8 : index + %c64_15 = arith.constant 64 : index + %4 = pto.partition_view %1, offsets = [%c0_12, %c0_13], sizes = [%c8_14, %c64_15] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_16 = arith.constant 0 : index + %c0_17 = arith.constant 0 : index + %c8_18 = arith.constant 8 : index + %c64_19 = arith.constant 64 : index + %5 = pto.partition_view %2, offsets = [%c0_16, %c0_17], sizes = [%c8_18, %c64_19] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %c4096_i64 = arith.constant 4096 : i64 + %6 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %7 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + %8 = pto.alloc_tile addr = %c4096_i64 : !pto.tile_buf + pto.tload ins(%3 : !pto.partition_tensor_view<8x64xf32>) outs(%6 : !pto.tile_buf) + pto.tload ins(%4 : !pto.partition_tensor_view<8x64xf32>) outs(%7 : !pto.tile_buf) + %9 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %10 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %11 = pto.castptr %c4096_i64 : i64 -> !pto.ptr + %c64_20 = arith.constant 64 : index + %12 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result, %updated_source = pto.vlds_post %9[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_21, %updated_source_22 = pto.vlds_post %10[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %13 = pto.vdiv %result, %result_21, %12 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %14 = pto.vsts_post %13, %11[%c64_20], %12 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %15 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_23, %updated_source_24 = pto.vlds_post %updated_source[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_25, %updated_source_26 = pto.vlds_post %updated_source_22[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %16 = pto.vdiv %result_23, %result_25, %15 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %17 = pto.vsts_post %16, %14[%c64_20], %15 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %18 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_27, %updated_source_28 = pto.vlds_post %updated_source_24[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_29, %updated_source_30 = pto.vlds_post %updated_source_26[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %19 = pto.vdiv %result_27, %result_29, %18 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %20 = pto.vsts_post %19, %17[%c64_20], %18 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %21 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_31, %updated_source_32 = pto.vlds_post %updated_source_28[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_33, %updated_source_34 = pto.vlds_post %updated_source_30[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %22 = pto.vdiv %result_31, %result_33, %21 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %23 = pto.vsts_post %22, %20[%c64_20], %21 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %24 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_35, %updated_source_36 = pto.vlds_post %updated_source_32[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_37, %updated_source_38 = pto.vlds_post %updated_source_34[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %25 = pto.vdiv %result_35, %result_37, %24 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %26 = pto.vsts_post %25, %23[%c64_20], %24 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %27 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_39, %updated_source_40 = pto.vlds_post %updated_source_36[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_41, %updated_source_42 = pto.vlds_post %updated_source_38[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %28 = pto.vdiv %result_39, %result_41, %27 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %29 = pto.vsts_post %28, %26[%c64_20], %27 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %30 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_43, %updated_source_44 = pto.vlds_post %updated_source_40[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_45, %updated_source_46 = pto.vlds_post %updated_source_42[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %31 = pto.vdiv %result_43, %result_45, %30 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %32 = pto.vsts_post %31, %29[%c64_20], %30 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %33 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_47, %updated_source_48 = pto.vlds_post %updated_source_44[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_49, %updated_source_50 = pto.vlds_post %updated_source_46[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %34 = pto.vdiv %result_47, %result_49, %33 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %35 = pto.vsts_post %34, %32[%c64_20], %33 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + pto.tstore ins(%8 : !pto.tile_buf) outs(%5 : !pto.partition_tensor_view<8x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/exp.pto b/ptodsl/lib/a5/generated/tile_ops/exp.pto new file mode 100644 index 00000000..0f47cc9a --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/exp.pto @@ -0,0 +1,75 @@ +module { + func.func @tile_op_exp(%arg0: !pto.ptr, %arg1: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c64_1 = arith.constant 64 : index + %c1_2 = arith.constant 1 : index + %c8_3 = arith.constant 8 : index + %c64_4 = arith.constant 64 : index + %1 = pto.make_tensor_view %arg1, shape = [%c8_3, %c64_4], strides = [%c64_1, %c1_2] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_5 = arith.constant 0 : index + %c8_6 = arith.constant 8 : index + %c64_7 = arith.constant 64 : index + %2 = pto.partition_view %0, offsets = [%c0, %c0_5], sizes = [%c8_6, %c64_7] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_8 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c8_10 = arith.constant 8 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %1, offsets = [%c0_8, %c0_9], sizes = [%c8_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %4 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %5 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + pto.tload ins(%2 : !pto.partition_tensor_view<8x64xf32>) outs(%4 : !pto.tile_buf) + %6 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %7 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %8 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c0_12 = arith.constant 0 : index + %9 = pto.vlds %6[%c0_12] : !pto.ptr -> !pto.vreg<64xf32> + %10 = pto.vexp %9, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %10, %7[%c0_12], %8 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %11 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c64_13 = arith.constant 64 : index + %12 = pto.vlds %6[%c64_13] : !pto.ptr -> !pto.vreg<64xf32> + %13 = pto.vexp %12, %11 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %13, %7[%c64_13], %11 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %14 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c128 = arith.constant 128 : index + %15 = pto.vlds %6[%c128] : !pto.ptr -> !pto.vreg<64xf32> + %16 = pto.vexp %15, %14 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %16, %7[%c128], %14 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %17 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c192 = arith.constant 192 : index + %18 = pto.vlds %6[%c192] : !pto.ptr -> !pto.vreg<64xf32> + %19 = pto.vexp %18, %17 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %19, %7[%c192], %17 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %20 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c256 = arith.constant 256 : index + %21 = pto.vlds %6[%c256] : !pto.ptr -> !pto.vreg<64xf32> + %22 = pto.vexp %21, %20 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %22, %7[%c256], %20 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %23 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c320 = arith.constant 320 : index + %24 = pto.vlds %6[%c320] : !pto.ptr -> !pto.vreg<64xf32> + %25 = pto.vexp %24, %23 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %25, %7[%c320], %23 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %26 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c384 = arith.constant 384 : index + %27 = pto.vlds %6[%c384] : !pto.ptr -> !pto.vreg<64xf32> + %28 = pto.vexp %27, %26 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %28, %7[%c384], %26 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %29 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c448 = arith.constant 448 : index + %30 = pto.vlds %6[%c448] : !pto.ptr -> !pto.vreg<64xf32> + %31 = pto.vexp %30, %29 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %31, %7[%c448], %29 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + pto.tstore ins(%5 : !pto.tile_buf) outs(%3 : !pto.partition_tensor_view<8x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/gather.pto b/ptodsl/lib/a5/generated/tile_ops/gather.pto new file mode 100644 index 00000000..01045b3a --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/gather.pto @@ -0,0 +1,55 @@ +module { + func.func @tile_op_gather(%arg0: !pto.ptr, %arg1: !pto.ptr, %arg2: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c1_0 = arith.constant 1 : index + %c64_1 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c1_0, %c64_1], strides = [%c64, %c1] : !pto.tensor_view + %c64_2 = arith.constant 64 : index + %c1_3 = arith.constant 1 : index + %c1_4 = arith.constant 1 : index + %c64_5 = arith.constant 64 : index + %1 = pto.make_tensor_view %arg1, shape = [%c1_4, %c64_5], strides = [%c64_2, %c1_3] : !pto.tensor_view + %c64_6 = arith.constant 64 : index + %c1_7 = arith.constant 1 : index + %c1_8 = arith.constant 1 : index + %c64_9 = arith.constant 64 : index + %2 = pto.make_tensor_view %arg2, shape = [%c1_8, %c64_9], strides = [%c64_6, %c1_7] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_10 = arith.constant 0 : index + %c1_11 = arith.constant 1 : index + %c64_12 = arith.constant 64 : index + %3 = pto.partition_view %0, offsets = [%c0, %c0_10], sizes = [%c1_11, %c64_12] : !pto.tensor_view -> !pto.partition_tensor_view<1x64xf32> + %c0_13 = arith.constant 0 : index + %c0_14 = arith.constant 0 : index + %c1_15 = arith.constant 1 : index + %c64_16 = arith.constant 64 : index + %4 = pto.partition_view %1, offsets = [%c0_13, %c0_14], sizes = [%c1_15, %c64_16] : !pto.tensor_view -> !pto.partition_tensor_view<1x64xui32> + %c0_17 = arith.constant 0 : index + %c0_18 = arith.constant 0 : index + %c1_19 = arith.constant 1 : index + %c64_20 = arith.constant 64 : index + %5 = pto.partition_view %2, offsets = [%c0_17, %c0_18], sizes = [%c1_19, %c64_20] : !pto.tensor_view -> !pto.partition_tensor_view<1x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c256_i64 = arith.constant 256 : i64 + %c512_i64 = arith.constant 512 : i64 + %6 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %7 = pto.alloc_tile addr = %c256_i64 : !pto.tile_buf + %8 = pto.alloc_tile addr = %c512_i64 : !pto.tile_buf + pto.tload ins(%3 : !pto.partition_tensor_view<1x64xf32>) outs(%6 : !pto.tile_buf) + pto.tload ins(%4 : !pto.partition_tensor_view<1x64xui32>) outs(%7 : !pto.tile_buf) + %9 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %10 = pto.castptr %c256_i64 : i64 -> !pto.ptr + %11 = pto.castptr %c512_i64 : i64 -> !pto.ptr + %c0_21 = arith.constant 0 : index + %12 = pto.pset_b32 "PAT_ALL" : !pto.mask + %13 = pto.vlds %10[%c0_21] : !pto.ptr -> !pto.vreg<64xui32> + %c64_22 = arith.constant 64 : index + %14 = pto.vgather2 %9, %13, %c64_22 : !pto.ptr, !pto.vreg<64xui32>, index -> !pto.vreg<64xf32> + pto.vsts %14, %11[%c0_21], %12 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + pto.tstore ins(%8 : !pto.tile_buf) outs(%5 : !pto.partition_tensor_view<1x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/log.pto b/ptodsl/lib/a5/generated/tile_ops/log.pto new file mode 100644 index 00000000..5b4f62bd --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/log.pto @@ -0,0 +1,75 @@ +module { + func.func @tile_op_log(%arg0: !pto.ptr, %arg1: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c64_1 = arith.constant 64 : index + %c1_2 = arith.constant 1 : index + %c8_3 = arith.constant 8 : index + %c64_4 = arith.constant 64 : index + %1 = pto.make_tensor_view %arg1, shape = [%c8_3, %c64_4], strides = [%c64_1, %c1_2] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_5 = arith.constant 0 : index + %c8_6 = arith.constant 8 : index + %c64_7 = arith.constant 64 : index + %2 = pto.partition_view %0, offsets = [%c0, %c0_5], sizes = [%c8_6, %c64_7] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_8 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c8_10 = arith.constant 8 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %1, offsets = [%c0_8, %c0_9], sizes = [%c8_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %4 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %5 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + pto.tload ins(%2 : !pto.partition_tensor_view<8x64xf32>) outs(%4 : !pto.tile_buf) + %6 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %7 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %8 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c0_12 = arith.constant 0 : index + %9 = pto.vlds %6[%c0_12] : !pto.ptr -> !pto.vreg<64xf32> + %10 = pto.vln %9, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %10, %7[%c0_12], %8 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %11 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c64_13 = arith.constant 64 : index + %12 = pto.vlds %6[%c64_13] : !pto.ptr -> !pto.vreg<64xf32> + %13 = pto.vln %12, %11 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %13, %7[%c64_13], %11 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %14 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c128 = arith.constant 128 : index + %15 = pto.vlds %6[%c128] : !pto.ptr -> !pto.vreg<64xf32> + %16 = pto.vln %15, %14 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %16, %7[%c128], %14 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %17 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c192 = arith.constant 192 : index + %18 = pto.vlds %6[%c192] : !pto.ptr -> !pto.vreg<64xf32> + %19 = pto.vln %18, %17 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %19, %7[%c192], %17 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %20 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c256 = arith.constant 256 : index + %21 = pto.vlds %6[%c256] : !pto.ptr -> !pto.vreg<64xf32> + %22 = pto.vln %21, %20 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %22, %7[%c256], %20 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %23 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c320 = arith.constant 320 : index + %24 = pto.vlds %6[%c320] : !pto.ptr -> !pto.vreg<64xf32> + %25 = pto.vln %24, %23 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %25, %7[%c320], %23 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %26 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c384 = arith.constant 384 : index + %27 = pto.vlds %6[%c384] : !pto.ptr -> !pto.vreg<64xf32> + %28 = pto.vln %27, %26 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %28, %7[%c384], %26 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %29 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c448 = arith.constant 448 : index + %30 = pto.vlds %6[%c448] : !pto.ptr -> !pto.vreg<64xf32> + %31 = pto.vln %30, %29 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %31, %7[%c448], %29 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + pto.tstore ins(%5 : !pto.tile_buf) outs(%3 : !pto.partition_tensor_view<8x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/mov.pto b/ptodsl/lib/a5/generated/tile_ops/mov.pto new file mode 100644 index 00000000..ade28826 --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/mov.pto @@ -0,0 +1,67 @@ +module { + func.func @tile_op_mov(%arg0: !pto.ptr, %arg1: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c64_1 = arith.constant 64 : index + %c1_2 = arith.constant 1 : index + %c8_3 = arith.constant 8 : index + %c64_4 = arith.constant 64 : index + %1 = pto.make_tensor_view %arg1, shape = [%c8_3, %c64_4], strides = [%c64_1, %c1_2] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_5 = arith.constant 0 : index + %c8_6 = arith.constant 8 : index + %c64_7 = arith.constant 64 : index + %2 = pto.partition_view %0, offsets = [%c0, %c0_5], sizes = [%c8_6, %c64_7] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_8 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c8_10 = arith.constant 8 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %1, offsets = [%c0_8, %c0_9], sizes = [%c8_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %4 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %5 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + pto.tload ins(%2 : !pto.partition_tensor_view<8x64xf32>) outs(%4 : !pto.tile_buf) + %6 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %7 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %8 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c0_12 = arith.constant 0 : index + %9 = pto.vlds %6[%c0_12] : !pto.ptr -> !pto.vreg<64xf32> + pto.vsts %9, %7[%c0_12], %8 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %10 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c64_13 = arith.constant 64 : index + %11 = pto.vlds %6[%c64_13] : !pto.ptr -> !pto.vreg<64xf32> + pto.vsts %11, %7[%c64_13], %10 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %12 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c128 = arith.constant 128 : index + %13 = pto.vlds %6[%c128] : !pto.ptr -> !pto.vreg<64xf32> + pto.vsts %13, %7[%c128], %12 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %14 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c192 = arith.constant 192 : index + %15 = pto.vlds %6[%c192] : !pto.ptr -> !pto.vreg<64xf32> + pto.vsts %15, %7[%c192], %14 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %16 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c256 = arith.constant 256 : index + %17 = pto.vlds %6[%c256] : !pto.ptr -> !pto.vreg<64xf32> + pto.vsts %17, %7[%c256], %16 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %18 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c320 = arith.constant 320 : index + %19 = pto.vlds %6[%c320] : !pto.ptr -> !pto.vreg<64xf32> + pto.vsts %19, %7[%c320], %18 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %20 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c384 = arith.constant 384 : index + %21 = pto.vlds %6[%c384] : !pto.ptr -> !pto.vreg<64xf32> + pto.vsts %21, %7[%c384], %20 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %22 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c448 = arith.constant 448 : index + %23 = pto.vlds %6[%c448] : !pto.ptr -> !pto.vreg<64xf32> + pto.vsts %23, %7[%c448], %22 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + pto.tstore ins(%5 : !pto.tile_buf) outs(%3 : !pto.partition_tensor_view<8x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/mrgsort.pto b/ptodsl/lib/a5/generated/tile_ops/mrgsort.pto new file mode 100644 index 00000000..0f0afe8e --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/mrgsort.pto @@ -0,0 +1,44 @@ +module { + func.func @tile_op_mrgsort(%arg0: !pto.ptr, %arg1: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c1_0 = arith.constant 1 : index + %c64_1 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c1_0, %c64_1], strides = [%c64, %c1] : !pto.tensor_view + %c64_2 = arith.constant 64 : index + %c1_3 = arith.constant 1 : index + %c1_4 = arith.constant 1 : index + %c64_5 = arith.constant 64 : index + %1 = pto.make_tensor_view %arg1, shape = [%c1_4, %c64_5], strides = [%c64_2, %c1_3] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_6 = arith.constant 0 : index + %c1_7 = arith.constant 1 : index + %c64_8 = arith.constant 64 : index + %2 = pto.partition_view %0, offsets = [%c0, %c0_6], sizes = [%c1_7, %c64_8] : !pto.tensor_view -> !pto.partition_tensor_view<1x64xf32> + %c0_9 = arith.constant 0 : index + %c0_10 = arith.constant 0 : index + %c1_11 = arith.constant 1 : index + %c64_12 = arith.constant 64 : index + %3 = pto.partition_view %1, offsets = [%c0_9, %c0_10], sizes = [%c1_11, %c64_12] : !pto.tensor_view -> !pto.partition_tensor_view<1x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c256_i64 = arith.constant 256 : i64 + %4 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %5 = pto.alloc_tile addr = %c256_i64 : !pto.tile_buf + pto.tload ins(%2 : !pto.partition_tensor_view<1x64xf32>) outs(%4 : !pto.tile_buf) + %6 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %7 = pto.castptr %c256_i64 : i64 -> !pto.ptr + %c16 = arith.constant 16 : index + %8 = pto.addptr %6, %c16 : -> + %c32 = arith.constant 32 : index + %9 = pto.addptr %6, %c32 : -> + %c48 = arith.constant 48 : index + %10 = pto.addptr %6, %c48 : -> + %c2251834173947912_i64 = arith.constant 2251834173947912 : i64 + %c3841_i64 = arith.constant 3841 : i64 + pto.vmrgsort4 %7, %6, %8, %9, %10, %c2251834173947912_i64, %c3841_i64 : !pto.ptr, !pto.ptr, !pto.ptr, !pto.ptr, !pto.ptr, i64, i64 + pto.tstore ins(%5 : !pto.tile_buf) outs(%3 : !pto.partition_tensor_view<1x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/mul.pto b/ptodsl/lib/a5/generated/tile_ops/mul.pto new file mode 100644 index 00000000..86af874f --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/mul.pto @@ -0,0 +1,90 @@ +module { + func.func @tile_op_mul(%arg0: !pto.ptr, %arg1: !pto.ptr, %arg2: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c64_1 = arith.constant 64 : index + %c1_2 = arith.constant 1 : index + %c8_3 = arith.constant 8 : index + %c64_4 = arith.constant 64 : index + %1 = pto.make_tensor_view %arg1, shape = [%c8_3, %c64_4], strides = [%c64_1, %c1_2] : !pto.tensor_view + %c64_5 = arith.constant 64 : index + %c1_6 = arith.constant 1 : index + %c8_7 = arith.constant 8 : index + %c64_8 = arith.constant 64 : index + %2 = pto.make_tensor_view %arg2, shape = [%c8_7, %c64_8], strides = [%c64_5, %c1_6] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c8_10 = arith.constant 8 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %0, offsets = [%c0, %c0_9], sizes = [%c8_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_12 = arith.constant 0 : index + %c0_13 = arith.constant 0 : index + %c8_14 = arith.constant 8 : index + %c64_15 = arith.constant 64 : index + %4 = pto.partition_view %1, offsets = [%c0_12, %c0_13], sizes = [%c8_14, %c64_15] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_16 = arith.constant 0 : index + %c0_17 = arith.constant 0 : index + %c8_18 = arith.constant 8 : index + %c64_19 = arith.constant 64 : index + %5 = pto.partition_view %2, offsets = [%c0_16, %c0_17], sizes = [%c8_18, %c64_19] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %c4096_i64 = arith.constant 4096 : i64 + %6 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %7 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + %8 = pto.alloc_tile addr = %c4096_i64 : !pto.tile_buf + pto.tload ins(%3 : !pto.partition_tensor_view<8x64xf32>) outs(%6 : !pto.tile_buf) + pto.tload ins(%4 : !pto.partition_tensor_view<8x64xf32>) outs(%7 : !pto.tile_buf) + %9 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %10 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %11 = pto.castptr %c4096_i64 : i64 -> !pto.ptr + %c64_20 = arith.constant 64 : index + %12 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result, %updated_source = pto.vlds_post %9[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_21, %updated_source_22 = pto.vlds_post %10[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %13 = pto.vmul %result, %result_21, %12 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %14 = pto.vsts_post %13, %11[%c64_20], %12 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %15 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_23, %updated_source_24 = pto.vlds_post %updated_source[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_25, %updated_source_26 = pto.vlds_post %updated_source_22[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %16 = pto.vmul %result_23, %result_25, %15 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %17 = pto.vsts_post %16, %14[%c64_20], %15 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %18 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_27, %updated_source_28 = pto.vlds_post %updated_source_24[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_29, %updated_source_30 = pto.vlds_post %updated_source_26[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %19 = pto.vmul %result_27, %result_29, %18 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %20 = pto.vsts_post %19, %17[%c64_20], %18 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %21 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_31, %updated_source_32 = pto.vlds_post %updated_source_28[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_33, %updated_source_34 = pto.vlds_post %updated_source_30[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %22 = pto.vmul %result_31, %result_33, %21 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %23 = pto.vsts_post %22, %20[%c64_20], %21 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %24 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_35, %updated_source_36 = pto.vlds_post %updated_source_32[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_37, %updated_source_38 = pto.vlds_post %updated_source_34[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %25 = pto.vmul %result_35, %result_37, %24 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %26 = pto.vsts_post %25, %23[%c64_20], %24 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %27 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_39, %updated_source_40 = pto.vlds_post %updated_source_36[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_41, %updated_source_42 = pto.vlds_post %updated_source_38[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %28 = pto.vmul %result_39, %result_41, %27 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %29 = pto.vsts_post %28, %26[%c64_20], %27 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %30 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_43, %updated_source_44 = pto.vlds_post %updated_source_40[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_45, %updated_source_46 = pto.vlds_post %updated_source_42[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %31 = pto.vmul %result_43, %result_45, %30 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %32 = pto.vsts_post %31, %29[%c64_20], %30 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %33 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_47, %updated_source_48 = pto.vlds_post %updated_source_44[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_49, %updated_source_50 = pto.vlds_post %updated_source_46[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %34 = pto.vmul %result_47, %result_49, %33 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %35 = pto.vsts_post %34, %32[%c64_20], %33 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + pto.tstore ins(%8 : !pto.tile_buf) outs(%5 : !pto.partition_tensor_view<8x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/or_.pto b/ptodsl/lib/a5/generated/tile_ops/or_.pto new file mode 100644 index 00000000..66add0fe --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/or_.pto @@ -0,0 +1,90 @@ +module { + func.func @tile_op_or_(%arg0: !pto.ptr, %arg1: !pto.ptr, %arg2: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c64_1 = arith.constant 64 : index + %c1_2 = arith.constant 1 : index + %c8_3 = arith.constant 8 : index + %c64_4 = arith.constant 64 : index + %1 = pto.make_tensor_view %arg1, shape = [%c8_3, %c64_4], strides = [%c64_1, %c1_2] : !pto.tensor_view + %c64_5 = arith.constant 64 : index + %c1_6 = arith.constant 1 : index + %c8_7 = arith.constant 8 : index + %c64_8 = arith.constant 64 : index + %2 = pto.make_tensor_view %arg2, shape = [%c8_7, %c64_8], strides = [%c64_5, %c1_6] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c8_10 = arith.constant 8 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %0, offsets = [%c0, %c0_9], sizes = [%c8_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_12 = arith.constant 0 : index + %c0_13 = arith.constant 0 : index + %c8_14 = arith.constant 8 : index + %c64_15 = arith.constant 64 : index + %4 = pto.partition_view %1, offsets = [%c0_12, %c0_13], sizes = [%c8_14, %c64_15] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_16 = arith.constant 0 : index + %c0_17 = arith.constant 0 : index + %c8_18 = arith.constant 8 : index + %c64_19 = arith.constant 64 : index + %5 = pto.partition_view %2, offsets = [%c0_16, %c0_17], sizes = [%c8_18, %c64_19] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %c4096_i64 = arith.constant 4096 : i64 + %6 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %7 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + %8 = pto.alloc_tile addr = %c4096_i64 : !pto.tile_buf + pto.tload ins(%3 : !pto.partition_tensor_view<8x64xf32>) outs(%6 : !pto.tile_buf) + pto.tload ins(%4 : !pto.partition_tensor_view<8x64xf32>) outs(%7 : !pto.tile_buf) + %9 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %10 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %11 = pto.castptr %c4096_i64 : i64 -> !pto.ptr + %c64_20 = arith.constant 64 : index + %12 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result, %updated_source = pto.vlds_post %9[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_21, %updated_source_22 = pto.vlds_post %10[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %13 = pto.vor %result, %result_21, %12 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %14 = pto.vsts_post %13, %11[%c64_20], %12 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %15 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_23, %updated_source_24 = pto.vlds_post %updated_source[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_25, %updated_source_26 = pto.vlds_post %updated_source_22[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %16 = pto.vor %result_23, %result_25, %15 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %17 = pto.vsts_post %16, %14[%c64_20], %15 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %18 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_27, %updated_source_28 = pto.vlds_post %updated_source_24[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_29, %updated_source_30 = pto.vlds_post %updated_source_26[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %19 = pto.vor %result_27, %result_29, %18 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %20 = pto.vsts_post %19, %17[%c64_20], %18 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %21 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_31, %updated_source_32 = pto.vlds_post %updated_source_28[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_33, %updated_source_34 = pto.vlds_post %updated_source_30[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %22 = pto.vor %result_31, %result_33, %21 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %23 = pto.vsts_post %22, %20[%c64_20], %21 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %24 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_35, %updated_source_36 = pto.vlds_post %updated_source_32[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_37, %updated_source_38 = pto.vlds_post %updated_source_34[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %25 = pto.vor %result_35, %result_37, %24 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %26 = pto.vsts_post %25, %23[%c64_20], %24 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %27 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_39, %updated_source_40 = pto.vlds_post %updated_source_36[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_41, %updated_source_42 = pto.vlds_post %updated_source_38[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %28 = pto.vor %result_39, %result_41, %27 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %29 = pto.vsts_post %28, %26[%c64_20], %27 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %30 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_43, %updated_source_44 = pto.vlds_post %updated_source_40[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_45, %updated_source_46 = pto.vlds_post %updated_source_42[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %31 = pto.vor %result_43, %result_45, %30 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %32 = pto.vsts_post %31, %29[%c64_20], %30 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %33 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_47, %updated_source_48 = pto.vlds_post %updated_source_44[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_49, %updated_source_50 = pto.vlds_post %updated_source_46[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %34 = pto.vor %result_47, %result_49, %33 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %35 = pto.vsts_post %34, %32[%c64_20], %33 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + pto.tstore ins(%8 : !pto.tile_buf) outs(%5 : !pto.partition_tensor_view<8x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/reciprocal.pto b/ptodsl/lib/a5/generated/tile_ops/reciprocal.pto new file mode 100644 index 00000000..3ed78ddc --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/reciprocal.pto @@ -0,0 +1,75 @@ +module { + func.func @tile_op_reciprocal(%arg0: !pto.ptr, %arg1: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c64_1 = arith.constant 64 : index + %c1_2 = arith.constant 1 : index + %c8_3 = arith.constant 8 : index + %c64_4 = arith.constant 64 : index + %1 = pto.make_tensor_view %arg1, shape = [%c8_3, %c64_4], strides = [%c64_1, %c1_2] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_5 = arith.constant 0 : index + %c8_6 = arith.constant 8 : index + %c64_7 = arith.constant 64 : index + %2 = pto.partition_view %0, offsets = [%c0, %c0_5], sizes = [%c8_6, %c64_7] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_8 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c8_10 = arith.constant 8 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %1, offsets = [%c0_8, %c0_9], sizes = [%c8_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %4 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %5 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + pto.tload ins(%2 : !pto.partition_tensor_view<8x64xf32>) outs(%4 : !pto.tile_buf) + %6 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %7 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %8 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c0_12 = arith.constant 0 : index + %9 = pto.vlds %6[%c0_12] : !pto.ptr -> !pto.vreg<64xf32> + %10 = pto.vrec %9, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %10, %7[%c0_12], %8 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %11 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c64_13 = arith.constant 64 : index + %12 = pto.vlds %6[%c64_13] : !pto.ptr -> !pto.vreg<64xf32> + %13 = pto.vrec %12, %11 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %13, %7[%c64_13], %11 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %14 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c128 = arith.constant 128 : index + %15 = pto.vlds %6[%c128] : !pto.ptr -> !pto.vreg<64xf32> + %16 = pto.vrec %15, %14 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %16, %7[%c128], %14 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %17 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c192 = arith.constant 192 : index + %18 = pto.vlds %6[%c192] : !pto.ptr -> !pto.vreg<64xf32> + %19 = pto.vrec %18, %17 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %19, %7[%c192], %17 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %20 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c256 = arith.constant 256 : index + %21 = pto.vlds %6[%c256] : !pto.ptr -> !pto.vreg<64xf32> + %22 = pto.vrec %21, %20 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %22, %7[%c256], %20 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %23 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c320 = arith.constant 320 : index + %24 = pto.vlds %6[%c320] : !pto.ptr -> !pto.vreg<64xf32> + %25 = pto.vrec %24, %23 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %25, %7[%c320], %23 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %26 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c384 = arith.constant 384 : index + %27 = pto.vlds %6[%c384] : !pto.ptr -> !pto.vreg<64xf32> + %28 = pto.vrec %27, %26 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %28, %7[%c384], %26 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %29 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c448 = arith.constant 448 : index + %30 = pto.vlds %6[%c448] : !pto.ptr -> !pto.vreg<64xf32> + %31 = pto.vrec %30, %29 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %31, %7[%c448], %29 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + pto.tstore ins(%5 : !pto.tile_buf) outs(%3 : !pto.partition_tensor_view<8x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/relu.pto b/ptodsl/lib/a5/generated/tile_ops/relu.pto new file mode 100644 index 00000000..001aa127 --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/relu.pto @@ -0,0 +1,75 @@ +module { + func.func @tile_op_relu(%arg0: !pto.ptr, %arg1: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c64_1 = arith.constant 64 : index + %c1_2 = arith.constant 1 : index + %c8_3 = arith.constant 8 : index + %c64_4 = arith.constant 64 : index + %1 = pto.make_tensor_view %arg1, shape = [%c8_3, %c64_4], strides = [%c64_1, %c1_2] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_5 = arith.constant 0 : index + %c8_6 = arith.constant 8 : index + %c64_7 = arith.constant 64 : index + %2 = pto.partition_view %0, offsets = [%c0, %c0_5], sizes = [%c8_6, %c64_7] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_8 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c8_10 = arith.constant 8 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %1, offsets = [%c0_8, %c0_9], sizes = [%c8_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %4 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %5 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + pto.tload ins(%2 : !pto.partition_tensor_view<8x64xf32>) outs(%4 : !pto.tile_buf) + %6 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %7 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %8 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c0_12 = arith.constant 0 : index + %9 = pto.vlds %6[%c0_12] : !pto.ptr -> !pto.vreg<64xf32> + %10 = pto.vrelu %9, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %10, %7[%c0_12], %8 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %11 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c64_13 = arith.constant 64 : index + %12 = pto.vlds %6[%c64_13] : !pto.ptr -> !pto.vreg<64xf32> + %13 = pto.vrelu %12, %11 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %13, %7[%c64_13], %11 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %14 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c128 = arith.constant 128 : index + %15 = pto.vlds %6[%c128] : !pto.ptr -> !pto.vreg<64xf32> + %16 = pto.vrelu %15, %14 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %16, %7[%c128], %14 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %17 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c192 = arith.constant 192 : index + %18 = pto.vlds %6[%c192] : !pto.ptr -> !pto.vreg<64xf32> + %19 = pto.vrelu %18, %17 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %19, %7[%c192], %17 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %20 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c256 = arith.constant 256 : index + %21 = pto.vlds %6[%c256] : !pto.ptr -> !pto.vreg<64xf32> + %22 = pto.vrelu %21, %20 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %22, %7[%c256], %20 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %23 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c320 = arith.constant 320 : index + %24 = pto.vlds %6[%c320] : !pto.ptr -> !pto.vreg<64xf32> + %25 = pto.vrelu %24, %23 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %25, %7[%c320], %23 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %26 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c384 = arith.constant 384 : index + %27 = pto.vlds %6[%c384] : !pto.ptr -> !pto.vreg<64xf32> + %28 = pto.vrelu %27, %26 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %28, %7[%c384], %26 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %29 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c448 = arith.constant 448 : index + %30 = pto.vlds %6[%c448] : !pto.ptr -> !pto.vreg<64xf32> + %31 = pto.vrelu %30, %29 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %31, %7[%c448], %29 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + pto.tstore ins(%5 : !pto.tile_buf) outs(%3 : !pto.partition_tensor_view<8x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/row_expand.pto b/ptodsl/lib/a5/generated/tile_ops/row_expand.pto new file mode 100644 index 00000000..e3fd1be3 --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/row_expand.pto @@ -0,0 +1,91 @@ +module { + func.func @tile_op_row_expand(%arg0: !pto.ptr, %arg1: !pto.ptr) { + %c1 = arith.constant 1 : index + %c1_0 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c1_1 = arith.constant 1 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c1_1], strides = [%c1, %c1_0] : !pto.tensor_view + %c64 = arith.constant 64 : index + %c1_2 = arith.constant 1 : index + %c8_3 = arith.constant 8 : index + %c64_4 = arith.constant 64 : index + %1 = pto.make_tensor_view %arg1, shape = [%c8_3, %c64_4], strides = [%c64, %c1_2] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_5 = arith.constant 0 : index + %c8_6 = arith.constant 8 : index + %c1_7 = arith.constant 1 : index + %2 = pto.partition_view %0, offsets = [%c0, %c0_5], sizes = [%c8_6, %c1_7] : !pto.tensor_view -> !pto.partition_tensor_view<8x1xf32> + %c0_8 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c8_10 = arith.constant 8 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %1, offsets = [%c0_8, %c0_9], sizes = [%c8_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %4 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %5 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + pto.tload ins(%2 : !pto.partition_tensor_view<8x1xf32>) outs(%4 : !pto.tile_buf) + %6 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %7 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %c0_12 = arith.constant 0 : index + %8 = pto.vldas %6[%c0_12] : !pto.ptr -> !pto.align + %result, %updated_align, %updated_source = pto.vldus %6[%c0_12], %8 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %9 = pto.vdup %result {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %10 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c0_13 = arith.constant 0 : index + pto.vsts %9, %7[%c0_13], %10 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c64_14 = arith.constant 64 : index + %11 = pto.vldas %6[%c64_14] : !pto.ptr -> !pto.align + %result_15, %updated_align_16, %updated_source_17 = pto.vldus %6[%c64_14], %11 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %12 = pto.vdup %result_15 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %13 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c64_18 = arith.constant 64 : index + pto.vsts %12, %7[%c64_18], %13 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c128 = arith.constant 128 : index + %14 = pto.vldas %6[%c128] : !pto.ptr -> !pto.align + %result_19, %updated_align_20, %updated_source_21 = pto.vldus %6[%c128], %14 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %15 = pto.vdup %result_19 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %16 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c128_22 = arith.constant 128 : index + pto.vsts %15, %7[%c128_22], %16 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c192 = arith.constant 192 : index + %17 = pto.vldas %6[%c192] : !pto.ptr -> !pto.align + %result_23, %updated_align_24, %updated_source_25 = pto.vldus %6[%c192], %17 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %18 = pto.vdup %result_23 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %19 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c192_26 = arith.constant 192 : index + pto.vsts %18, %7[%c192_26], %19 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c256 = arith.constant 256 : index + %20 = pto.vldas %6[%c256] : !pto.ptr -> !pto.align + %result_27, %updated_align_28, %updated_source_29 = pto.vldus %6[%c256], %20 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %21 = pto.vdup %result_27 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %22 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c256_30 = arith.constant 256 : index + pto.vsts %21, %7[%c256_30], %22 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c320 = arith.constant 320 : index + %23 = pto.vldas %6[%c320] : !pto.ptr -> !pto.align + %result_31, %updated_align_32, %updated_source_33 = pto.vldus %6[%c320], %23 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %24 = pto.vdup %result_31 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %25 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c320_34 = arith.constant 320 : index + pto.vsts %24, %7[%c320_34], %25 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c384 = arith.constant 384 : index + %26 = pto.vldas %6[%c384] : !pto.ptr -> !pto.align + %result_35, %updated_align_36, %updated_source_37 = pto.vldus %6[%c384], %26 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %27 = pto.vdup %result_35 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %28 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c384_38 = arith.constant 384 : index + pto.vsts %27, %7[%c384_38], %28 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c448 = arith.constant 448 : index + %29 = pto.vldas %6[%c448] : !pto.ptr -> !pto.align + %result_39, %updated_align_40, %updated_source_41 = pto.vldus %6[%c448], %29 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %30 = pto.vdup %result_39 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %31 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c448_42 = arith.constant 448 : index + pto.vsts %30, %7[%c448_42], %31 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + pto.tstore ins(%5 : !pto.tile_buf) outs(%3 : !pto.partition_tensor_view<8x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/row_expand_div.pto b/ptodsl/lib/a5/generated/tile_ops/row_expand_div.pto new file mode 100644 index 00000000..7f1574e2 --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/row_expand_div.pto @@ -0,0 +1,121 @@ +module { + func.func @tile_op_row_expand_div(%arg0: !pto.ptr, %arg1: !pto.ptr, %arg2: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c1_1 = arith.constant 1 : index + %c1_2 = arith.constant 1 : index + %c8_3 = arith.constant 8 : index + %c1_4 = arith.constant 1 : index + %1 = pto.make_tensor_view %arg1, shape = [%c8_3, %c1_4], strides = [%c1_1, %c1_2] : !pto.tensor_view + %c64_5 = arith.constant 64 : index + %c1_6 = arith.constant 1 : index + %c8_7 = arith.constant 8 : index + %c64_8 = arith.constant 64 : index + %2 = pto.make_tensor_view %arg2, shape = [%c8_7, %c64_8], strides = [%c64_5, %c1_6] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c8_10 = arith.constant 8 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %0, offsets = [%c0, %c0_9], sizes = [%c8_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_12 = arith.constant 0 : index + %c0_13 = arith.constant 0 : index + %c8_14 = arith.constant 8 : index + %c1_15 = arith.constant 1 : index + %4 = pto.partition_view %1, offsets = [%c0_12, %c0_13], sizes = [%c8_14, %c1_15] : !pto.tensor_view -> !pto.partition_tensor_view<8x1xf32> + %c0_16 = arith.constant 0 : index + %c0_17 = arith.constant 0 : index + %c8_18 = arith.constant 8 : index + %c64_19 = arith.constant 64 : index + %5 = pto.partition_view %2, offsets = [%c0_16, %c0_17], sizes = [%c8_18, %c64_19] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %c4096_i64 = arith.constant 4096 : i64 + %6 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %7 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + %8 = pto.alloc_tile addr = %c4096_i64 : !pto.tile_buf + pto.tload ins(%3 : !pto.partition_tensor_view<8x64xf32>) outs(%6 : !pto.tile_buf) + pto.tload ins(%4 : !pto.partition_tensor_view<8x1xf32>) outs(%7 : !pto.tile_buf) + %9 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %10 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %11 = pto.castptr %c4096_i64 : i64 -> !pto.ptr + %c0_20 = arith.constant 0 : index + %12 = pto.vldas %10[%c0_20] : !pto.ptr -> !pto.align + %result, %updated_align, %updated_source = pto.vldus %10[%c0_20], %12 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %13 = pto.vdup %result {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %14 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c0_21 = arith.constant 0 : index + %15 = pto.vlds %9[%c0_21] : !pto.ptr -> !pto.vreg<64xf32> + %16 = pto.vdiv %15, %13, %14 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %16, %11[%c0_21], %14 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c64_22 = arith.constant 64 : index + %17 = pto.vldas %10[%c64_22] : !pto.ptr -> !pto.align + %result_23, %updated_align_24, %updated_source_25 = pto.vldus %10[%c64_22], %17 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %18 = pto.vdup %result_23 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %19 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c64_26 = arith.constant 64 : index + %20 = pto.vlds %9[%c64_26] : !pto.ptr -> !pto.vreg<64xf32> + %21 = pto.vdiv %20, %18, %19 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %21, %11[%c64_26], %19 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c128 = arith.constant 128 : index + %22 = pto.vldas %10[%c128] : !pto.ptr -> !pto.align + %result_27, %updated_align_28, %updated_source_29 = pto.vldus %10[%c128], %22 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %23 = pto.vdup %result_27 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %24 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c128_30 = arith.constant 128 : index + %25 = pto.vlds %9[%c128_30] : !pto.ptr -> !pto.vreg<64xf32> + %26 = pto.vdiv %25, %23, %24 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %26, %11[%c128_30], %24 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c192 = arith.constant 192 : index + %27 = pto.vldas %10[%c192] : !pto.ptr -> !pto.align + %result_31, %updated_align_32, %updated_source_33 = pto.vldus %10[%c192], %27 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %28 = pto.vdup %result_31 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %29 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c192_34 = arith.constant 192 : index + %30 = pto.vlds %9[%c192_34] : !pto.ptr -> !pto.vreg<64xf32> + %31 = pto.vdiv %30, %28, %29 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %31, %11[%c192_34], %29 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c256 = arith.constant 256 : index + %32 = pto.vldas %10[%c256] : !pto.ptr -> !pto.align + %result_35, %updated_align_36, %updated_source_37 = pto.vldus %10[%c256], %32 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %33 = pto.vdup %result_35 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %34 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c256_38 = arith.constant 256 : index + %35 = pto.vlds %9[%c256_38] : !pto.ptr -> !pto.vreg<64xf32> + %36 = pto.vdiv %35, %33, %34 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %36, %11[%c256_38], %34 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c320 = arith.constant 320 : index + %37 = pto.vldas %10[%c320] : !pto.ptr -> !pto.align + %result_39, %updated_align_40, %updated_source_41 = pto.vldus %10[%c320], %37 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %38 = pto.vdup %result_39 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %39 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c320_42 = arith.constant 320 : index + %40 = pto.vlds %9[%c320_42] : !pto.ptr -> !pto.vreg<64xf32> + %41 = pto.vdiv %40, %38, %39 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %41, %11[%c320_42], %39 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c384 = arith.constant 384 : index + %42 = pto.vldas %10[%c384] : !pto.ptr -> !pto.align + %result_43, %updated_align_44, %updated_source_45 = pto.vldus %10[%c384], %42 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %43 = pto.vdup %result_43 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %44 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c384_46 = arith.constant 384 : index + %45 = pto.vlds %9[%c384_46] : !pto.ptr -> !pto.vreg<64xf32> + %46 = pto.vdiv %45, %43, %44 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %46, %11[%c384_46], %44 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c448 = arith.constant 448 : index + %47 = pto.vldas %10[%c448] : !pto.ptr -> !pto.align + %result_47, %updated_align_48, %updated_source_49 = pto.vldus %10[%c448], %47 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %48 = pto.vdup %result_47 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %49 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c448_50 = arith.constant 448 : index + %50 = pto.vlds %9[%c448_50] : !pto.ptr -> !pto.vreg<64xf32> + %51 = pto.vdiv %50, %48, %49 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %51, %11[%c448_50], %49 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + pto.tstore ins(%8 : !pto.tile_buf) outs(%5 : !pto.partition_tensor_view<8x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/row_expand_mul.pto b/ptodsl/lib/a5/generated/tile_ops/row_expand_mul.pto new file mode 100644 index 00000000..8f9631d5 --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/row_expand_mul.pto @@ -0,0 +1,121 @@ +module { + func.func @tile_op_row_expand_mul(%arg0: !pto.ptr, %arg1: !pto.ptr, %arg2: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c1_1 = arith.constant 1 : index + %c1_2 = arith.constant 1 : index + %c8_3 = arith.constant 8 : index + %c1_4 = arith.constant 1 : index + %1 = pto.make_tensor_view %arg1, shape = [%c8_3, %c1_4], strides = [%c1_1, %c1_2] : !pto.tensor_view + %c64_5 = arith.constant 64 : index + %c1_6 = arith.constant 1 : index + %c8_7 = arith.constant 8 : index + %c64_8 = arith.constant 64 : index + %2 = pto.make_tensor_view %arg2, shape = [%c8_7, %c64_8], strides = [%c64_5, %c1_6] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c8_10 = arith.constant 8 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %0, offsets = [%c0, %c0_9], sizes = [%c8_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_12 = arith.constant 0 : index + %c0_13 = arith.constant 0 : index + %c8_14 = arith.constant 8 : index + %c1_15 = arith.constant 1 : index + %4 = pto.partition_view %1, offsets = [%c0_12, %c0_13], sizes = [%c8_14, %c1_15] : !pto.tensor_view -> !pto.partition_tensor_view<8x1xf32> + %c0_16 = arith.constant 0 : index + %c0_17 = arith.constant 0 : index + %c8_18 = arith.constant 8 : index + %c64_19 = arith.constant 64 : index + %5 = pto.partition_view %2, offsets = [%c0_16, %c0_17], sizes = [%c8_18, %c64_19] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %c4096_i64 = arith.constant 4096 : i64 + %6 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %7 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + %8 = pto.alloc_tile addr = %c4096_i64 : !pto.tile_buf + pto.tload ins(%3 : !pto.partition_tensor_view<8x64xf32>) outs(%6 : !pto.tile_buf) + pto.tload ins(%4 : !pto.partition_tensor_view<8x1xf32>) outs(%7 : !pto.tile_buf) + %9 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %10 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %11 = pto.castptr %c4096_i64 : i64 -> !pto.ptr + %c0_20 = arith.constant 0 : index + %12 = pto.vldas %10[%c0_20] : !pto.ptr -> !pto.align + %result, %updated_align, %updated_source = pto.vldus %10[%c0_20], %12 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %13 = pto.vdup %result {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %14 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c0_21 = arith.constant 0 : index + %15 = pto.vlds %9[%c0_21] : !pto.ptr -> !pto.vreg<64xf32> + %16 = pto.vmul %15, %13, %14 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %16, %11[%c0_21], %14 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c64_22 = arith.constant 64 : index + %17 = pto.vldas %10[%c64_22] : !pto.ptr -> !pto.align + %result_23, %updated_align_24, %updated_source_25 = pto.vldus %10[%c64_22], %17 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %18 = pto.vdup %result_23 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %19 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c64_26 = arith.constant 64 : index + %20 = pto.vlds %9[%c64_26] : !pto.ptr -> !pto.vreg<64xf32> + %21 = pto.vmul %20, %18, %19 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %21, %11[%c64_26], %19 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c128 = arith.constant 128 : index + %22 = pto.vldas %10[%c128] : !pto.ptr -> !pto.align + %result_27, %updated_align_28, %updated_source_29 = pto.vldus %10[%c128], %22 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %23 = pto.vdup %result_27 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %24 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c128_30 = arith.constant 128 : index + %25 = pto.vlds %9[%c128_30] : !pto.ptr -> !pto.vreg<64xf32> + %26 = pto.vmul %25, %23, %24 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %26, %11[%c128_30], %24 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c192 = arith.constant 192 : index + %27 = pto.vldas %10[%c192] : !pto.ptr -> !pto.align + %result_31, %updated_align_32, %updated_source_33 = pto.vldus %10[%c192], %27 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %28 = pto.vdup %result_31 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %29 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c192_34 = arith.constant 192 : index + %30 = pto.vlds %9[%c192_34] : !pto.ptr -> !pto.vreg<64xf32> + %31 = pto.vmul %30, %28, %29 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %31, %11[%c192_34], %29 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c256 = arith.constant 256 : index + %32 = pto.vldas %10[%c256] : !pto.ptr -> !pto.align + %result_35, %updated_align_36, %updated_source_37 = pto.vldus %10[%c256], %32 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %33 = pto.vdup %result_35 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %34 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c256_38 = arith.constant 256 : index + %35 = pto.vlds %9[%c256_38] : !pto.ptr -> !pto.vreg<64xf32> + %36 = pto.vmul %35, %33, %34 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %36, %11[%c256_38], %34 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c320 = arith.constant 320 : index + %37 = pto.vldas %10[%c320] : !pto.ptr -> !pto.align + %result_39, %updated_align_40, %updated_source_41 = pto.vldus %10[%c320], %37 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %38 = pto.vdup %result_39 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %39 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c320_42 = arith.constant 320 : index + %40 = pto.vlds %9[%c320_42] : !pto.ptr -> !pto.vreg<64xf32> + %41 = pto.vmul %40, %38, %39 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %41, %11[%c320_42], %39 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c384 = arith.constant 384 : index + %42 = pto.vldas %10[%c384] : !pto.ptr -> !pto.align + %result_43, %updated_align_44, %updated_source_45 = pto.vldus %10[%c384], %42 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %43 = pto.vdup %result_43 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %44 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c384_46 = arith.constant 384 : index + %45 = pto.vlds %9[%c384_46] : !pto.ptr -> !pto.vreg<64xf32> + %46 = pto.vmul %45, %43, %44 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %46, %11[%c384_46], %44 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c448 = arith.constant 448 : index + %47 = pto.vldas %10[%c448] : !pto.ptr -> !pto.align + %result_47, %updated_align_48, %updated_source_49 = pto.vldus %10[%c448], %47 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %48 = pto.vdup %result_47 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %49 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c448_50 = arith.constant 448 : index + %50 = pto.vlds %9[%c448_50] : !pto.ptr -> !pto.vreg<64xf32> + %51 = pto.vmul %50, %48, %49 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %51, %11[%c448_50], %49 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + pto.tstore ins(%8 : !pto.tile_buf) outs(%5 : !pto.partition_tensor_view<8x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/row_expand_sub.pto b/ptodsl/lib/a5/generated/tile_ops/row_expand_sub.pto new file mode 100644 index 00000000..910a0c80 --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/row_expand_sub.pto @@ -0,0 +1,121 @@ +module { + func.func @tile_op_row_expand_sub(%arg0: !pto.ptr, %arg1: !pto.ptr, %arg2: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c1_1 = arith.constant 1 : index + %c1_2 = arith.constant 1 : index + %c8_3 = arith.constant 8 : index + %c1_4 = arith.constant 1 : index + %1 = pto.make_tensor_view %arg1, shape = [%c8_3, %c1_4], strides = [%c1_1, %c1_2] : !pto.tensor_view + %c64_5 = arith.constant 64 : index + %c1_6 = arith.constant 1 : index + %c8_7 = arith.constant 8 : index + %c64_8 = arith.constant 64 : index + %2 = pto.make_tensor_view %arg2, shape = [%c8_7, %c64_8], strides = [%c64_5, %c1_6] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c8_10 = arith.constant 8 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %0, offsets = [%c0, %c0_9], sizes = [%c8_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_12 = arith.constant 0 : index + %c0_13 = arith.constant 0 : index + %c8_14 = arith.constant 8 : index + %c1_15 = arith.constant 1 : index + %4 = pto.partition_view %1, offsets = [%c0_12, %c0_13], sizes = [%c8_14, %c1_15] : !pto.tensor_view -> !pto.partition_tensor_view<8x1xf32> + %c0_16 = arith.constant 0 : index + %c0_17 = arith.constant 0 : index + %c8_18 = arith.constant 8 : index + %c64_19 = arith.constant 64 : index + %5 = pto.partition_view %2, offsets = [%c0_16, %c0_17], sizes = [%c8_18, %c64_19] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %c4096_i64 = arith.constant 4096 : i64 + %6 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %7 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + %8 = pto.alloc_tile addr = %c4096_i64 : !pto.tile_buf + pto.tload ins(%3 : !pto.partition_tensor_view<8x64xf32>) outs(%6 : !pto.tile_buf) + pto.tload ins(%4 : !pto.partition_tensor_view<8x1xf32>) outs(%7 : !pto.tile_buf) + %9 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %10 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %11 = pto.castptr %c4096_i64 : i64 -> !pto.ptr + %c0_20 = arith.constant 0 : index + %12 = pto.vldas %10[%c0_20] : !pto.ptr -> !pto.align + %result, %updated_align, %updated_source = pto.vldus %10[%c0_20], %12 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %13 = pto.vdup %result {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %14 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c0_21 = arith.constant 0 : index + %15 = pto.vlds %9[%c0_21] : !pto.ptr -> !pto.vreg<64xf32> + %16 = pto.vsub %15, %13, %14 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %16, %11[%c0_21], %14 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c64_22 = arith.constant 64 : index + %17 = pto.vldas %10[%c64_22] : !pto.ptr -> !pto.align + %result_23, %updated_align_24, %updated_source_25 = pto.vldus %10[%c64_22], %17 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %18 = pto.vdup %result_23 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %19 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c64_26 = arith.constant 64 : index + %20 = pto.vlds %9[%c64_26] : !pto.ptr -> !pto.vreg<64xf32> + %21 = pto.vsub %20, %18, %19 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %21, %11[%c64_26], %19 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c128 = arith.constant 128 : index + %22 = pto.vldas %10[%c128] : !pto.ptr -> !pto.align + %result_27, %updated_align_28, %updated_source_29 = pto.vldus %10[%c128], %22 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %23 = pto.vdup %result_27 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %24 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c128_30 = arith.constant 128 : index + %25 = pto.vlds %9[%c128_30] : !pto.ptr -> !pto.vreg<64xf32> + %26 = pto.vsub %25, %23, %24 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %26, %11[%c128_30], %24 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c192 = arith.constant 192 : index + %27 = pto.vldas %10[%c192] : !pto.ptr -> !pto.align + %result_31, %updated_align_32, %updated_source_33 = pto.vldus %10[%c192], %27 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %28 = pto.vdup %result_31 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %29 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c192_34 = arith.constant 192 : index + %30 = pto.vlds %9[%c192_34] : !pto.ptr -> !pto.vreg<64xf32> + %31 = pto.vsub %30, %28, %29 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %31, %11[%c192_34], %29 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c256 = arith.constant 256 : index + %32 = pto.vldas %10[%c256] : !pto.ptr -> !pto.align + %result_35, %updated_align_36, %updated_source_37 = pto.vldus %10[%c256], %32 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %33 = pto.vdup %result_35 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %34 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c256_38 = arith.constant 256 : index + %35 = pto.vlds %9[%c256_38] : !pto.ptr -> !pto.vreg<64xf32> + %36 = pto.vsub %35, %33, %34 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %36, %11[%c256_38], %34 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c320 = arith.constant 320 : index + %37 = pto.vldas %10[%c320] : !pto.ptr -> !pto.align + %result_39, %updated_align_40, %updated_source_41 = pto.vldus %10[%c320], %37 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %38 = pto.vdup %result_39 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %39 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c320_42 = arith.constant 320 : index + %40 = pto.vlds %9[%c320_42] : !pto.ptr -> !pto.vreg<64xf32> + %41 = pto.vsub %40, %38, %39 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %41, %11[%c320_42], %39 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c384 = arith.constant 384 : index + %42 = pto.vldas %10[%c384] : !pto.ptr -> !pto.align + %result_43, %updated_align_44, %updated_source_45 = pto.vldus %10[%c384], %42 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %43 = pto.vdup %result_43 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %44 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c384_46 = arith.constant 384 : index + %45 = pto.vlds %9[%c384_46] : !pto.ptr -> !pto.vreg<64xf32> + %46 = pto.vsub %45, %43, %44 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %46, %11[%c384_46], %44 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %c448 = arith.constant 448 : index + %47 = pto.vldas %10[%c448] : !pto.ptr -> !pto.align + %result_47, %updated_align_48, %updated_source_49 = pto.vldus %10[%c448], %47 : !pto.ptr, !pto.align -> !pto.vreg<64xf32>, !pto.align, !pto.ptr + %48 = pto.vdup %result_47 {position = "POS_LOWEST"} : !pto.vreg<64xf32> -> !pto.vreg<64xf32> + %49 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c448_50 = arith.constant 448 : index + %50 = pto.vlds %9[%c448_50] : !pto.ptr -> !pto.vreg<64xf32> + %51 = pto.vsub %50, %48, %49 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %51, %11[%c448_50], %49 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + pto.tstore ins(%8 : !pto.tile_buf) outs(%5 : !pto.partition_tensor_view<8x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/row_max.pto b/ptodsl/lib/a5/generated/tile_ops/row_max.pto new file mode 100644 index 00000000..8bc4bafd --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/row_max.pto @@ -0,0 +1,118 @@ +module { + func.func @tile_op_row_max(%arg0: !pto.ptr, %arg1: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c1_1 = arith.constant 1 : index + %c1_2 = arith.constant 1 : index + %c8_3 = arith.constant 8 : index + %c1_4 = arith.constant 1 : index + %1 = pto.make_tensor_view %arg1, shape = [%c8_3, %c1_4], strides = [%c1_1, %c1_2] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_5 = arith.constant 0 : index + %c8_6 = arith.constant 8 : index + %c64_7 = arith.constant 64 : index + %2 = pto.partition_view %0, offsets = [%c0, %c0_5], sizes = [%c8_6, %c64_7] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_8 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c8_10 = arith.constant 8 : index + %c1_11 = arith.constant 1 : index + %3 = pto.partition_view %1, offsets = [%c0_8, %c0_9], sizes = [%c8_10, %c1_11] : !pto.tensor_view -> !pto.partition_tensor_view<8x1xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %4 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %5 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + pto.tload ins(%2 : !pto.partition_tensor_view<8x64xf32>) outs(%4 : !pto.tile_buf) + %6 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %7 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %8 = pto.pset_b32 "PAT_ALL" : !pto.mask + %cst = arith.constant 0xFF800000 : f32 + %9 = pto.vbr %cst : f32 -> !pto.vreg<64xf32> + %10 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c0_12 = arith.constant 0 : index + %11 = pto.vlds %6[%c0_12] : !pto.ptr -> !pto.vreg<64xf32> + %12 = pto.vsel %11, %9, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %13 = pto.vcmax %12, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %14 = pto.vmax %9, %13, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c0_13 = arith.constant 0 : index + %c1_i32 = arith.constant 1 : i32 + %mask, %scalar_out = pto.plt_b32 %c1_i32 : i32 -> !pto.mask, i32 + pto.vsts %14, %7[%c0_13], %mask {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %15 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c64_14 = arith.constant 64 : index + %16 = pto.vlds %6[%c64_14] : !pto.ptr -> !pto.vreg<64xf32> + %17 = pto.vsel %16, %9, %15 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %18 = pto.vcmax %17, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %19 = pto.vmax %9, %18, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_15 = arith.constant 64 : index + %c1_i32_16 = arith.constant 1 : i32 + %mask_17, %scalar_out_18 = pto.plt_b32 %c1_i32_16 : i32 -> !pto.mask, i32 + pto.vsts %19, %7[%c64_15], %mask_17 {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %20 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c128 = arith.constant 128 : index + %21 = pto.vlds %6[%c128] : !pto.ptr -> !pto.vreg<64xf32> + %22 = pto.vsel %21, %9, %20 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %23 = pto.vcmax %22, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %24 = pto.vmax %9, %23, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c128_19 = arith.constant 128 : index + %c1_i32_20 = arith.constant 1 : i32 + %mask_21, %scalar_out_22 = pto.plt_b32 %c1_i32_20 : i32 -> !pto.mask, i32 + pto.vsts %24, %7[%c128_19], %mask_21 {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %25 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c192 = arith.constant 192 : index + %26 = pto.vlds %6[%c192] : !pto.ptr -> !pto.vreg<64xf32> + %27 = pto.vsel %26, %9, %25 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %28 = pto.vcmax %27, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %29 = pto.vmax %9, %28, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c192_23 = arith.constant 192 : index + %c1_i32_24 = arith.constant 1 : i32 + %mask_25, %scalar_out_26 = pto.plt_b32 %c1_i32_24 : i32 -> !pto.mask, i32 + pto.vsts %29, %7[%c192_23], %mask_25 {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %30 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c256 = arith.constant 256 : index + %31 = pto.vlds %6[%c256] : !pto.ptr -> !pto.vreg<64xf32> + %32 = pto.vsel %31, %9, %30 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %33 = pto.vcmax %32, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %34 = pto.vmax %9, %33, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c256_27 = arith.constant 256 : index + %c1_i32_28 = arith.constant 1 : i32 + %mask_29, %scalar_out_30 = pto.plt_b32 %c1_i32_28 : i32 -> !pto.mask, i32 + pto.vsts %34, %7[%c256_27], %mask_29 {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %35 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c320 = arith.constant 320 : index + %36 = pto.vlds %6[%c320] : !pto.ptr -> !pto.vreg<64xf32> + %37 = pto.vsel %36, %9, %35 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %38 = pto.vcmax %37, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %39 = pto.vmax %9, %38, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c320_31 = arith.constant 320 : index + %c1_i32_32 = arith.constant 1 : i32 + %mask_33, %scalar_out_34 = pto.plt_b32 %c1_i32_32 : i32 -> !pto.mask, i32 + pto.vsts %39, %7[%c320_31], %mask_33 {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %40 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c384 = arith.constant 384 : index + %41 = pto.vlds %6[%c384] : !pto.ptr -> !pto.vreg<64xf32> + %42 = pto.vsel %41, %9, %40 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %43 = pto.vcmax %42, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %44 = pto.vmax %9, %43, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c384_35 = arith.constant 384 : index + %c1_i32_36 = arith.constant 1 : i32 + %mask_37, %scalar_out_38 = pto.plt_b32 %c1_i32_36 : i32 -> !pto.mask, i32 + pto.vsts %44, %7[%c384_35], %mask_37 {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %45 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c448 = arith.constant 448 : index + %46 = pto.vlds %6[%c448] : !pto.ptr -> !pto.vreg<64xf32> + %47 = pto.vsel %46, %9, %45 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %48 = pto.vcmax %47, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %49 = pto.vmax %9, %48, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c448_39 = arith.constant 448 : index + %c1_i32_40 = arith.constant 1 : i32 + %mask_41, %scalar_out_42 = pto.plt_b32 %c1_i32_40 : i32 -> !pto.mask, i32 + pto.vsts %49, %7[%c448_39], %mask_41 {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + pto.tstore ins(%5 : !pto.tile_buf) outs(%3 : !pto.partition_tensor_view<8x1xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/row_min.pto b/ptodsl/lib/a5/generated/tile_ops/row_min.pto new file mode 100644 index 00000000..394d2a95 --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/row_min.pto @@ -0,0 +1,118 @@ +module { + func.func @tile_op_row_min(%arg0: !pto.ptr, %arg1: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c1_1 = arith.constant 1 : index + %c1_2 = arith.constant 1 : index + %c8_3 = arith.constant 8 : index + %c1_4 = arith.constant 1 : index + %1 = pto.make_tensor_view %arg1, shape = [%c8_3, %c1_4], strides = [%c1_1, %c1_2] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_5 = arith.constant 0 : index + %c8_6 = arith.constant 8 : index + %c64_7 = arith.constant 64 : index + %2 = pto.partition_view %0, offsets = [%c0, %c0_5], sizes = [%c8_6, %c64_7] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_8 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c8_10 = arith.constant 8 : index + %c1_11 = arith.constant 1 : index + %3 = pto.partition_view %1, offsets = [%c0_8, %c0_9], sizes = [%c8_10, %c1_11] : !pto.tensor_view -> !pto.partition_tensor_view<8x1xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %4 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %5 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + pto.tload ins(%2 : !pto.partition_tensor_view<8x64xf32>) outs(%4 : !pto.tile_buf) + %6 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %7 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %8 = pto.pset_b32 "PAT_ALL" : !pto.mask + %cst = arith.constant 0x7F800000 : f32 + %9 = pto.vbr %cst : f32 -> !pto.vreg<64xf32> + %10 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c0_12 = arith.constant 0 : index + %11 = pto.vlds %6[%c0_12] : !pto.ptr -> !pto.vreg<64xf32> + %12 = pto.vsel %11, %9, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %13 = pto.vcmin %12, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %14 = pto.vmin %9, %13, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c0_13 = arith.constant 0 : index + %c1_i32 = arith.constant 1 : i32 + %mask, %scalar_out = pto.plt_b32 %c1_i32 : i32 -> !pto.mask, i32 + pto.vsts %14, %7[%c0_13], %mask {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %15 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c64_14 = arith.constant 64 : index + %16 = pto.vlds %6[%c64_14] : !pto.ptr -> !pto.vreg<64xf32> + %17 = pto.vsel %16, %9, %15 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %18 = pto.vcmin %17, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %19 = pto.vmin %9, %18, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_15 = arith.constant 64 : index + %c1_i32_16 = arith.constant 1 : i32 + %mask_17, %scalar_out_18 = pto.plt_b32 %c1_i32_16 : i32 -> !pto.mask, i32 + pto.vsts %19, %7[%c64_15], %mask_17 {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %20 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c128 = arith.constant 128 : index + %21 = pto.vlds %6[%c128] : !pto.ptr -> !pto.vreg<64xf32> + %22 = pto.vsel %21, %9, %20 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %23 = pto.vcmin %22, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %24 = pto.vmin %9, %23, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c128_19 = arith.constant 128 : index + %c1_i32_20 = arith.constant 1 : i32 + %mask_21, %scalar_out_22 = pto.plt_b32 %c1_i32_20 : i32 -> !pto.mask, i32 + pto.vsts %24, %7[%c128_19], %mask_21 {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %25 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c192 = arith.constant 192 : index + %26 = pto.vlds %6[%c192] : !pto.ptr -> !pto.vreg<64xf32> + %27 = pto.vsel %26, %9, %25 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %28 = pto.vcmin %27, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %29 = pto.vmin %9, %28, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c192_23 = arith.constant 192 : index + %c1_i32_24 = arith.constant 1 : i32 + %mask_25, %scalar_out_26 = pto.plt_b32 %c1_i32_24 : i32 -> !pto.mask, i32 + pto.vsts %29, %7[%c192_23], %mask_25 {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %30 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c256 = arith.constant 256 : index + %31 = pto.vlds %6[%c256] : !pto.ptr -> !pto.vreg<64xf32> + %32 = pto.vsel %31, %9, %30 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %33 = pto.vcmin %32, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %34 = pto.vmin %9, %33, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c256_27 = arith.constant 256 : index + %c1_i32_28 = arith.constant 1 : i32 + %mask_29, %scalar_out_30 = pto.plt_b32 %c1_i32_28 : i32 -> !pto.mask, i32 + pto.vsts %34, %7[%c256_27], %mask_29 {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %35 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c320 = arith.constant 320 : index + %36 = pto.vlds %6[%c320] : !pto.ptr -> !pto.vreg<64xf32> + %37 = pto.vsel %36, %9, %35 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %38 = pto.vcmin %37, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %39 = pto.vmin %9, %38, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c320_31 = arith.constant 320 : index + %c1_i32_32 = arith.constant 1 : i32 + %mask_33, %scalar_out_34 = pto.plt_b32 %c1_i32_32 : i32 -> !pto.mask, i32 + pto.vsts %39, %7[%c320_31], %mask_33 {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %40 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c384 = arith.constant 384 : index + %41 = pto.vlds %6[%c384] : !pto.ptr -> !pto.vreg<64xf32> + %42 = pto.vsel %41, %9, %40 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %43 = pto.vcmin %42, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %44 = pto.vmin %9, %43, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c384_35 = arith.constant 384 : index + %c1_i32_36 = arith.constant 1 : i32 + %mask_37, %scalar_out_38 = pto.plt_b32 %c1_i32_36 : i32 -> !pto.mask, i32 + pto.vsts %44, %7[%c384_35], %mask_37 {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %45 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c448 = arith.constant 448 : index + %46 = pto.vlds %6[%c448] : !pto.ptr -> !pto.vreg<64xf32> + %47 = pto.vsel %46, %9, %45 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %48 = pto.vcmin %47, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %49 = pto.vmin %9, %48, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c448_39 = arith.constant 448 : index + %c1_i32_40 = arith.constant 1 : i32 + %mask_41, %scalar_out_42 = pto.plt_b32 %c1_i32_40 : i32 -> !pto.mask, i32 + pto.vsts %49, %7[%c448_39], %mask_41 {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + pto.tstore ins(%5 : !pto.tile_buf) outs(%3 : !pto.partition_tensor_view<8x1xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/row_sum.pto b/ptodsl/lib/a5/generated/tile_ops/row_sum.pto new file mode 100644 index 00000000..61c20c50 --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/row_sum.pto @@ -0,0 +1,118 @@ +module { + func.func @tile_op_row_sum(%arg0: !pto.ptr, %arg1: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c1_1 = arith.constant 1 : index + %c1_2 = arith.constant 1 : index + %c8_3 = arith.constant 8 : index + %c1_4 = arith.constant 1 : index + %1 = pto.make_tensor_view %arg1, shape = [%c8_3, %c1_4], strides = [%c1_1, %c1_2] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_5 = arith.constant 0 : index + %c8_6 = arith.constant 8 : index + %c64_7 = arith.constant 64 : index + %2 = pto.partition_view %0, offsets = [%c0, %c0_5], sizes = [%c8_6, %c64_7] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_8 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c8_10 = arith.constant 8 : index + %c1_11 = arith.constant 1 : index + %3 = pto.partition_view %1, offsets = [%c0_8, %c0_9], sizes = [%c8_10, %c1_11] : !pto.tensor_view -> !pto.partition_tensor_view<8x1xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %4 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %5 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + pto.tload ins(%2 : !pto.partition_tensor_view<8x64xf32>) outs(%4 : !pto.tile_buf) + %6 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %7 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %8 = pto.pset_b32 "PAT_ALL" : !pto.mask + %cst = arith.constant 0.000000e+00 : f32 + %9 = pto.vbr %cst : f32 -> !pto.vreg<64xf32> + %10 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c0_12 = arith.constant 0 : index + %11 = pto.vlds %6[%c0_12] : !pto.ptr -> !pto.vreg<64xf32> + %12 = pto.vsel %11, %9, %10 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %13 = pto.vcadd %12, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %14 = pto.vadd %9, %13, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c0_13 = arith.constant 0 : index + %c1_i32 = arith.constant 1 : i32 + %mask, %scalar_out = pto.plt_b32 %c1_i32 : i32 -> !pto.mask, i32 + pto.vsts %14, %7[%c0_13], %mask {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %15 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c64_14 = arith.constant 64 : index + %16 = pto.vlds %6[%c64_14] : !pto.ptr -> !pto.vreg<64xf32> + %17 = pto.vsel %16, %9, %15 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %18 = pto.vcadd %17, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %19 = pto.vadd %9, %18, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c64_15 = arith.constant 64 : index + %c1_i32_16 = arith.constant 1 : i32 + %mask_17, %scalar_out_18 = pto.plt_b32 %c1_i32_16 : i32 -> !pto.mask, i32 + pto.vsts %19, %7[%c64_15], %mask_17 {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %20 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c128 = arith.constant 128 : index + %21 = pto.vlds %6[%c128] : !pto.ptr -> !pto.vreg<64xf32> + %22 = pto.vsel %21, %9, %20 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %23 = pto.vcadd %22, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %24 = pto.vadd %9, %23, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c128_19 = arith.constant 128 : index + %c1_i32_20 = arith.constant 1 : i32 + %mask_21, %scalar_out_22 = pto.plt_b32 %c1_i32_20 : i32 -> !pto.mask, i32 + pto.vsts %24, %7[%c128_19], %mask_21 {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %25 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c192 = arith.constant 192 : index + %26 = pto.vlds %6[%c192] : !pto.ptr -> !pto.vreg<64xf32> + %27 = pto.vsel %26, %9, %25 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %28 = pto.vcadd %27, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %29 = pto.vadd %9, %28, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c192_23 = arith.constant 192 : index + %c1_i32_24 = arith.constant 1 : i32 + %mask_25, %scalar_out_26 = pto.plt_b32 %c1_i32_24 : i32 -> !pto.mask, i32 + pto.vsts %29, %7[%c192_23], %mask_25 {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %30 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c256 = arith.constant 256 : index + %31 = pto.vlds %6[%c256] : !pto.ptr -> !pto.vreg<64xf32> + %32 = pto.vsel %31, %9, %30 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %33 = pto.vcadd %32, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %34 = pto.vadd %9, %33, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c256_27 = arith.constant 256 : index + %c1_i32_28 = arith.constant 1 : i32 + %mask_29, %scalar_out_30 = pto.plt_b32 %c1_i32_28 : i32 -> !pto.mask, i32 + pto.vsts %34, %7[%c256_27], %mask_29 {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %35 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c320 = arith.constant 320 : index + %36 = pto.vlds %6[%c320] : !pto.ptr -> !pto.vreg<64xf32> + %37 = pto.vsel %36, %9, %35 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %38 = pto.vcadd %37, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %39 = pto.vadd %9, %38, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c320_31 = arith.constant 320 : index + %c1_i32_32 = arith.constant 1 : i32 + %mask_33, %scalar_out_34 = pto.plt_b32 %c1_i32_32 : i32 -> !pto.mask, i32 + pto.vsts %39, %7[%c320_31], %mask_33 {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %40 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c384 = arith.constant 384 : index + %41 = pto.vlds %6[%c384] : !pto.ptr -> !pto.vreg<64xf32> + %42 = pto.vsel %41, %9, %40 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %43 = pto.vcadd %42, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %44 = pto.vadd %9, %43, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c384_35 = arith.constant 384 : index + %c1_i32_36 = arith.constant 1 : i32 + %mask_37, %scalar_out_38 = pto.plt_b32 %c1_i32_36 : i32 -> !pto.mask, i32 + pto.vsts %44, %7[%c384_35], %mask_37 {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %45 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c448 = arith.constant 448 : index + %46 = pto.vlds %6[%c448] : !pto.ptr -> !pto.vreg<64xf32> + %47 = pto.vsel %46, %9, %45 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %48 = pto.vcadd %47, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %49 = pto.vadd %9, %48, %8 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %c448_39 = arith.constant 448 : index + %c1_i32_40 = arith.constant 1 : i32 + %mask_41, %scalar_out_42 = pto.plt_b32 %c1_i32_40 : i32 -> !pto.mask, i32 + pto.vsts %49, %7[%c448_39], %mask_41 {dist = "ONEPT_B32"} : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + pto.tstore ins(%5 : !pto.tile_buf) outs(%3 : !pto.partition_tensor_view<8x1xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/rsqrt.pto b/ptodsl/lib/a5/generated/tile_ops/rsqrt.pto new file mode 100644 index 00000000..17752171 --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/rsqrt.pto @@ -0,0 +1,83 @@ +module { + func.func @tile_op_rsqrt(%arg0: !pto.ptr, %arg1: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c64_1 = arith.constant 64 : index + %c1_2 = arith.constant 1 : index + %c8_3 = arith.constant 8 : index + %c64_4 = arith.constant 64 : index + %1 = pto.make_tensor_view %arg1, shape = [%c8_3, %c64_4], strides = [%c64_1, %c1_2] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_5 = arith.constant 0 : index + %c8_6 = arith.constant 8 : index + %c64_7 = arith.constant 64 : index + %2 = pto.partition_view %0, offsets = [%c0, %c0_5], sizes = [%c8_6, %c64_7] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_8 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c8_10 = arith.constant 8 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %1, offsets = [%c0_8, %c0_9], sizes = [%c8_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %4 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %5 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + pto.tload ins(%2 : !pto.partition_tensor_view<8x64xf32>) outs(%4 : !pto.tile_buf) + %6 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %7 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %8 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c0_12 = arith.constant 0 : index + %9 = pto.vlds %6[%c0_12] : !pto.ptr -> !pto.vreg<64xf32> + %10 = pto.vsqrt %9, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %11 = pto.vrec %10, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %11, %7[%c0_12], %8 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %12 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c64_13 = arith.constant 64 : index + %13 = pto.vlds %6[%c64_13] : !pto.ptr -> !pto.vreg<64xf32> + %14 = pto.vsqrt %13, %12 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %15 = pto.vrec %14, %12 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %15, %7[%c64_13], %12 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %16 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c128 = arith.constant 128 : index + %17 = pto.vlds %6[%c128] : !pto.ptr -> !pto.vreg<64xf32> + %18 = pto.vsqrt %17, %16 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %19 = pto.vrec %18, %16 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %19, %7[%c128], %16 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %20 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c192 = arith.constant 192 : index + %21 = pto.vlds %6[%c192] : !pto.ptr -> !pto.vreg<64xf32> + %22 = pto.vsqrt %21, %20 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %23 = pto.vrec %22, %20 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %23, %7[%c192], %20 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %24 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c256 = arith.constant 256 : index + %25 = pto.vlds %6[%c256] : !pto.ptr -> !pto.vreg<64xf32> + %26 = pto.vsqrt %25, %24 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %27 = pto.vrec %26, %24 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %27, %7[%c256], %24 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %28 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c320 = arith.constant 320 : index + %29 = pto.vlds %6[%c320] : !pto.ptr -> !pto.vreg<64xf32> + %30 = pto.vsqrt %29, %28 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %31 = pto.vrec %30, %28 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %31, %7[%c320], %28 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %32 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c384 = arith.constant 384 : index + %33 = pto.vlds %6[%c384] : !pto.ptr -> !pto.vreg<64xf32> + %34 = pto.vsqrt %33, %32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %35 = pto.vrec %34, %32 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %35, %7[%c384], %32 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %36 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c448 = arith.constant 448 : index + %37 = pto.vlds %6[%c448] : !pto.ptr -> !pto.vreg<64xf32> + %38 = pto.vsqrt %37, %36 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %39 = pto.vrec %38, %36 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %39, %7[%c448], %36 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + pto.tstore ins(%5 : !pto.tile_buf) outs(%3 : !pto.partition_tensor_view<8x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/sort32.pto b/ptodsl/lib/a5/generated/tile_ops/sort32.pto new file mode 100644 index 00000000..c739860e --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/sort32.pto @@ -0,0 +1,57 @@ +module { + func.func @tile_op_sort32(%arg0: !pto.ptr, %arg1: !pto.ptr, %arg2: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c1_0 = arith.constant 1 : index + %c64_1 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c1_0, %c64_1], strides = [%c64, %c1] : !pto.tensor_view + %c64_2 = arith.constant 64 : index + %c1_3 = arith.constant 1 : index + %c1_4 = arith.constant 1 : index + %c64_5 = arith.constant 64 : index + %1 = pto.make_tensor_view %arg1, shape = [%c1_4, %c64_5], strides = [%c64_2, %c1_3] : !pto.tensor_view + %c128 = arith.constant 128 : index + %c1_6 = arith.constant 1 : index + %c1_7 = arith.constant 1 : index + %c128_8 = arith.constant 128 : index + %2 = pto.make_tensor_view %arg2, shape = [%c1_7, %c128_8], strides = [%c128, %c1_6] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c1_10 = arith.constant 1 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %0, offsets = [%c0, %c0_9], sizes = [%c1_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<1x64xf32> + %c0_12 = arith.constant 0 : index + %c0_13 = arith.constant 0 : index + %c1_14 = arith.constant 1 : index + %c64_15 = arith.constant 64 : index + %4 = pto.partition_view %1, offsets = [%c0_12, %c0_13], sizes = [%c1_14, %c64_15] : !pto.tensor_view -> !pto.partition_tensor_view<1x64xui32> + %c0_16 = arith.constant 0 : index + %c0_17 = arith.constant 0 : index + %c1_18 = arith.constant 1 : index + %c128_19 = arith.constant 128 : index + %5 = pto.partition_view %2, offsets = [%c0_16, %c0_17], sizes = [%c1_18, %c128_19] : !pto.tensor_view -> !pto.partition_tensor_view<1x128xf32> + %c0_i64 = arith.constant 0 : i64 + %c256_i64 = arith.constant 256 : i64 + %c512_i64 = arith.constant 512 : i64 + %6 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %7 = pto.alloc_tile addr = %c256_i64 : !pto.tile_buf + %8 = pto.alloc_tile addr = %c512_i64 : !pto.tile_buf + pto.tload ins(%3 : !pto.partition_tensor_view<1x64xf32>) outs(%6 : !pto.tile_buf) + pto.tload ins(%4 : !pto.partition_tensor_view<1x64xui32>) outs(%7 : !pto.tile_buf) + %9 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %10 = pto.castptr %c256_i64 : i64 -> !pto.ptr + %11 = pto.castptr %c512_i64 : i64 -> !pto.ptr + %c2 = arith.constant 2 : index + %c0_20 = arith.constant 0 : index + %12 = pto.addptr %9, %c0_20 : -> + %c0_21 = arith.constant 0 : index + %13 = pto.addptr %10, %c0_21 : -> + %c0_22 = arith.constant 0 : index + %14 = pto.addptr %11, %c0_22 : -> + pto.vbitsort %14, %12, %13, %c2 : !pto.ptr, !pto.ptr, !pto.ptr, index + pto.tstore ins(%8 : !pto.tile_buf) outs(%5 : !pto.partition_tensor_view<1x128xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/sqrt.pto b/ptodsl/lib/a5/generated/tile_ops/sqrt.pto new file mode 100644 index 00000000..4de655e0 --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/sqrt.pto @@ -0,0 +1,75 @@ +module { + func.func @tile_op_sqrt(%arg0: !pto.ptr, %arg1: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c64_1 = arith.constant 64 : index + %c1_2 = arith.constant 1 : index + %c8_3 = arith.constant 8 : index + %c64_4 = arith.constant 64 : index + %1 = pto.make_tensor_view %arg1, shape = [%c8_3, %c64_4], strides = [%c64_1, %c1_2] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_5 = arith.constant 0 : index + %c8_6 = arith.constant 8 : index + %c64_7 = arith.constant 64 : index + %2 = pto.partition_view %0, offsets = [%c0, %c0_5], sizes = [%c8_6, %c64_7] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_8 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c8_10 = arith.constant 8 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %1, offsets = [%c0_8, %c0_9], sizes = [%c8_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %4 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %5 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + pto.tload ins(%2 : !pto.partition_tensor_view<8x64xf32>) outs(%4 : !pto.tile_buf) + %6 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %7 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %8 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c0_12 = arith.constant 0 : index + %9 = pto.vlds %6[%c0_12] : !pto.ptr -> !pto.vreg<64xf32> + %10 = pto.vsqrt %9, %8 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %10, %7[%c0_12], %8 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %11 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c64_13 = arith.constant 64 : index + %12 = pto.vlds %6[%c64_13] : !pto.ptr -> !pto.vreg<64xf32> + %13 = pto.vsqrt %12, %11 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %13, %7[%c64_13], %11 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %14 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c128 = arith.constant 128 : index + %15 = pto.vlds %6[%c128] : !pto.ptr -> !pto.vreg<64xf32> + %16 = pto.vsqrt %15, %14 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %16, %7[%c128], %14 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %17 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c192 = arith.constant 192 : index + %18 = pto.vlds %6[%c192] : !pto.ptr -> !pto.vreg<64xf32> + %19 = pto.vsqrt %18, %17 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %19, %7[%c192], %17 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %20 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c256 = arith.constant 256 : index + %21 = pto.vlds %6[%c256] : !pto.ptr -> !pto.vreg<64xf32> + %22 = pto.vsqrt %21, %20 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %22, %7[%c256], %20 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %23 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c320 = arith.constant 320 : index + %24 = pto.vlds %6[%c320] : !pto.ptr -> !pto.vreg<64xf32> + %25 = pto.vsqrt %24, %23 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %25, %7[%c320], %23 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %26 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c384 = arith.constant 384 : index + %27 = pto.vlds %6[%c384] : !pto.ptr -> !pto.vreg<64xf32> + %28 = pto.vsqrt %27, %26 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %28, %7[%c384], %26 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %29 = pto.pset_b32 "PAT_ALL" : !pto.mask + %c448 = arith.constant 448 : index + %30 = pto.vlds %6[%c448] : !pto.ptr -> !pto.vreg<64xf32> + %31 = pto.vsqrt %30, %29 : !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %31, %7[%c448], %29 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + pto.tstore ins(%5 : !pto.tile_buf) outs(%3 : !pto.partition_tensor_view<8x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/generated/tile_ops/sub.pto b/ptodsl/lib/a5/generated/tile_ops/sub.pto new file mode 100644 index 00000000..e81c72ae --- /dev/null +++ b/ptodsl/lib/a5/generated/tile_ops/sub.pto @@ -0,0 +1,90 @@ +module { + func.func @tile_op_sub(%arg0: !pto.ptr, %arg1: !pto.ptr, %arg2: !pto.ptr) { + %c64 = arith.constant 64 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64_0 = arith.constant 64 : index + %0 = pto.make_tensor_view %arg0, shape = [%c8, %c64_0], strides = [%c64, %c1] : !pto.tensor_view + %c64_1 = arith.constant 64 : index + %c1_2 = arith.constant 1 : index + %c8_3 = arith.constant 8 : index + %c64_4 = arith.constant 64 : index + %1 = pto.make_tensor_view %arg1, shape = [%c8_3, %c64_4], strides = [%c64_1, %c1_2] : !pto.tensor_view + %c64_5 = arith.constant 64 : index + %c1_6 = arith.constant 1 : index + %c8_7 = arith.constant 8 : index + %c64_8 = arith.constant 64 : index + %2 = pto.make_tensor_view %arg2, shape = [%c8_7, %c64_8], strides = [%c64_5, %c1_6] : !pto.tensor_view + pto.section.vector { + %c0 = arith.constant 0 : index + %c0_9 = arith.constant 0 : index + %c8_10 = arith.constant 8 : index + %c64_11 = arith.constant 64 : index + %3 = pto.partition_view %0, offsets = [%c0, %c0_9], sizes = [%c8_10, %c64_11] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_12 = arith.constant 0 : index + %c0_13 = arith.constant 0 : index + %c8_14 = arith.constant 8 : index + %c64_15 = arith.constant 64 : index + %4 = pto.partition_view %1, offsets = [%c0_12, %c0_13], sizes = [%c8_14, %c64_15] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_16 = arith.constant 0 : index + %c0_17 = arith.constant 0 : index + %c8_18 = arith.constant 8 : index + %c64_19 = arith.constant 64 : index + %5 = pto.partition_view %2, offsets = [%c0_16, %c0_17], sizes = [%c8_18, %c64_19] : !pto.tensor_view -> !pto.partition_tensor_view<8x64xf32> + %c0_i64 = arith.constant 0 : i64 + %c2048_i64 = arith.constant 2048 : i64 + %c4096_i64 = arith.constant 4096 : i64 + %6 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %7 = pto.alloc_tile addr = %c2048_i64 : !pto.tile_buf + %8 = pto.alloc_tile addr = %c4096_i64 : !pto.tile_buf + pto.tload ins(%3 : !pto.partition_tensor_view<8x64xf32>) outs(%6 : !pto.tile_buf) + pto.tload ins(%4 : !pto.partition_tensor_view<8x64xf32>) outs(%7 : !pto.tile_buf) + %9 = pto.castptr %c0_i64 : i64 -> !pto.ptr + %10 = pto.castptr %c2048_i64 : i64 -> !pto.ptr + %11 = pto.castptr %c4096_i64 : i64 -> !pto.ptr + %c64_20 = arith.constant 64 : index + %12 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result, %updated_source = pto.vlds_post %9[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_21, %updated_source_22 = pto.vlds_post %10[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %13 = pto.vsub %result, %result_21, %12 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %14 = pto.vsts_post %13, %11[%c64_20], %12 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %15 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_23, %updated_source_24 = pto.vlds_post %updated_source[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_25, %updated_source_26 = pto.vlds_post %updated_source_22[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %16 = pto.vsub %result_23, %result_25, %15 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %17 = pto.vsts_post %16, %14[%c64_20], %15 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %18 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_27, %updated_source_28 = pto.vlds_post %updated_source_24[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_29, %updated_source_30 = pto.vlds_post %updated_source_26[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %19 = pto.vsub %result_27, %result_29, %18 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %20 = pto.vsts_post %19, %17[%c64_20], %18 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %21 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_31, %updated_source_32 = pto.vlds_post %updated_source_28[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_33, %updated_source_34 = pto.vlds_post %updated_source_30[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %22 = pto.vsub %result_31, %result_33, %21 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %23 = pto.vsts_post %22, %20[%c64_20], %21 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %24 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_35, %updated_source_36 = pto.vlds_post %updated_source_32[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_37, %updated_source_38 = pto.vlds_post %updated_source_34[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %25 = pto.vsub %result_35, %result_37, %24 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %26 = pto.vsts_post %25, %23[%c64_20], %24 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %27 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_39, %updated_source_40 = pto.vlds_post %updated_source_36[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_41, %updated_source_42 = pto.vlds_post %updated_source_38[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %28 = pto.vsub %result_39, %result_41, %27 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %29 = pto.vsts_post %28, %26[%c64_20], %27 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %30 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_43, %updated_source_44 = pto.vlds_post %updated_source_40[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_45, %updated_source_46 = pto.vlds_post %updated_source_42[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %31 = pto.vsub %result_43, %result_45, %30 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %32 = pto.vsts_post %31, %29[%c64_20], %30 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %33 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_47, %updated_source_48 = pto.vlds_post %updated_source_44[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_49, %updated_source_50 = pto.vlds_post %updated_source_46[%c64_20] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %34 = pto.vsub %result_47, %result_49, %33 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %35 = pto.vsts_post %34, %32[%c64_20], %33 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + pto.tstore ins(%8 : !pto.tile_buf) outs(%5 : !pto.partition_tensor_view<8x64xf32>) + } + return + } +} diff --git a/ptodsl/lib/a5/tile_op_kernels.py b/ptodsl/lib/a5/tile_op_kernels.py new file mode 100644 index 00000000..d4da63e4 --- /dev/null +++ b/ptodsl/lib/a5/tile_op_kernels.py @@ -0,0 +1,498 @@ +from ... import pto, to_ir_module +from . import ops +from .tile_micro_coverage import TILE_MICRO_COVERAGE + + +_MATRIX_SHAPE = [8, 64] +_ROW_VECTOR_SHAPE = [1, 64] +_ROW_REDUCE_OUT_SHAPE = [8, 1] +_COL_REDUCE_OUT_SHAPE = [1, 64] +_ROW_EXPAND_SRC_SHAPE = [8, 1] +_SORT_SHAPE = [1, 64] +_SORT32_OUT_SHAPE = [1, 128] +_BLOCK_LEN = 16 + + +def _resolve_dtype(dtype_name): + return getattr(pto, dtype_name) + + +def _tensor(ptr_value, *, shape, dtype): + return pto.make_tensor(ptr_value, shape=shape, dtype=dtype) + + +def _tile(view, shape): + return view.slice([0, 0], shape).raw + + +def _binary_kernel(kernel_name, helper, *, dtype_name, shape): + def meta_data(): + dtype = _resolve_dtype(dtype_name) + return {"ptr_t": pto.ptr(dtype)} + + def kernel(src0: "ptr_t", src1: "ptr_t", dst: "ptr_t") -> None: + dtype = _resolve_dtype(dtype_name) + lhs = _tensor(src0, shape=shape, dtype=dtype) + rhs = _tensor(src1, shape=shape, dtype=dtype) + out = _tensor(dst, shape=shape, dtype=dtype) + + with pto.vector_section(): + helper( + _tile(lhs, shape), + _tile(rhs, shape), + _tile(out, shape), + dtype=dtype, + tile_shape=shape, + ) + + kernel.__name__ = kernel_name + return to_ir_module(meta_data=meta_data)(kernel) + + +def _move_kernel(kernel_name, helper, *, dtype_name, shape): + def meta_data(): + dtype = _resolve_dtype(dtype_name) + return {"ptr_t": pto.ptr(dtype)} + + def kernel(src: "ptr_t", dst: "ptr_t") -> None: + dtype = _resolve_dtype(dtype_name) + src_view = _tensor(src, shape=shape, dtype=dtype) + out_view = _tensor(dst, shape=shape, dtype=dtype) + + with pto.vector_section(): + helper( + _tile(src_view, shape), + _tile(out_view, shape), + dtype=dtype, + tile_shape=shape, + ) + + kernel.__name__ = kernel_name + return to_ir_module(meta_data=meta_data)(kernel) + + +def _unary_kernel(kernel_name, helper, *, dtype_name, shape): + def meta_data(): + dtype = _resolve_dtype(dtype_name) + return {"ptr_t": pto.ptr(dtype)} + + def kernel(src: "ptr_t", dst: "ptr_t") -> None: + dtype = _resolve_dtype(dtype_name) + inp = _tensor(src, shape=shape, dtype=dtype) + out = _tensor(dst, shape=shape, dtype=dtype) + + with pto.vector_section(): + helper( + _tile(inp, shape), + _tile(out, shape), + dtype=dtype, + tile_shape=shape, + ) + + kernel.__name__ = kernel_name + return to_ir_module(meta_data=meta_data)(kernel) + + +def _row_reduce_kernel(kernel_name, helper, *, dtype_name): + def meta_data(): + dtype = _resolve_dtype(dtype_name) + return { + "src_ptr_t": pto.ptr(dtype), + "dst_ptr_t": pto.ptr(dtype), + } + + def kernel(src: "src_ptr_t", dst: "dst_ptr_t") -> None: + dtype = _resolve_dtype(dtype_name) + src_view = _tensor(src, shape=_MATRIX_SHAPE, dtype=dtype) + out_view = _tensor(dst, shape=_ROW_REDUCE_OUT_SHAPE, dtype=dtype) + + with pto.vector_section(): + helper( + _tile(src_view, _MATRIX_SHAPE), + _tile(out_view, _ROW_REDUCE_OUT_SHAPE), + dtype=dtype, + tile_shape=_MATRIX_SHAPE, + ) + + kernel.__name__ = kernel_name + return to_ir_module(meta_data=meta_data)(kernel) + + +def _col_reduce_kernel(kernel_name, helper, *, dtype_name): + def meta_data(): + dtype = _resolve_dtype(dtype_name) + return { + "src_ptr_t": pto.ptr(dtype), + "dst_ptr_t": pto.ptr(dtype), + } + + def kernel(src: "src_ptr_t", dst: "dst_ptr_t") -> None: + dtype = _resolve_dtype(dtype_name) + src_view = _tensor(src, shape=_MATRIX_SHAPE, dtype=dtype) + out_view = _tensor(dst, shape=_COL_REDUCE_OUT_SHAPE, dtype=dtype) + + with pto.vector_section(): + helper( + _tile(src_view, _MATRIX_SHAPE), + _tile(out_view, _COL_REDUCE_OUT_SHAPE), + dtype=dtype, + tile_shape=_MATRIX_SHAPE, + ) + + kernel.__name__ = kernel_name + return to_ir_module(meta_data=meta_data)(kernel) + + +def _row_expand_kernel(kernel_name, helper, *, dtype_name): + def meta_data(): + dtype = _resolve_dtype(dtype_name) + return { + "src_ptr_t": pto.ptr(dtype), + "dst_ptr_t": pto.ptr(dtype), + } + + def kernel(src: "src_ptr_t", dst: "dst_ptr_t") -> None: + dtype = _resolve_dtype(dtype_name) + src_view = _tensor(src, shape=_ROW_EXPAND_SRC_SHAPE, dtype=dtype) + out_view = _tensor(dst, shape=_MATRIX_SHAPE, dtype=dtype) + + with pto.vector_section(): + helper( + _tile(src_view, _ROW_EXPAND_SRC_SHAPE), + _tile(out_view, _MATRIX_SHAPE), + dtype=dtype, + tile_shape=_MATRIX_SHAPE, + ) + + kernel.__name__ = kernel_name + return to_ir_module(meta_data=meta_data)(kernel) + + +def _row_expand_binary_kernel(kernel_name, helper, *, dtype_name): + def meta_data(): + dtype = _resolve_dtype(dtype_name) + return { + "base_ptr_t": pto.ptr(dtype), + "expand_ptr_t": pto.ptr(dtype), + "dst_ptr_t": pto.ptr(dtype), + } + + def kernel(base: "base_ptr_t", expand: "expand_ptr_t", dst: "dst_ptr_t") -> None: + dtype = _resolve_dtype(dtype_name) + base_view = _tensor(base, shape=_MATRIX_SHAPE, dtype=dtype) + expand_view = _tensor(expand, shape=_ROW_EXPAND_SRC_SHAPE, dtype=dtype) + out_view = _tensor(dst, shape=_MATRIX_SHAPE, dtype=dtype) + + with pto.vector_section(): + helper( + _tile(base_view, _MATRIX_SHAPE), + _tile(expand_view, _ROW_EXPAND_SRC_SHAPE), + _tile(out_view, _MATRIX_SHAPE), + dtype=dtype, + tile_shape=_MATRIX_SHAPE, + ) + + kernel.__name__ = kernel_name + return to_ir_module(meta_data=meta_data)(kernel) + + +def _col_expand_kernel(kernel_name, helper, *, dtype_name): + def meta_data(): + dtype = _resolve_dtype(dtype_name) + return { + "src_ptr_t": pto.ptr(dtype), + "dst_ptr_t": pto.ptr(dtype), + } + + def kernel(src: "src_ptr_t", dst: "dst_ptr_t") -> None: + dtype = _resolve_dtype(dtype_name) + src_view = _tensor(src, shape=_ROW_VECTOR_SHAPE, dtype=dtype) + out_view = _tensor(dst, shape=_MATRIX_SHAPE, dtype=dtype) + + with pto.vector_section(): + helper( + _tile(src_view, _ROW_VECTOR_SHAPE), + _tile(out_view, _MATRIX_SHAPE), + dtype=dtype, + tile_shape=_MATRIX_SHAPE, + ) + + kernel.__name__ = kernel_name + return to_ir_module(meta_data=meta_data)(kernel) + + +def _gather_kernel(kernel_name, *, dtype_name): + def meta_data(): + dtype = _resolve_dtype(dtype_name) + return { + "src_ptr_t": pto.ptr(dtype), + "idx_ptr_t": pto.ptr(_resolve_dtype("uint32")), + "dst_ptr_t": pto.ptr(dtype), + } + + def kernel(src: "src_ptr_t", idx: "idx_ptr_t", dst: "dst_ptr_t") -> None: + dtype = _resolve_dtype(dtype_name) + index_dtype = _resolve_dtype("uint32") + src_view = _tensor(src, shape=_SORT_SHAPE, dtype=dtype) + idx_view = _tensor(idx, shape=_SORT_SHAPE, dtype=index_dtype) + out_view = _tensor(dst, shape=_SORT_SHAPE, dtype=dtype) + + with pto.vector_section(): + ops.tgather( + _tile(src_view, _SORT_SHAPE), + _tile(idx_view, _SORT_SHAPE), + _tile(out_view, _SORT_SHAPE), + dtype=dtype, + index_dtype=index_dtype, + tile_shape=_SORT_SHAPE, + ) + + kernel.__name__ = kernel_name + return to_ir_module(meta_data=meta_data)(kernel) + + +def _mrgsort_kernel(kernel_name, *, dtype_name): + def meta_data(): + dtype = _resolve_dtype(dtype_name) + return { + "src_ptr_t": pto.ptr(dtype), + "dst_ptr_t": pto.ptr(dtype), + } + + def kernel(src: "src_ptr_t", dst: "dst_ptr_t") -> None: + dtype = _resolve_dtype(dtype_name) + src_view = _tensor(src, shape=_SORT_SHAPE, dtype=dtype) + out_view = _tensor(dst, shape=_SORT_SHAPE, dtype=dtype) + + with pto.vector_section(): + ops.tmrgsort( + _tile(src_view, _SORT_SHAPE), + _tile(out_view, _SORT_SHAPE), + dtype=dtype, + tile_shape=_SORT_SHAPE, + block_len=_BLOCK_LEN, + ) + + kernel.__name__ = kernel_name + return to_ir_module(meta_data=meta_data)(kernel) + + +def _sort32_kernel(kernel_name, *, dtype_name): + def meta_data(): + dtype = _resolve_dtype(dtype_name) + return { + "src_ptr_t": pto.ptr(dtype), + "idx_ptr_t": pto.ptr(_resolve_dtype("uint32")), + "dst_ptr_t": pto.ptr(dtype), + } + + def kernel(src: "src_ptr_t", idx: "idx_ptr_t", dst: "dst_ptr_t") -> None: + dtype = _resolve_dtype(dtype_name) + index_dtype = _resolve_dtype("uint32") + src_view = _tensor(src, shape=_SORT_SHAPE, dtype=dtype) + idx_view = _tensor(idx, shape=_SORT_SHAPE, dtype=index_dtype) + out_view = _tensor(dst, shape=_SORT32_OUT_SHAPE, dtype=dtype) + + with pto.vector_section(): + ops.tsort32( + _tile(src_view, _SORT_SHAPE), + _tile(idx_view, _SORT_SHAPE), + _tile(out_view, _SORT32_OUT_SHAPE), + dtype=dtype, + tile_shape=_SORT_SHAPE, + ) + + kernel.__name__ = kernel_name + return to_ir_module(meta_data=meta_data)(kernel) + + +TILE_OP_KERNEL_SPECS = { + "mov": { + "builder": lambda: _move_kernel( + "tile_op_mov", ops.tmov, dtype_name="float32", shape=_MATRIX_SHAPE + ), + "expected_tokens": ["pto.vlds", "pto.vsts"], + }, + "add": { + "builder": lambda: _binary_kernel( + "tile_op_add", ops.tadd, dtype_name="float32", shape=_MATRIX_SHAPE + ), + "expected_tokens": ["pto.vadd"], + }, + "sub": { + "builder": lambda: _binary_kernel( + "tile_op_sub", ops.tsub, dtype_name="float32", shape=_MATRIX_SHAPE + ), + "expected_tokens": ["pto.vsub"], + }, + "div": { + "builder": lambda: _binary_kernel( + "tile_op_div", ops.tdiv, dtype_name="float32", shape=_MATRIX_SHAPE + ), + "expected_tokens": ["pto.vdiv"], + }, + "mul": { + "builder": lambda: _binary_kernel( + "tile_op_mul", ops.tmul, dtype_name="float32", shape=_MATRIX_SHAPE + ), + "expected_tokens": ["pto.vmul"], + }, + "or_": { + "builder": lambda: _binary_kernel( + "tile_op_or_", ops.tor_, dtype_name="float32", shape=_MATRIX_SHAPE + ), + "expected_tokens": ["pto.vor"], + }, + "gather": { + "builder": lambda: _gather_kernel("tile_op_gather", dtype_name="float32"), + "expected_tokens": ["pto.vgather2"], + }, + "exp": { + "builder": lambda: _unary_kernel( + "tile_op_exp", ops.texp, dtype_name="float32", shape=_MATRIX_SHAPE + ), + "expected_tokens": ["pto.vexp"], + }, + "log": { + "builder": lambda: _unary_kernel( + "tile_op_log", ops.tlog, dtype_name="float32", shape=_MATRIX_SHAPE + ), + "expected_tokens": ["pto.vln"], + }, + "relu": { + "builder": lambda: _unary_kernel( + "tile_op_relu", ops.trelu, dtype_name="float32", shape=_MATRIX_SHAPE + ), + "expected_tokens": ["pto.vrelu"], + }, + "abs": { + "builder": lambda: _unary_kernel( + "tile_op_abs", ops.tabs, dtype_name="float32", shape=_MATRIX_SHAPE + ), + "expected_tokens": ["pto.vabs"], + }, + "sqrt": { + "builder": lambda: _unary_kernel( + "tile_op_sqrt", ops.tsqrt, dtype_name="float32", shape=_MATRIX_SHAPE + ), + "expected_tokens": ["pto.vsqrt"], + }, + "rsqrt": { + "builder": lambda: _unary_kernel( + "tile_op_rsqrt", ops.trsqrt, dtype_name="float32", shape=_MATRIX_SHAPE + ), + "expected_tokens": ["pto.vsqrt", "pto.vrec"], + }, + "reciprocal": { + "builder": lambda: _unary_kernel( + "tile_op_reciprocal", + ops.trecip, + dtype_name="float32", + shape=_MATRIX_SHAPE, + ), + "expected_tokens": ["pto.vrec"], + }, + "row_sum": { + "builder": lambda: _row_reduce_kernel( + "tile_op_row_sum", ops.trow_sum, dtype_name="float32" + ), + "expected_tokens": ["pto.vcadd"], + }, + "row_min": { + "builder": lambda: _row_reduce_kernel( + "tile_op_row_min", ops.trow_min, dtype_name="float32" + ), + "expected_tokens": ["pto.vcmin"], + }, + "row_max": { + "builder": lambda: _row_reduce_kernel( + "tile_op_row_max", ops.trow_max, dtype_name="float32" + ), + "expected_tokens": ["pto.vcmax"], + }, + "row_expand": { + "builder": lambda: _row_expand_kernel( + "tile_op_row_expand", ops.trow_expand, dtype_name="float32" + ), + "expected_tokens": ["pto.vdup"], + }, + "row_expand_sub": { + "builder": lambda: _row_expand_binary_kernel( + "tile_op_row_expand_sub", ops.trow_expand_sub, dtype_name="float32" + ), + "expected_tokens": ["pto.vsub"], + }, + "row_expand_div": { + "builder": lambda: _row_expand_binary_kernel( + "tile_op_row_expand_div", ops.trow_expand_div, dtype_name="float32" + ), + "expected_tokens": ["pto.vdiv"], + }, + "row_expand_mul": { + "builder": lambda: _row_expand_binary_kernel( + "tile_op_row_expand_mul", ops.trow_expand_mul, dtype_name="float32" + ), + "expected_tokens": ["pto.vmul"], + }, + "col_sum": { + "builder": lambda: _col_reduce_kernel( + "tile_op_col_sum", ops.tcol_sum, dtype_name="float32" + ), + "expected_tokens": ["pto.vadd"], + }, + "col_min": { + "builder": lambda: _col_reduce_kernel( + "tile_op_col_min", ops.tcol_min, dtype_name="float32" + ), + "expected_tokens": ["pto.vmin"], + }, + "col_max": { + "builder": lambda: _col_reduce_kernel( + "tile_op_col_max", ops.tcol_max, dtype_name="float32" + ), + "expected_tokens": ["pto.vmax"], + }, + "col_expand": { + "builder": lambda: _col_expand_kernel( + "tile_op_col_expand", ops.tcol_expand, dtype_name="float32" + ), + "expected_tokens": ["pto.vlds", "pto.vsts"], + }, + "mrgsort": { + "builder": lambda: _mrgsort_kernel("tile_op_mrgsort", dtype_name="float32"), + "expected_tokens": ["pto.vmrgsort4"], + }, + "sort32": { + "builder": lambda: _sort32_kernel("tile_op_sort32", dtype_name="float32"), + "expected_tokens": ["pto.vbitsort"], + }, +} + +TILE_OP_KERNEL_BUILDERS = { + name: spec["builder"] for name, spec in TILE_OP_KERNEL_SPECS.items() +} + + +def tile_op_generation_index_markdown(): + lines = [ + "# Tile Op PTO Generation", + "", + "| tile op | status | artifact | note |", + "| --- | --- | --- | --- |", + ] + for op_name, entry in TILE_MICRO_COVERAGE.items(): + if op_name in TILE_OP_KERNEL_BUILDERS: + artifact = f"`tile_ops/{op_name}.pto`" + status = "generated" + else: + artifact = "-" + status = entry["status"] + lines.append(f"| `{op_name}` | `{status}` | {artifact} | {entry['note']} |") + return "\n".join(lines) + "\n" + + +__all__ = [ + "TILE_OP_KERNEL_BUILDERS", + "TILE_OP_KERNEL_SPECS", + "tile_op_generation_index_markdown", +] diff --git a/scripts/generate_tile_op_pto.py b/scripts/generate_tile_op_pto.py new file mode 100644 index 00000000..2d19aa76 --- /dev/null +++ b/scripts/generate_tile_op_pto.py @@ -0,0 +1,39 @@ +#!/usr/bin/env python3 + +import pathlib +import sys + + +_ROOT = pathlib.Path(__file__).resolve().parents[1] +if str(_ROOT) not in sys.path: + sys.path.insert(0, str(_ROOT)) + +from ptodsl.lib import a5 + + +_DEFAULT_OUTPUT_DIR = _ROOT / "ptodsl" / "lib" / "a5" / "generated" / "tile_ops" + + +def emit_tile_ops(*, output_dir): + output_dir.mkdir(parents=True, exist_ok=True) + generated = [] + for op_name, builder in a5.TILE_OP_KERNEL_BUILDERS.items(): + module = builder() + path = output_dir / f"{op_name}.pto" + path.write_text(f"{module}\n", encoding="utf-8") + generated.append(path) + + index_path = output_dir / "TILE_OP_GENERATION_INDEX.md" + index_path.write_text(a5.tile_op_generation_index_markdown(), encoding="utf-8") + generated.append(index_path) + return generated + + +def main(): + generated = emit_tile_ops(output_dir=_DEFAULT_OUTPUT_DIR) + for path in generated: + print(path) + + +if __name__ == "__main__": + main() diff --git a/tests/regression/test_generated_tile_op_pto.py b/tests/regression/test_generated_tile_op_pto.py new file mode 100644 index 00000000..cea4ac4a --- /dev/null +++ b/tests/regression/test_generated_tile_op_pto.py @@ -0,0 +1,30 @@ +from pathlib import Path + +from ptodsl.lib import a5 + + +_GENERATED_DIR = ( + Path(__file__).resolve().parents[2] + / "ptodsl" + / "lib" + / "a5" + / "generated" + / "tile_ops" +) + + +def test_generated_tile_op_pto_files_exist_and_show_micro_ops(): + for op_name, spec in a5.TILE_OP_KERNEL_SPECS.items(): + text = (_GENERATED_DIR / f"{op_name}.pto").read_text(encoding="utf-8") + assert f"func.func @tile_op_{op_name}" in text + assert "pto.section.vector" in text + for token in spec["expected_tokens"]: + assert token in text + + +def test_tile_op_generation_index_covers_public_tile_ops(): + text = (_GENERATED_DIR / "TILE_OP_GENERATION_INDEX.md").read_text(encoding="utf-8") + for op_name in a5.TILE_MICRO_COVERAGE: + assert f"`{op_name}`" in text + assert "`matmul` | `blocked`" in text + assert "`subset` | `not_applicable`" in text