diff --git a/ptodsl/api/pto.py b/ptodsl/api/pto.py index f2e2d0ac..158bca41 100644 --- a/ptodsl/api/pto.py +++ b/ptodsl/api/pto.py @@ -21,6 +21,7 @@ TensorType, TileBufConfig, TileBufType, + ptr, __getattr__, ) @@ -34,6 +35,7 @@ "int16", "int32", "PtrType", + "ptr", "TensorType", "SubTensorType", "TileBufConfig", diff --git a/ptodsl/api/type_def.py b/ptodsl/api/type_def.py index 251303f6..f0ef6872 100644 --- a/ptodsl/api/type_def.py +++ b/ptodsl/api/type_def.py @@ -11,8 +11,26 @@ def __getattr__(name): raise AttributeError(f"module '{__name__}' has no attribute '{name}'") -def PtrType(dtype): - return _pto.PtrType.get(dtype) +def _resolve_address_space(memory_space): + if memory_space is None: + return _pto.AddressSpace.GM + if not isinstance(memory_space, str): + return memory_space + normalized = memory_space.strip().upper() + try: + return getattr(_pto.AddressSpace, normalized) + except AttributeError as exc: + raise ValueError(f"Unsupported memory_space '{memory_space}'.") from exc + + +def PtrType(dtype, memory_space=None): + if memory_space is None: + return _pto.PtrType.get(dtype) + return _pto.PtrType.get(dtype, _resolve_address_space(memory_space)) + + +def ptr(dtype, *, space=None): + return PtrType(dtype, memory_space=space) def TensorType(*, rank, dtype): @@ -99,6 +117,7 @@ def TileBufType(*, shape, dtype, memory_space, valid_shape=None, config=None): __all__ = [ "PtrType", + "ptr", "TensorType", "SubTensorType", "TileBufConfig", diff --git a/ptodsl/language.py b/ptodsl/language.py index 5e9cab1d..5ee76d5e 100644 --- a/ptodsl/language.py +++ b/ptodsl/language.py @@ -178,8 +178,26 @@ def __getattr__(name): raise AttributeError(f"module '{__name__}' has no attribute '{name}'") -def PtrType(dtype): - return pto.PtrType.get(dtype) +def _resolve_address_space(memory_space): + if memory_space is None: + return pto.AddressSpace.GM + if not isinstance(memory_space, str): + return memory_space + normalized = memory_space.strip().upper() + try: + return getattr(pto.AddressSpace, normalized) + except AttributeError as exc: + raise ValueError(f"Unsupported memory_space '{memory_space}'.") from exc + + +def PtrType(dtype, memory_space=None): + if memory_space is None: + return pto.PtrType.get(dtype) + return pto.PtrType.get(dtype, _resolve_address_space(memory_space)) + + +def ptr(dtype, *, space=None): + return PtrType(dtype, memory_space=space) def TensorType(*, rank, dtype): diff --git a/ptodsl/lib/a5/A5_HEADER_COVERAGE.md b/ptodsl/lib/a5/A5_HEADER_COVERAGE.md new file mode 100644 index 00000000..b864afea --- /dev/null +++ b/ptodsl/lib/a5/A5_HEADER_COVERAGE.md @@ -0,0 +1,127 @@ +# A5 Header Coverage + +- Total A5 headers tracked: `116` +- Implemented: `49` +- Partial: `2` +- Native only: `11` +- Pending: `35` +- Blocked/meta: `19` + +| header | status | helper | note | +| --- | --- | --- | --- | +| `MGather` | `pending` | `-` | Memory gather helper is not yet represented in the PTODSL A5 layer. | +| `MScatter` | `pending` | `-` | Memory scatter helper is not yet represented in the PTODSL A5 layer. | +| `TAdd` | `implemented` | `tadd` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TAddS` | `implemented` | `tadds` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TAlias` | `blocked` | `-` | C++ helper/meta header, not a tile micro-instruction kernel surface. | +| `TAnd` | `implemented` | `tand` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TAndS` | `implemented` | `tands` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TAssign` | `blocked` | `-` | C++ helper/meta header, not a tile micro-instruction kernel surface. | +| `TAxpy` | `partial` | `taxpy` | Same-dtype vector-micro path is implemented via vmula; the C++ mixed f32<-f16 variant is still missing. | +| `TBinOp` | `implemented` | `tbinary._binary_tile_vop` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TBinSOp` | `implemented` | `tscalar._scalar_tile_vop` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TCmp` | `native` | `compare` | Still forwarded to the native PTO tile builder; packed predicate tile lowering is not rewritten yet. | +| `TCmps` | `native` | `compare` | Still forwarded to the native PTO tile builder; scalar compare packing is not rewritten yet. | +| `TColArgMax` | `pending` | `-` | Arg-reduction micro lowering is not implemented yet. | +| `TColArgMin` | `pending` | `-` | Arg-reduction micro lowering is not implemented yet. | +| `TColExpand` | `implemented` | `tcol_expand` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TColExpandAdd` | `implemented` | `tcol_expand_add` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TColExpandBinOp` | `pending` | `-` | Generic binary broadcast frontend is not exposed yet. | +| `TColExpandDiv` | `implemented` | `tcol_expand_div` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TColExpandExpdif` | `pending` | `-` | Specialized exp-diff broadcast lowering is not implemented yet. | +| `TColExpandMax` | `implemented` | `tcol_expand_max` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TColExpandMin` | `implemented` | `tcol_expand_min` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TColExpandMul` | `implemented` | `tcol_expand_mul` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TColExpandSub` | `implemented` | `tcol_expand_sub` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TColMax` | `implemented` | `tcol_max` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TColMin` | `implemented` | `tcol_min` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TColProd` | `blocked` | `-` | No column-product micro lowering is wired yet. | +| `TColReduceIdx` | `pending` | `-` | Indexed column reduction is not implemented yet. | +| `TColReduceOps` | `implemented` | `treduce._tcol_reduce` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TColSum` | `implemented` | `tcol_sum` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TConcat` | `native` | `concat` | Still forwarded to the native PTO tile builder, not rewritten to micro ops yet. | +| `TCvt` | `pending` | `-` | Tile conversion helper is not implemented in the A5 micro layer yet. | +| `TDeQuant` | `pending` | `-` | Quantization/dequantization path is not implemented yet. | +| `TDiv` | `implemented` | `tdiv` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TDivS` | `implemented` | `tdivs` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TExpandS` | `pending` | `-` | Scalar expand helper is not implemented yet. | +| `TExtract` | `native` | `extract` | Still forwarded to the native PTO tile builder. | +| `TFMod` | `pending` | `-` | Fmod lowering is not implemented yet. | +| `TFModS` | `pending` | `-` | Scalar fmod lowering is not implemented yet. | +| `TFillPad` | `pending` | `-` | Pad/fill helper is not implemented yet. | +| `TGather` | `partial` | `tgather` | Indexed gather is implemented via vgather2; mask-pattern gather still needs missing vsqz-style micro support. | +| `TGatherB` | `pending` | `-` | GatherB lowering is not implemented yet, even though vgatherb exists in the micro surface. | +| `TGetScaleAddr` | `pending` | `-` | Scale-address helper is not represented in the PTODSL A5 layer. | +| `THistogram` | `pending` | `-` | Histogram lowering is not implemented yet. | +| `TImg2col` | `blocked` | `-` | Hardware layout/state programming path, not a straightforward vector-micro rewrite target. | +| `TInsert` | `native` | `insert` | Still forwarded to the native PTO tile builder. | +| `TLRelu` | `implemented` | `tlrelu` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TLoad` | `native` | `load_tile` | Structural staging helper, not a compute rewrite target. | +| `TMatmul` | `blocked` | `-` | Cube/L0 path is not a pure vector-micro rewrite target. | +| `TMax` | `implemented` | `tmax` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TMaxs` | `implemented` | `tmaxs` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TMin` | `implemented` | `tmin` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TMins` | `implemented` | `tmins` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TMov` | `implemented` | `tmov` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TMrgSort` | `implemented` | `tmrgsort` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TMul` | `implemented` | `tmul` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TMulS` | `implemented` | `tmuls` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TOr` | `implemented` | `tor_` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TOrS` | `implemented` | `tors` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TPack` | `pending` | `-` | Pack lowering is not implemented yet. | +| `TPartAdd` | `pending` | `-` | Part-op lowering is not implemented yet. | +| `TPartBinOps` | `pending` | `-` | Part-op lowering is not implemented yet. | +| `TPartMax` | `pending` | `-` | Part-op lowering is not implemented yet. | +| `TPartMin` | `pending` | `-` | Part-op lowering is not implemented yet. | +| `TPartMul` | `pending` | `-` | Part-op lowering is not implemented yet. | +| `TPop` | `blocked` | `-` | Runtime buffer stack/state helper, not a direct vector tile rewrite target. | +| `TPrefetch` | `blocked` | `-` | Prefetch/runtime helper, not a direct vector tile rewrite target. | +| `TPrelu` | `pending` | `-` | PReLU lowering is not implemented yet. | +| `TPrint` | `native` | `native print` | Still forwarded to the native PTO tile builder. | +| `TPush` | `blocked` | `-` | Runtime buffer stack/state helper, not a direct vector tile rewrite target. | +| `TQuant` | `pending` | `-` | Quantization path is not implemented yet. | +| `TRandom` | `pending` | `-` | Random-number helper is not implemented yet. | +| `TRem` | `pending` | `-` | Remainder lowering is not implemented yet. | +| `TRemS` | `pending` | `-` | Scalar remainder lowering is not implemented yet. | +| `TReshape` | `native` | `native reshape` | View/layout helper, not rewritten in the A5 micro layer. | +| `TRowExpand` | `implemented` | `trow_expand` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TRowExpandAdd` | `implemented` | `trow_expand_add` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TRowExpandBinOp` | `pending` | `-` | Generic row-broadcast binary frontend is not exposed yet. | +| `TRowExpandDiv` | `implemented` | `trow_expand_div` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TRowExpandExpdif` | `pending` | `-` | Specialized exp-diff row-broadcast lowering is not implemented yet. | +| `TRowExpandMax` | `implemented` | `trow_expand_max` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TRowExpandMin` | `implemented` | `trow_expand_min` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TRowExpandMul` | `implemented` | `trow_expand_mul` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TRowExpandSub` | `implemented` | `trow_expand_sub` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TRowProd` | `blocked` | `-` | No row-product micro lowering is wired yet. | +| `TRowReduce` | `implemented` | `treduce._trow_reduce` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TRowReduceIdx` | `pending` | `-` | Indexed row reduction is not implemented yet. | +| `TRsqrt` | `implemented` | `trsqrt` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TScatter` | `pending` | `-` | Scatter lowering is not implemented yet, even though vscatter exists in the micro surface. | +| `TSel` | `pending` | `-` | Packed-mask select lowering is not implemented yet. | +| `TSels` | `pending` | `-` | Scalar/mask select lowering is not implemented yet. | +| `TSetFmatrix` | `blocked` | `-` | Hardware state setup header, not a straightforward vector-micro rewrite target. | +| `TSetImg2colPadding` | `blocked` | `-` | Hardware state setup header, not a straightforward vector-micro rewrite target. | +| `TSetImg2colRpt` | `blocked` | `-` | Hardware state setup header, not a straightforward vector-micro rewrite target. | +| `TShl` | `implemented` | `tshl` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TShlS` | `implemented` | `tshls` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TShr` | `implemented` | `tshr` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TShrS` | `implemented` | `tshrs` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TSort32` | `implemented` | `tsort32` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TStore` | `native` | `store_tile` | Structural staging helper, not a compute rewrite target. | +| `TSub` | `implemented` | `tsub` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TSubS` | `implemented` | `tsubs` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TSubView` | `blocked` | `-` | View helper, not a tile compute rewrite target. | +| `TSync` | `pending` | `-` | Synchronization helper is not represented in the A5 library layer yet. | +| `TTrans` | `native` | `trans` | Still forwarded to the native PTO tile builder. | +| `TTri` | `pending` | `-` | Triangular helper is not implemented yet. | +| `TUnaryOp` | `implemented` | `tunary._unary_tile_vop` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TXor` | `implemented` | `txor` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `TXorS` | `implemented` | `txors` | Rewritten with PTO micro instructions in the PTODSL A5 layer. | +| `Tci` | `native` | `native tci` | Still forwarded to the native PTO tile builder. | +| `common` | `blocked` | `-` | A5 shared infrastructure header. | +| `custom/Div754` | `blocked` | `-` | Custom implementation helper header. | +| `custom/TSyncCVID` | `blocked` | `-` | Custom sync helper header. | +| `custom/TSync_Custom` | `blocked` | `-` | Custom sync helper header. | +| `datatype` | `blocked` | `-` | A5 shared datatype infrastructure header. | +| `utils` | `blocked` | `-` | A5 shared utility infrastructure header. | diff --git a/ptodsl/lib/a5/README.md b/ptodsl/lib/a5/README.md index 204bb8e7..1a824e3b 100644 --- a/ptodsl/lib/a5/README.md +++ b/ptodsl/lib/a5/README.md @@ -8,22 +8,29 @@ The scope of this layout is: - Small, readable files that show how a tile helper is written from PTO micro opcodes such as `pto.vlds`, `pto.vadd`, and `pto.vsts` -- A5-flavored aliases such as `TLoad`, `TAdd`, `TMatmul`, and `TStore` +- Canonical tile helper names such as `tadd`, `trow_sum`, and `tgather`, plus + matching A5 header aliases such as `TAdd` where parity with `pto-isa` matters - Example builder kernels that emit `.pto` through PTODSL -- A checked-in generation flow for reproducible `.pto` artifacts +- A checked-in generation flow for reproducible `.pto` artifacts and HIVM LLVM + sidecars for pure micro kernels Entry points: - [`tbinary.py`](./tbinary.py): tile binary helpers such as `tadd`, `tsub`, `tmul`, - `tdiv`, and `tor_`, written with PTO vector micro ops + `tdiv`, `tmax`, `tmin`, and `tor_`, written with PTO vector micro ops +- [`tscalar.py`](./tscalar.py): scalar tile helpers such as `tadds`, `tmaxs`, + `tlrelu`, and `taxpy` - [`tunary.py`](./tunary.py): tile unary helpers such as `texp`, `tlog`, `trelu`, `tsqrt`, `trsqrt`, and `trecip` -- [`texpand.py`](./texpand.py): row and column broadcast helpers +- [`texpand.py`](./texpand.py): row and column broadcast helpers, including add/sub/div/mul/max/min variants - [`treduce.py`](./treduce.py): row and column reduction helpers - [`tsort.py`](./tsort.py): gather and sort helpers - [`native.py`](./native.py): helpers that still map directly to tile/cube ops - [`ops.py`](./ops.py): the public A5 surface that re-exports the split helpers -- [`kernels.py`](./kernels.py): translated example kernels +- [`a5_header_coverage.py`](./a5_header_coverage.py): tracked status for the wider A5 header inventory +- [`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 - [`generated`](./generated): emitted `.pto` artifacts from `scripts/generate_a5_pto.py` Regenerate the current artifacts with: @@ -33,5 +40,17 @@ 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-cpp` is best-effort: the tile-based kernels lower through local `ptoas`, -while the direct micro-only kernel currently remains `.pto`-only in this environment. +To also emit HIVM LLVM for the pure micro kernels: + +```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_a5_pto.py --emit-hivm-llvm +``` + +`--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.*` +- `--emit-hivm-llvm` only writes `.ll` sidecars for kernels listed in + `a5.HIVM_LLVM_KERNELS` +- tensor-view and tile-buffer frontend examples remain useful PTODSL coverage, + but they are not yet the canonical PTOAS VPTO/HIVM path diff --git a/ptodsl/lib/a5/__init__.py b/ptodsl/lib/a5/__init__.py index 3b273e59..87a72335 100644 --- a/ptodsl/lib/a5/__init__.py +++ b/ptodsl/lib/a5/__init__.py @@ -1,8 +1,11 @@ -from . import native, ops, tbinary, texpand, treduce, tsort, tunary +from . import native, ops, tbinary, texpand, treduce, tscalar, tsort, tunary +from .a5_header_coverage import A5_HEADER_COVERAGE, a5_header_coverage_markdown from .kernels import ( + HIVM_LLVM_KERNELS, KERNEL_BUILDERS, build_cube_matmul, build_elementwise_add, + build_hivm_vadd_demo, build_mxfp8_matmul, build_templated_elementwise_add, build_vector_copy, @@ -15,10 +18,14 @@ ) __all__ = list(ops.__all__) + [ + "A5_HEADER_COVERAGE", + "HIVM_LLVM_KERNELS", "KERNEL_BUILDERS", "TILE_MICRO_COVERAGE", + "a5_header_coverage_markdown", "build_cube_matmul", "build_elementwise_add", + "build_hivm_vadd_demo", "build_mxfp8_matmul", "build_templated_elementwise_add", "build_vector_copy", @@ -28,6 +35,7 @@ "tbinary", "texpand", "treduce", + "tscalar", "tsort", "tunary", ] diff --git a/ptodsl/lib/a5/_common.py b/ptodsl/lib/a5/_common.py index 9114639f..de80c57b 100644 --- a/ptodsl/lib/a5/_common.py +++ b/ptodsl/lib/a5/_common.py @@ -4,7 +4,7 @@ import re from mlir.dialects import arith, pto -from mlir.ir import IntegerAttr, IntegerType +from mlir.ir import IndexType, IntegerAttr, IntegerType from ... import const_expr, language as dsl, range_constexpr, scalar as s from ...api.scalar import _unwrap @@ -27,6 +27,17 @@ "i8": {"i8", "int8"}, } +_ADDRESS_SPACE_BY_NAME = { + "ACC": pto.AddressSpace.ACC, + "BIAS": pto.AddressSpace.BIAS, + "GM": pto.AddressSpace.GM, + "LEFT": pto.AddressSpace.LEFT, + "MAT": pto.AddressSpace.MAT, + "RIGHT": pto.AddressSpace.RIGHT, + "SCALING": pto.AddressSpace.SCALING, + "VEC": pto.AddressSpace.VEC, +} + def _call(op, *args, **kwargs): return op( @@ -40,7 +51,10 @@ def raw(value): def _space_enum(space): - return getattr(pto.AddressSpace, str(space).upper()) + normalized = str(space).upper() + if normalized not in _ADDRESS_SPACE_BY_NAME: + raise ValueError(f"Unsupported address space '{space}'.") + return _ADDRESS_SPACE_BY_NAME[normalized] def ptr(dtype, *, space="GM"): @@ -73,10 +87,21 @@ def const_i32(value): return arith.ConstantOp(i32, IntegerAttr.get(i32, value)).result +def const_index(value): + return arith.ConstantOp(IndexType.get(), value).result + + def const_float(dtype, value): return arith.ConstantOp(dtype, value).result +def const_scalar(dtype, value): + token = dtype_token(dtype) + if token.startswith("f") or token == "bf16": + return arith.ConstantOp(dtype, float(value)).result + return arith.ConstantOp(dtype, int(value)).result + + def row_major_strides(shape): strides = [None] * len(shape) stride = s.const(1) @@ -91,6 +116,26 @@ def _index_value(value): return s.const(value) if isinstance(value, int) else value +def _to_index(value): + if isinstance(value, int): + return const_index(value) + value = _unwrap(value) + value_type = value.type if hasattr(value, "type") else "" + if str(value_type) == "index": + return value + return arith.IndexCastOp(IndexType.get(), value).result + + +def _to_i32(value): + if isinstance(value, int): + return const_i32(value) + value = _unwrap(value) + value_type = value.type if hasattr(value, "type") else "" + if str(value_type) == "index": + return arith.IndexCastOp(IntegerType.get_signless(32), value).result + return value + + def make_tensor(ptr_value, *, shape, dtype): tensor_type = dsl.TensorType(rank=len(shape), dtype=dtype) return dsl.as_tensor( @@ -113,7 +158,7 @@ def slice_tensor(source, *, offsets, sizes, dtype): def alloc_tile_buffer( dtype, - shape, + tile_shape, *, space="VEC", valid_shape=None, @@ -122,20 +167,32 @@ def alloc_tile_buffer( valid_row=None, valid_col=None, ): + if valid_shape is None: + valid_shape = list(tile_shape) tile_type = dsl.TileBufType( - shape=shape, + shape=tile_shape, dtype=dtype, memory_space=space, valid_shape=valid_shape, config=config, ) kwargs = {} + row_is_dynamic = ( + isinstance(valid_shape, (list, tuple)) + and len(valid_shape) > 0 + and valid_shape[0] == -1 + ) + col_is_dynamic = ( + isinstance(valid_shape, (list, tuple)) + and len(valid_shape) > 1 + and valid_shape[1] == -1 + ) if addr is not None: kwargs["addr"] = _unwrap(addr) - if valid_row is not None: - kwargs["valid_row"] = _unwrap(valid_row) - if valid_col is not None: - kwargs["valid_col"] = _unwrap(valid_col) + if valid_row is not None and row_is_dynamic: + kwargs["valid_row"] = _to_index(valid_row) + if valid_col is not None and col_is_dynamic: + kwargs["valid_col"] = _to_index(valid_col) return pto.AllocTileOp(tile_type, **kwargs).result @@ -159,22 +216,29 @@ def load_tile( tile_buffer=None, *, dtype=None, + tile_shape=None, shape=None, space="VEC", valid_shape=None, config=None, + valid_row=None, + valid_col=None, ): if tile_buffer is None: - if dtype is None or shape is None: + if tile_shape is None: + tile_shape = shape + if dtype is None or tile_shape is None: raise ValueError( - "`load_tile(...)` requires either `tile_buffer=` or both `dtype=` and `shape=`." + "`load_tile(...)` requires either `tile_buffer=` or both `dtype=` and `tile_shape=`." ) tile_buffer = alloc_tile_buffer( dtype, - shape, + tile_shape, space=space, valid_shape=valid_shape, config=config, + valid_row=valid_row, + valid_col=valid_col, ) load_view(view, tile_buffer) return tile_buffer @@ -214,7 +278,7 @@ def resolve_lanes(dtype, lanes): def extract_static_tensor_shape(value): raw = _unwrap(value) - type_obj = getattr(raw, "type", None) + type_obj = raw.type if hasattr(raw, "type") else None if type_obj is None: return None text = str(type_obj) @@ -238,7 +302,7 @@ def extract_static_tensor_shape(value): def extract_tensor_dtype_token(value): raw = _unwrap(value) - type_obj = getattr(raw, "type", None) + type_obj = raw.type if hasattr(raw, "type") else None if type_obj is None: return None text = str(type_obj).lower() @@ -283,6 +347,32 @@ def require_static_matrix_shape(shape, *, context): return rows, cols +def resolve_tile_spec( + *, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + context, +): + if tile_shape is None: + tile_shape = shape + if tile_shape is None: + raise ValueError(f"{context} requires `tile_shape=`.") + rows, cols = require_static_matrix_shape(tile_shape, context=context) + if valid_row is None: + valid_row = rows + if valid_col is None: + valid_col = cols + if valid_shape is None: + valid_shape = [ + valid_row if isinstance(valid_row, int) else -1, + valid_col if isinstance(valid_col, int) else -1, + ] + return rows, cols, valid_row, valid_col, valid_shape + + def full_mask(dtype): width = dtype_byte_width(dtype) if width == 4: @@ -296,7 +386,7 @@ def full_mask(dtype): def tail_mask(dtype, active_lanes): i32 = IntegerType.get_signless(32) - active = const_i32(active_lanes) + active = _to_i32(active_lanes) width = dtype_byte_width(dtype) if width == 4: mask, _ = pto.plt_b32(mask_type(), i32, active) @@ -312,11 +402,51 @@ def tail_mask(dtype, active_lanes): def mask_for_chunk(dtype, active_lanes): lanes = micro_lane_count(dtype) - if active_lanes == lanes: + if isinstance(active_lanes, int) and active_lanes == lanes: return full_mask(dtype) return tail_mask(dtype, active_lanes) +def clamp_active_lanes(active_lanes, lanes): + if isinstance(active_lanes, int): + return max(0, min(active_lanes, lanes)) + active_i32 = _to_i32(active_lanes) + zero = const_i32(0) + lanes_i32 = const_i32(lanes) + is_positive = arith.CmpIOp(arith.CmpIPredicate.sgt, active_i32, zero).result + nonnegative = arith.SelectOp(is_positive, active_i32, zero).result + fits = arith.CmpIOp(arith.CmpIPredicate.slt, nonnegative, lanes_i32).result + return arith.SelectOp(fits, nonnegative, lanes_i32).result + + +def flat_active_lanes(valid_row, valid_col, offset, lanes): + if isinstance(valid_row, int) and isinstance(valid_col, int): + return clamp_active_lanes(valid_row * valid_col - offset, lanes) + total = arith.MulIOp(_to_index(valid_row), _to_index(valid_col)).result + remaining = arith.SubIOp(total, const_index(offset)).result + return clamp_active_lanes(remaining, lanes) + + +def matrix_active_lanes(valid_row, valid_col, row, col, lanes): + if ( + isinstance(valid_row, int) + and isinstance(valid_col, int) + and isinstance(row, int) + and isinstance(col, int) + ): + if row >= valid_row: + return 0 + return clamp_active_lanes(valid_col - col, lanes) + + valid_col_idx = _to_index(valid_col) + remaining = arith.SubIOp(valid_col_idx, const_index(col)).result + active = clamp_active_lanes(remaining, lanes) + row_valid = arith.CmpIOp( + arith.CmpIPredicate.slt, const_index(row), _to_index(valid_row) + ).result + return arith.SelectOp(row_valid, _to_i32(active), const_i32(0)).result + + def onept_dist(dtype): width = dtype_byte_width(dtype) if width == 4: @@ -353,6 +483,28 @@ def normalize_vf_impl_kind(impl): return aliases[normalized] +def check_tscalar_operands(src_view, out_view, *, dtype, shape, context, allowed=None): + rows, cols = require_static_matrix_shape(shape, context=context) + require_supported_dtype( + dtype, + allowed=allowed + or {"f32", "f16", "bf16", "i32", "u32", "i16", "u16", "i8", "u8"}, + message=f"Fix: {context} has invalid data type.", + ) + for view, label in ((src_view, "src"), (out_view, "dst")): + require_view_shape( + view, + [rows, cols], + message=f"Fix: {context} input tile {label} valid shape mismatch.", + ) + require_view_dtype( + view, + dtype, + message=f"Fix: {context} input and output tile data type mismatch.", + ) + return rows, cols + + def check_tbinop_operands(lhs_view, rhs_view, out_view, *, dtype, shape, context): rows, cols = require_static_matrix_shape(shape, context=context) require_supported_dtype( @@ -591,6 +743,7 @@ def check_sort32_operands(src_view, idx_view, out_view, *, dtype, shape): "_call", "align_type", "alloc_tile_buffer", + "clamp_active_lanes", "check_col_expand_operands", "check_col_reduce_operands", "check_gather_operands", @@ -598,17 +751,22 @@ def check_sort32_operands(src_view, idx_view, out_view, *, dtype, shape): "check_row_expand_operands", "check_row_reduce_operands", "check_sort32_operands", + "check_tscalar_operands", "check_tbinop_operands", "const_expr", "const_float", + "const_scalar", + "const_index", "const_i32", "const_i64", "dtype_byte_width", "dtype_token", + "flat_active_lanes", "full_mask", "load_tile", "load_view", "make_tensor", + "matrix_active_lanes", "mask_for_chunk", "mask_type", "micro_lane_count", @@ -617,6 +775,7 @@ def check_sort32_operands(src_view, idx_view, out_view, *, dtype, shape): "onept_dist", "ptr", "range_constexpr", + "resolve_tile_spec", "resolve_lanes", "s", "slice_tensor", diff --git a/ptodsl/lib/a5/a5_header_coverage.py b/ptodsl/lib/a5/a5_header_coverage.py new file mode 100644 index 00000000..c20570b3 --- /dev/null +++ b/ptodsl/lib/a5/a5_header_coverage.py @@ -0,0 +1,335 @@ +"""Coverage view over the broader A5 header surface from pto-isa/include/pto/npu/a5.""" + +A5_HEADER_INVENTORY = [ + "MGather", + "MScatter", + "TAdd", + "TAddS", + "TAlias", + "TAnd", + "TAndS", + "TAssign", + "TAxpy", + "TBinOp", + "TBinSOp", + "TCmp", + "TCmps", + "TColArgMax", + "TColArgMin", + "TColExpand", + "TColExpandAdd", + "TColExpandBinOp", + "TColExpandDiv", + "TColExpandExpdif", + "TColExpandMax", + "TColExpandMin", + "TColExpandMul", + "TColExpandSub", + "TColMax", + "TColMin", + "TColProd", + "TColReduceIdx", + "TColReduceOps", + "TColSum", + "TConcat", + "TCvt", + "TDeQuant", + "TDiv", + "TDivS", + "TExpandS", + "TExtract", + "TFMod", + "TFModS", + "TFillPad", + "TGather", + "TGatherB", + "TGetScaleAddr", + "THistogram", + "TImg2col", + "TInsert", + "TLRelu", + "TLoad", + "TMatmul", + "TMax", + "TMaxs", + "TMin", + "TMins", + "TMov", + "TMrgSort", + "TMul", + "TMulS", + "TOr", + "TOrS", + "TPack", + "TPartAdd", + "TPartBinOps", + "TPartMax", + "TPartMin", + "TPartMul", + "TPop", + "TPrefetch", + "TPrelu", + "TPrint", + "TPush", + "TQuant", + "TRandom", + "TRem", + "TRemS", + "TReshape", + "TRowExpand", + "TRowExpandAdd", + "TRowExpandBinOp", + "TRowExpandDiv", + "TRowExpandExpdif", + "TRowExpandMax", + "TRowExpandMin", + "TRowExpandMul", + "TRowExpandSub", + "TRowProd", + "TRowReduce", + "TRowReduceIdx", + "TRsqrt", + "TScatter", + "TSel", + "TSels", + "TSetFmatrix", + "TSetImg2colPadding", + "TSetImg2colRpt", + "TShl", + "TShlS", + "TShr", + "TShrS", + "TSort32", + "TStore", + "TSub", + "TSubS", + "TSubView", + "TSync", + "TTrans", + "TTri", + "TUnaryOp", + "TXor", + "TXorS", + "Tci", + "common", + "custom/Div754", + "custom/TSyncCVID", + "custom/TSync_Custom", + "datatype", + "utils", +] + +_MICRO_HELPERS = { + "TAdd": "tadd", + "TAddS": "tadds", + "TAnd": "tand", + "TAndS": "tands", + "TBinOp": "tbinary._binary_tile_vop", + "TBinSOp": "tscalar._scalar_tile_vop", + "TColExpand": "tcol_expand", + "TColExpandAdd": "tcol_expand_add", + "TColExpandDiv": "tcol_expand_div", + "TColExpandMax": "tcol_expand_max", + "TColExpandMin": "tcol_expand_min", + "TColExpandMul": "tcol_expand_mul", + "TColExpandSub": "tcol_expand_sub", + "TColMax": "tcol_max", + "TColMin": "tcol_min", + "TColReduceOps": "treduce._tcol_reduce", + "TColSum": "tcol_sum", + "TDiv": "tdiv", + "TDivS": "tdivs", + "TLRelu": "tlrelu", + "TMax": "tmax", + "TMaxs": "tmaxs", + "TMin": "tmin", + "TMins": "tmins", + "TMov": "tmov", + "TMrgSort": "tmrgsort", + "TMul": "tmul", + "TMulS": "tmuls", + "TOr": "tor_", + "TOrS": "tors", + "TRowExpand": "trow_expand", + "TRowExpandAdd": "trow_expand_add", + "TRowExpandDiv": "trow_expand_div", + "TRowExpandMax": "trow_expand_max", + "TRowExpandMin": "trow_expand_min", + "TRowExpandMul": "trow_expand_mul", + "TRowExpandSub": "trow_expand_sub", + "TRowReduce": "treduce._trow_reduce", + "TRsqrt": "trsqrt", + "TShl": "tshl", + "TShlS": "tshls", + "TShr": "tshr", + "TShrS": "tshrs", + "TSort32": "tsort32", + "TSub": "tsub", + "TSubS": "tsubs", + "TUnaryOp": "tunary._unary_tile_vop", + "TXor": "txor", + "TXorS": "txors", +} + +_PARTIAL_HELPERS = { + "TAxpy": ( + "taxpy", + "Same-dtype vector-micro path is implemented via vmula; the C++ mixed f32<-f16 variant is still missing.", + ), + "TGather": ( + "tgather", + "Indexed gather is implemented via vgather2; mask-pattern gather still needs missing vsqz-style micro support.", + ), +} + +_NATIVE_HELPERS = { + "TConcat": ( + "concat", + "Still forwarded to the native PTO tile builder, not rewritten to micro ops yet.", + ), + "TCmp": ( + "compare", + "Still forwarded to the native PTO tile builder; packed predicate tile lowering is not rewritten yet.", + ), + "TCmps": ( + "compare", + "Still forwarded to the native PTO tile builder; scalar compare packing is not rewritten yet.", + ), + "TExtract": ("extract", "Still forwarded to the native PTO tile builder."), + "TInsert": ("insert", "Still forwarded to the native PTO tile builder."), + "TLoad": ("load_tile", "Structural staging helper, not a compute rewrite target."), + "TPrint": ("native print", "Still forwarded to the native PTO tile builder."), + "TReshape": ( + "native reshape", + "View/layout helper, not rewritten in the A5 micro layer.", + ), + "TStore": ( + "store_tile", + "Structural staging helper, not a compute rewrite target.", + ), + "TTrans": ("trans", "Still forwarded to the native PTO tile builder."), + "Tci": ("native tci", "Still forwarded to the native PTO tile builder."), +} + +_BLOCKED_HEADERS = { + "TAlias": "C++ helper/meta header, not a tile micro-instruction kernel surface.", + "TAssign": "C++ helper/meta header, not a tile micro-instruction kernel surface.", + "TColProd": "No column-product micro lowering is wired yet.", + "TImg2col": "Hardware layout/state programming path, not a straightforward vector-micro rewrite target.", + "TMatmul": "Cube/L0 path is not a pure vector-micro rewrite target.", + "TPop": "Runtime buffer stack/state helper, not a direct vector tile rewrite target.", + "TPrefetch": "Prefetch/runtime helper, not a direct vector tile rewrite target.", + "TPush": "Runtime buffer stack/state helper, not a direct vector tile rewrite target.", + "TRowProd": "No row-product micro lowering is wired yet.", + "TSetFmatrix": "Hardware state setup header, not a straightforward vector-micro rewrite target.", + "TSetImg2colPadding": "Hardware state setup header, not a straightforward vector-micro rewrite target.", + "TSetImg2colRpt": "Hardware state setup header, not a straightforward vector-micro rewrite target.", + "TSubView": "View helper, not a tile compute rewrite target.", + "common": "A5 shared infrastructure header.", + "custom/Div754": "Custom implementation helper header.", + "custom/TSyncCVID": "Custom sync helper header.", + "custom/TSync_Custom": "Custom sync helper header.", + "datatype": "A5 shared datatype infrastructure header.", + "utils": "A5 shared utility infrastructure header.", +} + +_PENDING_HEADERS = { + "MGather": "Memory gather helper is not yet represented in the PTODSL A5 layer.", + "MScatter": "Memory scatter helper is not yet represented in the PTODSL A5 layer.", + "TColArgMax": "Arg-reduction micro lowering is not implemented yet.", + "TColArgMin": "Arg-reduction micro lowering is not implemented yet.", + "TColExpandBinOp": "Generic binary broadcast frontend is not exposed yet.", + "TColExpandExpdif": "Specialized exp-diff broadcast lowering is not implemented yet.", + "TColReduceIdx": "Indexed column reduction is not implemented yet.", + "TCvt": "Tile conversion helper is not implemented in the A5 micro layer yet.", + "TDeQuant": "Quantization/dequantization path is not implemented yet.", + "TExpandS": "Scalar expand helper is not implemented yet.", + "TFMod": "Fmod lowering is not implemented yet.", + "TFModS": "Scalar fmod lowering is not implemented yet.", + "TFillPad": "Pad/fill helper is not implemented yet.", + "TGatherB": "GatherB lowering is not implemented yet, even though vgatherb exists in the micro surface.", + "TGetScaleAddr": "Scale-address helper is not represented in the PTODSL A5 layer.", + "THistogram": "Histogram lowering is not implemented yet.", + "TPack": "Pack lowering is not implemented yet.", + "TPartAdd": "Part-op lowering is not implemented yet.", + "TPartBinOps": "Part-op lowering is not implemented yet.", + "TPartMax": "Part-op lowering is not implemented yet.", + "TPartMin": "Part-op lowering is not implemented yet.", + "TPartMul": "Part-op lowering is not implemented yet.", + "TPrelu": "PReLU lowering is not implemented yet.", + "TQuant": "Quantization path is not implemented yet.", + "TRandom": "Random-number helper is not implemented yet.", + "TRem": "Remainder lowering is not implemented yet.", + "TRemS": "Scalar remainder lowering is not implemented yet.", + "TRowExpandBinOp": "Generic row-broadcast binary frontend is not exposed yet.", + "TRowExpandExpdif": "Specialized exp-diff row-broadcast lowering is not implemented yet.", + "TRowReduceIdx": "Indexed row reduction is not implemented yet.", + "TScatter": "Scatter lowering is not implemented yet, even though vscatter exists in the micro surface.", + "TSel": "Packed-mask select lowering is not implemented yet.", + "TSels": "Scalar/mask select lowering is not implemented yet.", + "TSync": "Synchronization helper is not represented in the A5 library layer yet.", + "TTri": "Triangular helper is not implemented yet.", +} + + +def _entry(name): + if name in _MICRO_HELPERS: + return { + "status": "implemented", + "helper": _MICRO_HELPERS[name], + "note": "Rewritten with PTO micro instructions in the PTODSL A5 layer.", + } + if name in _PARTIAL_HELPERS: + helper, note = _PARTIAL_HELPERS[name] + return {"status": "partial", "helper": helper, "note": note} + if name in _NATIVE_HELPERS: + helper, note = _NATIVE_HELPERS[name] + return {"status": "native", "helper": helper, "note": note} + if name in _BLOCKED_HEADERS: + return {"status": "blocked", "helper": None, "note": _BLOCKED_HEADERS[name]} + if name in _PENDING_HEADERS: + return {"status": "pending", "helper": None, "note": _PENDING_HEADERS[name]} + raise KeyError(f"Missing A5 header coverage classification for '{name}'.") + + +A5_HEADER_COVERAGE = {name: _entry(name) for name in A5_HEADER_INVENTORY} + + +def a5_header_coverage_summary(): + counts = {} + for entry in A5_HEADER_COVERAGE.values(): + status = entry["status"] + counts[status] = counts.get(status, 0) + 1 + return counts + + +def a5_header_coverage_markdown(): + counts = a5_header_coverage_summary() + lines = [ + "# A5 Header Coverage", + "", + f"- Total A5 headers tracked: `{len(A5_HEADER_INVENTORY)}`", + f"- Implemented: `{counts.get('implemented', 0)}`", + f"- Partial: `{counts.get('partial', 0)}`", + f"- Native only: `{counts.get('native', 0)}`", + f"- Pending: `{counts.get('pending', 0)}`", + f"- Blocked/meta: `{counts.get('blocked', 0)}`", + "", + "| header | status | helper | note |", + "| --- | --- | --- | --- |", + ] + for name in A5_HEADER_INVENTORY: + entry = A5_HEADER_COVERAGE[name] + helper = entry["helper"] or "-" + lines.append( + f"| `{name}` | `{entry['status']}` | `{helper}` | {entry['note']} |" + ) + return "\n".join(lines) + "\n" + + +__all__ = [ + "A5_HEADER_COVERAGE", + "A5_HEADER_INVENTORY", + "a5_header_coverage_markdown", + "a5_header_coverage_summary", +] diff --git a/ptodsl/lib/a5/generated/a5_elementwise_add.pto b/ptodsl/lib/a5/generated/a5_elementwise_add.pto index e25d712a..c47a56ec 100644 --- a/ptodsl/lib/a5/generated/a5_elementwise_add.pto +++ b/ptodsl/lib/a5/generated/a5_elementwise_add.pto @@ -39,230 +39,87 @@ module { %15 = pto.castptr %c0_i64 : i64 -> !pto.ptr %16 = pto.castptr %c4096_i64 : i64 -> !pto.ptr %17 = pto.castptr %c8192_i64 : i64 -> !pto.ptr - %c32_i32 = arith.constant 32 : i32 - %mask, %scalar_out = pto.plt_b32 %c32_i32 : i32 -> !pto.mask, i32 - %c0_12 = arith.constant 0 : index - %18 = pto.vlds %15[%c0_12] : !pto.ptr -> !pto.vreg<64xf32> - %19 = pto.vlds %16[%c0_12] : !pto.ptr -> !pto.vreg<64xf32> - %20 = pto.vadd %18, %19, %mask : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %20, %17[%c0_12], %mask : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_13 = arith.constant 32 : i32 - %mask_14, %scalar_out_15 = pto.plt_b32 %c32_i32_13 : i32 -> !pto.mask, i32 - %c32_16 = arith.constant 32 : index - %21 = pto.vlds %15[%c32_16] : !pto.ptr -> !pto.vreg<64xf32> - %22 = pto.vlds %16[%c32_16] : !pto.ptr -> !pto.vreg<64xf32> - %23 = pto.vadd %21, %22, %mask_14 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %23, %17[%c32_16], %mask_14 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_17 = arith.constant 32 : i32 - %mask_18, %scalar_out_19 = pto.plt_b32 %c32_i32_17 : i32 -> !pto.mask, i32 %c64 = arith.constant 64 : index - %24 = pto.vlds %15[%c64] : !pto.ptr -> !pto.vreg<64xf32> - %25 = pto.vlds %16[%c64] : !pto.ptr -> !pto.vreg<64xf32> - %26 = pto.vadd %24, %25, %mask_18 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %26, %17[%c64], %mask_18 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_20 = arith.constant 32 : i32 - %mask_21, %scalar_out_22 = pto.plt_b32 %c32_i32_20 : i32 -> !pto.mask, i32 - %c96 = arith.constant 96 : index - %27 = pto.vlds %15[%c96] : !pto.ptr -> !pto.vreg<64xf32> - %28 = pto.vlds %16[%c96] : !pto.ptr -> !pto.vreg<64xf32> - %29 = pto.vadd %27, %28, %mask_21 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %29, %17[%c96], %mask_21 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_23 = arith.constant 32 : i32 - %mask_24, %scalar_out_25 = pto.plt_b32 %c32_i32_23 : i32 -> !pto.mask, i32 - %c128 = arith.constant 128 : index - %30 = pto.vlds %15[%c128] : !pto.ptr -> !pto.vreg<64xf32> - %31 = pto.vlds %16[%c128] : !pto.ptr -> !pto.vreg<64xf32> - %32 = pto.vadd %30, %31, %mask_24 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %32, %17[%c128], %mask_24 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_26 = arith.constant 32 : i32 - %mask_27, %scalar_out_28 = pto.plt_b32 %c32_i32_26 : i32 -> !pto.mask, i32 - %c160 = arith.constant 160 : index - %33 = pto.vlds %15[%c160] : !pto.ptr -> !pto.vreg<64xf32> - %34 = pto.vlds %16[%c160] : !pto.ptr -> !pto.vreg<64xf32> - %35 = pto.vadd %33, %34, %mask_27 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %35, %17[%c160], %mask_27 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_29 = arith.constant 32 : i32 - %mask_30, %scalar_out_31 = pto.plt_b32 %c32_i32_29 : i32 -> !pto.mask, i32 - %c192 = arith.constant 192 : index - %36 = pto.vlds %15[%c192] : !pto.ptr -> !pto.vreg<64xf32> - %37 = pto.vlds %16[%c192] : !pto.ptr -> !pto.vreg<64xf32> - %38 = pto.vadd %36, %37, %mask_30 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %38, %17[%c192], %mask_30 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_32 = arith.constant 32 : i32 - %mask_33, %scalar_out_34 = pto.plt_b32 %c32_i32_32 : i32 -> !pto.mask, i32 - %c224 = arith.constant 224 : index - %39 = pto.vlds %15[%c224] : !pto.ptr -> !pto.vreg<64xf32> - %40 = pto.vlds %16[%c224] : !pto.ptr -> !pto.vreg<64xf32> - %41 = pto.vadd %39, %40, %mask_33 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %41, %17[%c224], %mask_33 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_35 = arith.constant 32 : i32 - %mask_36, %scalar_out_37 = pto.plt_b32 %c32_i32_35 : i32 -> !pto.mask, i32 - %c256 = arith.constant 256 : index - %42 = pto.vlds %15[%c256] : !pto.ptr -> !pto.vreg<64xf32> - %43 = pto.vlds %16[%c256] : !pto.ptr -> !pto.vreg<64xf32> - %44 = pto.vadd %42, %43, %mask_36 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %44, %17[%c256], %mask_36 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_38 = arith.constant 32 : i32 - %mask_39, %scalar_out_40 = pto.plt_b32 %c32_i32_38 : i32 -> !pto.mask, i32 - %c288 = arith.constant 288 : index - %45 = pto.vlds %15[%c288] : !pto.ptr -> !pto.vreg<64xf32> - %46 = pto.vlds %16[%c288] : !pto.ptr -> !pto.vreg<64xf32> - %47 = pto.vadd %45, %46, %mask_39 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %47, %17[%c288], %mask_39 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_41 = arith.constant 32 : i32 - %mask_42, %scalar_out_43 = pto.plt_b32 %c32_i32_41 : i32 -> !pto.mask, i32 - %c320 = arith.constant 320 : index - %48 = pto.vlds %15[%c320] : !pto.ptr -> !pto.vreg<64xf32> - %49 = pto.vlds %16[%c320] : !pto.ptr -> !pto.vreg<64xf32> - %50 = pto.vadd %48, %49, %mask_42 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %50, %17[%c320], %mask_42 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_44 = arith.constant 32 : i32 - %mask_45, %scalar_out_46 = pto.plt_b32 %c32_i32_44 : i32 -> !pto.mask, i32 - %c352 = arith.constant 352 : index - %51 = pto.vlds %15[%c352] : !pto.ptr -> !pto.vreg<64xf32> - %52 = pto.vlds %16[%c352] : !pto.ptr -> !pto.vreg<64xf32> - %53 = pto.vadd %51, %52, %mask_45 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %53, %17[%c352], %mask_45 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_47 = arith.constant 32 : i32 - %mask_48, %scalar_out_49 = pto.plt_b32 %c32_i32_47 : i32 -> !pto.mask, i32 - %c384 = arith.constant 384 : index - %54 = pto.vlds %15[%c384] : !pto.ptr -> !pto.vreg<64xf32> - %55 = pto.vlds %16[%c384] : !pto.ptr -> !pto.vreg<64xf32> - %56 = pto.vadd %54, %55, %mask_48 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %56, %17[%c384], %mask_48 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_50 = arith.constant 32 : i32 - %mask_51, %scalar_out_52 = pto.plt_b32 %c32_i32_50 : i32 -> !pto.mask, i32 - %c416 = arith.constant 416 : index - %57 = pto.vlds %15[%c416] : !pto.ptr -> !pto.vreg<64xf32> - %58 = pto.vlds %16[%c416] : !pto.ptr -> !pto.vreg<64xf32> - %59 = pto.vadd %57, %58, %mask_51 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %59, %17[%c416], %mask_51 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_53 = arith.constant 32 : i32 - %mask_54, %scalar_out_55 = pto.plt_b32 %c32_i32_53 : i32 -> !pto.mask, i32 - %c448 = arith.constant 448 : index - %60 = pto.vlds %15[%c448] : !pto.ptr -> !pto.vreg<64xf32> - %61 = pto.vlds %16[%c448] : !pto.ptr -> !pto.vreg<64xf32> - %62 = pto.vadd %60, %61, %mask_54 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %62, %17[%c448], %mask_54 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_56 = arith.constant 32 : i32 - %mask_57, %scalar_out_58 = pto.plt_b32 %c32_i32_56 : i32 -> !pto.mask, i32 - %c480 = arith.constant 480 : index - %63 = pto.vlds %15[%c480] : !pto.ptr -> !pto.vreg<64xf32> - %64 = pto.vlds %16[%c480] : !pto.ptr -> !pto.vreg<64xf32> - %65 = pto.vadd %63, %64, %mask_57 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %65, %17[%c480], %mask_57 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_59 = arith.constant 32 : i32 - %mask_60, %scalar_out_61 = pto.plt_b32 %c32_i32_59 : i32 -> !pto.mask, i32 - %c512 = arith.constant 512 : index - %66 = pto.vlds %15[%c512] : !pto.ptr -> !pto.vreg<64xf32> - %67 = pto.vlds %16[%c512] : !pto.ptr -> !pto.vreg<64xf32> - %68 = pto.vadd %66, %67, %mask_60 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %68, %17[%c512], %mask_60 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_62 = arith.constant 32 : i32 - %mask_63, %scalar_out_64 = pto.plt_b32 %c32_i32_62 : i32 -> !pto.mask, i32 - %c544 = arith.constant 544 : index - %69 = pto.vlds %15[%c544] : !pto.ptr -> !pto.vreg<64xf32> - %70 = pto.vlds %16[%c544] : !pto.ptr -> !pto.vreg<64xf32> - %71 = pto.vadd %69, %70, %mask_63 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %71, %17[%c544], %mask_63 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_65 = arith.constant 32 : i32 - %mask_66, %scalar_out_67 = pto.plt_b32 %c32_i32_65 : i32 -> !pto.mask, i32 - %c576 = arith.constant 576 : index - %72 = pto.vlds %15[%c576] : !pto.ptr -> !pto.vreg<64xf32> - %73 = pto.vlds %16[%c576] : !pto.ptr -> !pto.vreg<64xf32> - %74 = pto.vadd %72, %73, %mask_66 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %74, %17[%c576], %mask_66 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_68 = arith.constant 32 : i32 - %mask_69, %scalar_out_70 = pto.plt_b32 %c32_i32_68 : i32 -> !pto.mask, i32 - %c608 = arith.constant 608 : index - %75 = pto.vlds %15[%c608] : !pto.ptr -> !pto.vreg<64xf32> - %76 = pto.vlds %16[%c608] : !pto.ptr -> !pto.vreg<64xf32> - %77 = pto.vadd %75, %76, %mask_69 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %77, %17[%c608], %mask_69 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_71 = arith.constant 32 : i32 - %mask_72, %scalar_out_73 = pto.plt_b32 %c32_i32_71 : i32 -> !pto.mask, i32 - %c640 = arith.constant 640 : index - %78 = pto.vlds %15[%c640] : !pto.ptr -> !pto.vreg<64xf32> - %79 = pto.vlds %16[%c640] : !pto.ptr -> !pto.vreg<64xf32> - %80 = pto.vadd %78, %79, %mask_72 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %80, %17[%c640], %mask_72 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_74 = arith.constant 32 : i32 - %mask_75, %scalar_out_76 = pto.plt_b32 %c32_i32_74 : i32 -> !pto.mask, i32 - %c672 = arith.constant 672 : index - %81 = pto.vlds %15[%c672] : !pto.ptr -> !pto.vreg<64xf32> - %82 = pto.vlds %16[%c672] : !pto.ptr -> !pto.vreg<64xf32> - %83 = pto.vadd %81, %82, %mask_75 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %83, %17[%c672], %mask_75 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_77 = arith.constant 32 : i32 - %mask_78, %scalar_out_79 = pto.plt_b32 %c32_i32_77 : i32 -> !pto.mask, i32 - %c704 = arith.constant 704 : index - %84 = pto.vlds %15[%c704] : !pto.ptr -> !pto.vreg<64xf32> - %85 = pto.vlds %16[%c704] : !pto.ptr -> !pto.vreg<64xf32> - %86 = pto.vadd %84, %85, %mask_78 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %86, %17[%c704], %mask_78 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_80 = arith.constant 32 : i32 - %mask_81, %scalar_out_82 = pto.plt_b32 %c32_i32_80 : i32 -> !pto.mask, i32 - %c736 = arith.constant 736 : index - %87 = pto.vlds %15[%c736] : !pto.ptr -> !pto.vreg<64xf32> - %88 = pto.vlds %16[%c736] : !pto.ptr -> !pto.vreg<64xf32> - %89 = pto.vadd %87, %88, %mask_81 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %89, %17[%c736], %mask_81 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_83 = arith.constant 32 : i32 - %mask_84, %scalar_out_85 = pto.plt_b32 %c32_i32_83 : i32 -> !pto.mask, i32 - %c768 = arith.constant 768 : index - %90 = pto.vlds %15[%c768] : !pto.ptr -> !pto.vreg<64xf32> - %91 = pto.vlds %16[%c768] : !pto.ptr -> !pto.vreg<64xf32> - %92 = pto.vadd %90, %91, %mask_84 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %92, %17[%c768], %mask_84 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_86 = arith.constant 32 : i32 - %mask_87, %scalar_out_88 = pto.plt_b32 %c32_i32_86 : i32 -> !pto.mask, i32 - %c800 = arith.constant 800 : index - %93 = pto.vlds %15[%c800] : !pto.ptr -> !pto.vreg<64xf32> - %94 = pto.vlds %16[%c800] : !pto.ptr -> !pto.vreg<64xf32> - %95 = pto.vadd %93, %94, %mask_87 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %95, %17[%c800], %mask_87 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_89 = arith.constant 32 : i32 - %mask_90, %scalar_out_91 = pto.plt_b32 %c32_i32_89 : i32 -> !pto.mask, i32 - %c832 = arith.constant 832 : index - %96 = pto.vlds %15[%c832] : !pto.ptr -> !pto.vreg<64xf32> - %97 = pto.vlds %16[%c832] : !pto.ptr -> !pto.vreg<64xf32> - %98 = pto.vadd %96, %97, %mask_90 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %98, %17[%c832], %mask_90 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_92 = arith.constant 32 : i32 - %mask_93, %scalar_out_94 = pto.plt_b32 %c32_i32_92 : i32 -> !pto.mask, i32 - %c864 = arith.constant 864 : index - %99 = pto.vlds %15[%c864] : !pto.ptr -> !pto.vreg<64xf32> - %100 = pto.vlds %16[%c864] : !pto.ptr -> !pto.vreg<64xf32> - %101 = pto.vadd %99, %100, %mask_93 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %101, %17[%c864], %mask_93 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_95 = arith.constant 32 : i32 - %mask_96, %scalar_out_97 = pto.plt_b32 %c32_i32_95 : i32 -> !pto.mask, i32 - %c896 = arith.constant 896 : index - %102 = pto.vlds %15[%c896] : !pto.ptr -> !pto.vreg<64xf32> - %103 = pto.vlds %16[%c896] : !pto.ptr -> !pto.vreg<64xf32> - %104 = pto.vadd %102, %103, %mask_96 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %104, %17[%c896], %mask_96 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_98 = arith.constant 32 : i32 - %mask_99, %scalar_out_100 = pto.plt_b32 %c32_i32_98 : i32 -> !pto.mask, i32 - %c928 = arith.constant 928 : index - %105 = pto.vlds %15[%c928] : !pto.ptr -> !pto.vreg<64xf32> - %106 = pto.vlds %16[%c928] : !pto.ptr -> !pto.vreg<64xf32> - %107 = pto.vadd %105, %106, %mask_99 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %107, %17[%c928], %mask_99 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_101 = arith.constant 32 : i32 - %mask_102, %scalar_out_103 = pto.plt_b32 %c32_i32_101 : i32 -> !pto.mask, i32 - %c960 = arith.constant 960 : index - %108 = pto.vlds %15[%c960] : !pto.ptr -> !pto.vreg<64xf32> - %109 = pto.vlds %16[%c960] : !pto.ptr -> !pto.vreg<64xf32> - %110 = pto.vadd %108, %109, %mask_102 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %110, %17[%c960], %mask_102 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - %c32_i32_104 = arith.constant 32 : i32 - %mask_105, %scalar_out_106 = pto.plt_b32 %c32_i32_104 : i32 -> !pto.mask, i32 - %c992 = arith.constant 992 : index - %111 = pto.vlds %15[%c992] : !pto.ptr -> !pto.vreg<64xf32> - %112 = pto.vlds %16[%c992] : !pto.ptr -> !pto.vreg<64xf32> - %113 = pto.vadd %111, %112, %mask_105 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - pto.vsts %113, %17[%c992], %mask_105 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + %18 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result, %updated_source = pto.vlds_post %15[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_12, %updated_source_13 = pto.vlds_post %16[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %19 = pto.vadd %result, %result_12, %18 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %20 = pto.vsts_post %19, %17[%c64], %18 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %21 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_14, %updated_source_15 = pto.vlds_post %updated_source[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_16, %updated_source_17 = pto.vlds_post %updated_source_13[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %22 = pto.vadd %result_14, %result_16, %21 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %23 = pto.vsts_post %22, %20[%c64], %21 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %24 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_18, %updated_source_19 = pto.vlds_post %updated_source_15[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_20, %updated_source_21 = pto.vlds_post %updated_source_17[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %25 = pto.vadd %result_18, %result_20, %24 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %26 = pto.vsts_post %25, %23[%c64], %24 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %27 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_22, %updated_source_23 = pto.vlds_post %updated_source_19[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_24, %updated_source_25 = pto.vlds_post %updated_source_21[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %28 = pto.vadd %result_22, %result_24, %27 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %29 = pto.vsts_post %28, %26[%c64], %27 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %30 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_26, %updated_source_27 = pto.vlds_post %updated_source_23[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_28, %updated_source_29 = pto.vlds_post %updated_source_25[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %31 = pto.vadd %result_26, %result_28, %30 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %32 = pto.vsts_post %31, %29[%c64], %30 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %33 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_30, %updated_source_31 = pto.vlds_post %updated_source_27[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_32, %updated_source_33 = pto.vlds_post %updated_source_29[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %34 = pto.vadd %result_30, %result_32, %33 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %35 = pto.vsts_post %34, %32[%c64], %33 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %36 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_34, %updated_source_35 = pto.vlds_post %updated_source_31[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_36, %updated_source_37 = pto.vlds_post %updated_source_33[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %37 = pto.vadd %result_34, %result_36, %36 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %38 = pto.vsts_post %37, %35[%c64], %36 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %39 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_38, %updated_source_39 = pto.vlds_post %updated_source_35[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_40, %updated_source_41 = pto.vlds_post %updated_source_37[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %40 = pto.vadd %result_38, %result_40, %39 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %41 = pto.vsts_post %40, %38[%c64], %39 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %42 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_42, %updated_source_43 = pto.vlds_post %updated_source_39[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_44, %updated_source_45 = pto.vlds_post %updated_source_41[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %43 = pto.vadd %result_42, %result_44, %42 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %44 = pto.vsts_post %43, %41[%c64], %42 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %45 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_46, %updated_source_47 = pto.vlds_post %updated_source_43[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_48, %updated_source_49 = pto.vlds_post %updated_source_45[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %46 = pto.vadd %result_46, %result_48, %45 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %47 = pto.vsts_post %46, %44[%c64], %45 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %48 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_50, %updated_source_51 = pto.vlds_post %updated_source_47[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_52, %updated_source_53 = pto.vlds_post %updated_source_49[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %49 = pto.vadd %result_50, %result_52, %48 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %50 = pto.vsts_post %49, %47[%c64], %48 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %51 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_54, %updated_source_55 = pto.vlds_post %updated_source_51[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_56, %updated_source_57 = pto.vlds_post %updated_source_53[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %52 = pto.vadd %result_54, %result_56, %51 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %53 = pto.vsts_post %52, %50[%c64], %51 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %54 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_58, %updated_source_59 = pto.vlds_post %updated_source_55[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_60, %updated_source_61 = pto.vlds_post %updated_source_57[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %55 = pto.vadd %result_58, %result_60, %54 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %56 = pto.vsts_post %55, %53[%c64], %54 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %57 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_62, %updated_source_63 = pto.vlds_post %updated_source_59[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_64, %updated_source_65 = pto.vlds_post %updated_source_61[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %58 = pto.vadd %result_62, %result_64, %57 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %59 = pto.vsts_post %58, %56[%c64], %57 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %60 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_66, %updated_source_67 = pto.vlds_post %updated_source_63[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_68, %updated_source_69 = pto.vlds_post %updated_source_65[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %61 = pto.vadd %result_66, %result_68, %60 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %62 = pto.vsts_post %61, %59[%c64], %60 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr + %63 = pto.pset_b32 "PAT_ALL" : !pto.mask + %result_70, %updated_source_71 = pto.vlds_post %updated_source_67[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %result_72, %updated_source_73 = pto.vlds_post %updated_source_69[%c64] : !pto.ptr -> !pto.vreg<64xf32>, !pto.ptr + %64 = pto.vadd %result_70, %result_72, %63 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + %65 = pto.vsts_post %64, %62[%c64], %63 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask -> !pto.ptr pto.tstore ins(%14 : !pto.tile_buf) outs(%11 : !pto.partition_tensor_view<32x32xf32>) } return diff --git a/ptodsl/lib/a5/generated/a5_hivm_vadd_demo.ll b/ptodsl/lib/a5/generated/a5_hivm_vadd_demo.ll new file mode 100644 index 00000000..e18a70aa --- /dev/null +++ b/ptodsl/lib/a5/generated/a5_hivm_vadd_demo.ll @@ -0,0 +1,36 @@ +; ModuleID = 'ptoas.hivm.official' +source_filename = "ptoas.hivm.official" + +declare void @llvm.hivm.vstsx1.v64f32(<64 x float>, ptr addrspace(6), i32, i32, i32, <256 x i1>) + +declare <64 x float> @llvm.hivm.vadd.v64f32.x(<64 x float>, <64 x float>, <256 x i1>) + +declare <64 x float> @llvm.hivm.vldsx1.v64f32(ptr addrspace(6), i32, i32, i32) + +declare { <256 x i1>, i32 } @llvm.hivm.plt.b32.v300(i32) + +define void @a5_hivm_vadd_demo(ptr addrspace(6) %0, ptr addrspace(6) %1, ptr addrspace(6) %2, i64 %3) #0 { + %5 = call { <256 x i1>, i32 } @llvm.hivm.plt.b32.v300(i32 64) + %6 = extractvalue { <256 x i1>, i32 } %5, 0 + %7 = extractvalue { <256 x i1>, i32 } %5, 1 + %8 = trunc i64 %3 to i32 + %9 = mul i32 %8, 4 + %10 = call <64 x float> @llvm.hivm.vldsx1.v64f32(ptr addrspace(6) %0, i32 %9, i32 0, i32 0) + %11 = trunc i64 %3 to i32 + %12 = mul i32 %11, 4 + %13 = call <64 x float> @llvm.hivm.vldsx1.v64f32(ptr addrspace(6) %1, i32 %12, i32 0, i32 0) + %14 = call <64 x float> @llvm.hivm.vadd.v64f32.x(<64 x float> %10, <64 x float> %13, <256 x i1> %6) + %15 = trunc i64 %3 to i32 + %16 = mul i32 %15, 4 + call void @llvm.hivm.vstsx1.v64f32(<64 x float> %14, ptr addrspace(6) %2, i32 %16, i32 2, i32 0, <256 x i1> %6) + ret void +} + +attributes #0 = { "target-cpu"="dav-c310-vec" "target-features"="+ATOMIC,+ArchV130,+AregRedefinable,+ArithmeticBf16,+AtomicForB8 ,+F8e4m3,+F8e5m2,+F8e8m0,+FFTSBlk,+Fp4e1m2x2,+Fp4e2m1x2,+LDExtRefine,+MOVX8,+SPR7bits,+SyncV,+dav-c310-vec" } + +!llvm.module.flags = !{!0} +!hivm.annotations = !{!1, !2} + +!0 = !{i32 2, !"Debug Info Version", i32 3} +!1 = !{ptr @a5_hivm_vadd_demo, !"kernel", i32 1} +!2 = !{ptr @a5_hivm_vadd_demo, !"kernel_with_simd", i32 1} diff --git a/ptodsl/lib/a5/generated/a5_hivm_vadd_demo.pto b/ptodsl/lib/a5/generated/a5_hivm_vadd_demo.pto new file mode 100644 index 00000000..4299bc0a --- /dev/null +++ b/ptodsl/lib/a5/generated/a5_hivm_vadd_demo.pto @@ -0,0 +1,10 @@ +module { + func.func @a5_hivm_vadd_demo(%arg0: !pto.ptr, %arg1: !pto.ptr, %arg2: !pto.ptr, %arg3: index) { + %0 = pto.pset_b32 "PAT_ALL" : !pto.mask + %1 = pto.vlds %arg0[%arg3] : !pto.ptr -> !pto.vreg<64xf32> + %2 = pto.vlds %arg1[%arg3] : !pto.ptr -> !pto.vreg<64xf32> + %3 = pto.vadd %1, %2, %0 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> + pto.vsts %3, %arg2[%arg3], %0 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask + return + } +} diff --git a/ptodsl/lib/a5/generated/a5_vector_copy.ll b/ptodsl/lib/a5/generated/a5_vector_copy.ll new file mode 100644 index 00000000..6e0ec23a --- /dev/null +++ b/ptodsl/lib/a5/generated/a5_vector_copy.ll @@ -0,0 +1,30 @@ +; ModuleID = 'ptoas.hivm.official' +source_filename = "ptoas.hivm.official" + +declare void @llvm.hivm.vstsx1.v64f32(<64 x float>, ptr addrspace(6), i32, i32, i32, <256 x i1>) + +declare { <256 x i1>, i32 } @llvm.hivm.plt.b32.v300(i32) + +declare <64 x float> @llvm.hivm.vldsx1.v64f32(ptr addrspace(6), i32, i32, i32) + +define void @a5_vector_copy(ptr addrspace(6) %0, ptr addrspace(6) %1, i64 %2) #0 { + %4 = trunc i64 %2 to i32 + %5 = mul i32 %4, 4 + %6 = call <64 x float> @llvm.hivm.vldsx1.v64f32(ptr addrspace(6) %0, i32 %5, i32 0, i32 0) + %7 = call { <256 x i1>, i32 } @llvm.hivm.plt.b32.v300(i32 64) + %8 = extractvalue { <256 x i1>, i32 } %7, 0 + %9 = extractvalue { <256 x i1>, i32 } %7, 1 + %10 = trunc i64 %2 to i32 + %11 = mul i32 %10, 4 + call void @llvm.hivm.vstsx1.v64f32(<64 x float> %6, ptr addrspace(6) %1, i32 %11, i32 2, i32 0, <256 x i1> %8) + ret void +} + +attributes #0 = { "target-cpu"="dav-c310-vec" "target-features"="+ATOMIC,+ArchV130,+AregRedefinable,+ArithmeticBf16,+AtomicForB8 ,+F8e4m3,+F8e5m2,+F8e8m0,+FFTSBlk,+Fp4e1m2x2,+Fp4e2m1x2,+LDExtRefine,+MOVX8,+SPR7bits,+SyncV,+dav-c310-vec" } + +!llvm.module.flags = !{!0} +!hivm.annotations = !{!1, !2} + +!0 = !{i32 2, !"Debug Info Version", i32 3} +!1 = !{ptr @a5_vector_copy, !"kernel", i32 1} +!2 = !{ptr @a5_vector_copy, !"kernel_with_simd", i32 1} diff --git a/ptodsl/lib/a5/generated/a5_vector_copy.pto b/ptodsl/lib/a5/generated/a5_vector_copy.pto index 7558636d..2845696f 100644 --- a/ptodsl/lib/a5/generated/a5_vector_copy.pto +++ b/ptodsl/lib/a5/generated/a5_vector_copy.pto @@ -1,10 +1,8 @@ module { func.func @a5_vector_copy(%arg0: !pto.ptr, %arg1: !pto.ptr, %arg2: index) { - pto.section.vector { - %0 = pto.vlds %arg0[%arg2] : !pto.ptr -> !pto.vreg<64xf32> - %1 = pto.pset_b32 "PAT_ALL" : !pto.mask - pto.vsts %0, %arg1[%arg2], %1 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask - } + %0 = pto.vlds %arg0[%arg2] : !pto.ptr -> !pto.vreg<64xf32> + %1 = pto.pset_b32 "PAT_ALL" : !pto.mask + pto.vsts %0, %arg1[%arg2], %1 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask return } } diff --git a/ptodsl/lib/a5/kernels.py b/ptodsl/lib/a5/kernels.py index b66e5370..bc946cd1 100644 --- a/ptodsl/lib/a5/kernels.py +++ b/ptodsl/lib/a5/kernels.py @@ -2,24 +2,42 @@ from mlir.ir import IndexType from ... import Constexpr, language as dsl, scalar as s, to_ir_module +from ...api.scalar import _unwrap from ...language import make_mxfp8 from ._common import ( VF_IMPL_DEFAULT, alloc_tile_buffer, load_tile, + mask_type, make_tensor, ptr, slice_tensor, store_tile, + vreg_type, ) from . import ops +_DSL_DTYPE_BY_NAME = { + "float32": lambda: dsl.float32, + "float16": lambda: dsl.float16, + "bfloat16": lambda: dsl.bfloat16, + "int32": lambda: dsl.int32, + "int16": lambda: dsl.int16, + "int8": lambda: dsl.int8, + "uint32": lambda: dsl.uint32, + "uint16": lambda: dsl.uint16, + "uint8": lambda: dsl.uint8, + "bool": lambda: dsl.bool, +} + def _resolve_dtype(dtype, default_name): if dtype is None: - return getattr(dsl, default_name) + return _DSL_DTYPE_BY_NAME[default_name]() if isinstance(dtype, str): - return getattr(dsl, dtype) + if dtype not in _DSL_DTYPE_BY_NAME: + raise ValueError(f"Unsupported PTODSL dtype '{dtype}'.") + return _DSL_DTYPE_BY_NAME[dtype]() return dtype @@ -69,7 +87,7 @@ def a5_elementwise_add( rhs_tile, out_tile, dtype=element_dtype, - shape=[tile_rows, tile_cols], + tile_shape=[tile_rows, tile_cols], ) return a5_elementwise_add @@ -106,7 +124,7 @@ def a5_templated_elementwise_add( slice_tensor(rhs, offsets=[0, 0], sizes=shape, dtype=element_dtype), slice_tensor(out, offsets=[0, 0], sizes=shape, dtype=element_dtype), dtype=element_dtype, - shape=shape, + tile_shape=shape, impl=VF_IMPL, ) @@ -126,12 +144,34 @@ def meta_data(): @to_ir_module(meta_data=meta_data) def a5_vector_copy(src: "ptr_t", dst: "ptr_t", offset: "index_t") -> None: element_dtype = _resolve_dtype(dtype, "float32") - with dsl.vector_section(): - ops.vector_copy(src, dst, offset, lanes=lanes, dtype=element_dtype) + ops.vector_copy(src, dst, offset, lanes=lanes, dtype=element_dtype) return a5_vector_copy +def build_hivm_vadd_demo(*, lanes=64, dtype=None): + def meta_data(): + element_dtype = _resolve_dtype(dtype, "float32") + return { + "ptr_t": ptr(element_dtype, space="VEC"), + "index_t": IndexType.get(), + } + + @to_ir_module(meta_data=meta_data) + def a5_hivm_vadd_demo( + src0: "ptr_t", src1: "ptr_t", dst: "ptr_t", offset: "index_t" + ) -> None: + element_dtype = _resolve_dtype(dtype, "float32") + vector_type = vreg_type(lanes, element_dtype) + mask = raw_pto.pset_b32(mask_type(), "PAT_ALL") + lhs = raw_pto.vlds(vector_type, _unwrap(src0), _unwrap(offset)) + rhs = raw_pto.vlds(vector_type, _unwrap(src1), _unwrap(offset)) + out = raw_pto.vadd(vector_type, lhs, rhs, mask) + raw_pto.vsts(out, _unwrap(dst), _unwrap(offset), mask) + + return a5_hivm_vadd_demo + + def build_mxfp8_matmul(*, m=16, k=64, n=32, lhs_variant="e5m2", rhs_variant="e5m2"): def meta_data(): mx = make_mxfp8(lhs=lhs_variant, rhs=rhs_variant) @@ -259,7 +299,7 @@ def a5_cube_matmul( lhs, offsets=[0, 0], sizes=[m, k], dtype=lhs_element_dtype ), dtype=lhs_element_dtype, - shape=[m, k], + tile_shape=[m, k], space="MAT", ) rhs_mat = load_tile( @@ -267,7 +307,7 @@ def a5_cube_matmul( rhs, offsets=[0, 0], sizes=[k, n], dtype=rhs_element_dtype ), dtype=rhs_element_dtype, - shape=[k, n], + tile_shape=[k, n], space="MAT", ) lhs_tile = alloc_tile_buffer(lhs_element_dtype, [m, k], space="LEFT") @@ -288,15 +328,23 @@ def a5_cube_matmul( KERNEL_BUILDERS = { "a5_elementwise_add": build_elementwise_add, + "a5_hivm_vadd_demo": build_hivm_vadd_demo, "a5_vector_copy": build_vector_copy, "a5_cube_matmul": build_cube_matmul, } +HIVM_LLVM_KERNELS = { + "a5_hivm_vadd_demo", + "a5_vector_copy", +} + __all__ = [ "KERNEL_BUILDERS", + "HIVM_LLVM_KERNELS", "build_cube_matmul", "build_elementwise_add", + "build_hivm_vadd_demo", "build_vector_copy", "build_mxfp8_matmul", "build_templated_elementwise_add", diff --git a/ptodsl/lib/a5/native.py b/ptodsl/lib/a5/native.py index 1fd2e937..c41db477 100644 --- a/ptodsl/lib/a5/native.py +++ b/ptodsl/lib/a5/native.py @@ -18,6 +18,15 @@ vreg_type, ) +_CMP_MODE_BY_NAME = { + "EQ": pto.CmpMode.EQ, + "NE": pto.CmpMode.NE, + "LT": pto.CmpMode.LT, + "LE": pto.CmpMode.LE, + "GT": pto.CmpMode.GT, + "GE": pto.CmpMode.GE, +} + def adds(src, scalar, out): _call(pto.TAddSOp, src, scalar, out) @@ -91,7 +100,7 @@ def shrs(src, scalar, out): def compare(src0, src1, out, *, mode): cmp_mode = ( - pto.CmpModeAttr.get(getattr(pto.CmpMode, mode.upper())) + pto.CmpModeAttr.get(_CMP_MODE_BY_NAME[mode.upper()]) if isinstance(mode, str) else mode ) diff --git a/ptodsl/lib/a5/ops.py b/ptodsl/lib/a5/ops.py index 043e240f..1270b99c 100644 --- a/ptodsl/lib/a5/ops.py +++ b/ptodsl/lib/a5/ops.py @@ -8,15 +8,9 @@ VF_IMPL_DEFAULT, ) from .native import ( - adds, - and_, - col_expand_max, - col_expand_min, - col_expand_mul, col_prod, compare, concat, - divs, extract, full_mask_b32, insert, @@ -27,32 +21,43 @@ matmul_mx, matmul_mx_acc, matmul_mx_bias, - max, - maxs, - min, - mins, move_tile, - muls, row_prod, scatter, select, - shl, - shls, - shr, - shrs, store_tile, - subs, trans, vector_copy, vload, vstore, - xor, ) -from .tbinary import tadd, tdiv, tmov, tmul, tor_, tsub +from .tbinary import ( + tand, + tadd, + tdiv, + tmax, + tmin, + tmov, + tmul, + tor_, + tshl, + tshr, + tsub, + txor, +) from .texpand import ( tcol_expand, + tcol_expand_add, + tcol_expand_div, + tcol_expand_max, + tcol_expand_min, + tcol_expand_mul, + tcol_expand_sub, trow_expand, + trow_expand_add, trow_expand_div, + trow_expand_max, + trow_expand_min, trow_expand_mul, trow_expand_sub, ) @@ -64,69 +69,62 @@ trow_min, trow_sum, ) +from .tscalar import ( + taxpy, + tadds, + tands, + tdivs, + tlrelu, + tmaxs, + tmins, + tmuls, + tors, + tshls, + tshrs, + tsubs, + txors, +) from .tsort import tgather, tmrgsort, tsort32 from .tunary import tabs, texp, tlog, trecip, trelu, trsqrt, tsqrt -# Readable aliases that match the public tile op names. -mov = tmov -add = tadd -sub = tsub -mul = tmul -div = tdiv -or_ = tor_ -gather = tgather -exp = texp -log = tlog -relu = trelu -abs = tabs -sqrt = tsqrt -rsqrt = trsqrt -reciprocal = trecip -row_sum = trow_sum -row_min = trow_min -row_max = trow_max -row_expand = trow_expand -row_expand_sub = trow_expand_sub -row_expand_mul = trow_expand_mul -row_expand_div = trow_expand_div -col_sum = tcol_sum -col_min = tcol_min -col_max = tcol_max -col_expand = tcol_expand -mrgsort = tmrgsort -sort32 = tsort32 - # A5-style aliases. TLoad = load_tile TStore = store_tile TMov = tmov TAdd = tadd -TAddS = adds +TAddS = tadds TSub = tsub -TSubS = subs +TSubS = tsubs TMul = tmul -TMulS = muls +TMulS = tmuls TDiv = tdiv -TDivS = divs -TMax = max -TMaxS = maxs -TMin = min -TMinS = mins -TAnd = and_ +TDivS = tdivs +TMax = tmax +TMaxS = tmaxs +TMaxs = tmaxs +TMin = tmin +TMinS = tmins +TMins = tmins +TAnd = tand +TAndS = tands TOr = tor_ -TXor = xor -TShl = shl -TShlS = shls -TShr = shr -TShrS = shrs +TOrS = tors +TXor = txor +TXorS = txors +TShl = tshl +TShlS = tshls +TShr = tshr +TShrS = tshrs TCmp = compare TExp = texp TLog = tlog TRelu = trelu +TLRelu = tlrelu TAbs = tabs TSqrt = tsqrt TRsqrt = trsqrt TRecip = trecip +TAxpy = taxpy TGather = tgather TScatter = scatter TSel = select @@ -137,16 +135,22 @@ TRowMin = trow_min TRowMax = trow_max TRowExpand = trow_expand +TRowExpandAdd = trow_expand_add TRowExpandSub = trow_expand_sub TRowExpandMul = trow_expand_mul TRowExpandDiv = trow_expand_div +TRowExpandMax = trow_expand_max +TRowExpandMin = trow_expand_min TColSum = tcol_sum TColMin = tcol_min TColMax = tcol_max TColExpand = tcol_expand -TColExpandMul = col_expand_mul -TColExpandMax = col_expand_max -TColExpandMin = col_expand_min +TColExpandAdd = tcol_expand_add +TColExpandSub = tcol_expand_sub +TColExpandDiv = tcol_expand_div +TColExpandMul = tcol_expand_mul +TColExpandMax = tcol_expand_max +TColExpandMin = tcol_expand_min TMrgSort = tmrgsort TSort32 = tsort32 TTrans = trans @@ -167,10 +171,15 @@ "TAdd", "TAddS", "TAnd", + "TAndS", + "TAxpy", "TColExpand", + "TColExpandAdd", + "TColExpandDiv", "TColExpandMax", "TColExpandMin", "TColExpandMul", + "TColExpandSub", "TColMax", "TColMin", "TColSum", @@ -182,6 +191,7 @@ "TExtract", "TGather", "TInsert", + "TLRelu", "TLoad", "TLog", "TMatmul", @@ -192,17 +202,23 @@ "TMatmulMxBias", "TMax", "TMaxS", + "TMaxs", "TMin", "TMinS", + "TMins", "TMov", "TMrgSort", "TMul", "TMulS", "TOr", + "TOrS", "TRecip", "TRelu", "TRowExpand", + "TRowExpandAdd", "TRowExpandDiv", + "TRowExpandMax", + "TRowExpandMin", "TRowExpandMul", "TRowExpandSub", "TRowMax", @@ -222,96 +238,82 @@ "TSubS", "TTrans", "TXor", - "add", - "adds", - "and_", - "col_expand", - "col_expand_max", - "col_expand_min", - "col_expand_mul", - "col_max", - "col_min", + "TXorS", "col_prod", - "col_sum", "compare", "concat", - "div", - "divs", - "exp", "extract", "full_mask_b32", - "gather", "insert", "load_tile", - "log", "matmul", "matmul_acc", "matmul_bias", "matmul_mx", "matmul_mx_acc", "matmul_mx_bias", - "max", - "maxs", - "min", - "mins", "move_tile", - "mov", - "mrgsort", - "mul", - "muls", - "or_", - "reciprocal", - "relu", - "row_expand", - "row_expand_div", - "row_expand_mul", - "row_expand_sub", - "row_max", - "row_min", "row_prod", - "row_sum", - "rsqrt", "scatter", "select", - "shl", - "shls", - "shr", - "shrs", - "sort32", - "sqrt", "store_tile", - "sub", - "subs", "tabs", "tadd", + "tadds", + "tand", + "tands", + "taxpy", "tcol_expand", + "tcol_expand_add", + "tcol_expand_div", + "tcol_expand_max", + "tcol_expand_min", + "tcol_expand_mul", + "tcol_expand_sub", "tcol_max", "tcol_min", "tcol_sum", "tdiv", + "tdivs", "texp", "tgather", + "tlrelu", "tlog", + "tmax", + "tmaxs", + "tmin", + "tmins", "tmov", "tmrgsort", "tmul", + "tmuls", "tor_", + "tors", "trans", "trecip", "trelu", "trow_expand", + "trow_expand_add", "trow_expand_div", + "trow_expand_max", + "trow_expand_min", "trow_expand_mul", "trow_expand_sub", "trow_max", "trow_min", "trow_sum", "trsqrt", + "tshl", + "tshls", + "tshr", + "tshrs", "tsort32", "tsqrt", "tsub", + "tsubs", + "txor", + "txors", "vector_copy", "vload", "vstore", - "xor", ] diff --git a/ptodsl/lib/a5/tbinary.py b/ptodsl/lib/a5/tbinary.py index 6eb9ce6a..133ed9e1 100644 --- a/ptodsl/lib/a5/tbinary.py +++ b/ptodsl/lib/a5/tbinary.py @@ -4,8 +4,6 @@ directly in terms of `pto.vlds`, `pto.vadd`, and `pto.vsts`. """ -import builtins - from mlir.dialects import pto from ._common import ( @@ -19,11 +17,14 @@ const_expr, const_i64, dtype_byte_width, + flat_active_lanes, mask_for_chunk, + matrix_active_lanes, normalize_vf_impl_kind, ptr, raw, range_constexpr, + resolve_tile_spec, resolve_lanes, s, store_view, @@ -38,7 +39,11 @@ def tadd( out_view, *, dtype, - shape, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, lanes=None, base_addr=0, impl=VF_IMPL_DEFAULT, @@ -48,10 +53,15 @@ def tadd( rhs_view, out_view, dtype=dtype, + tile_shape=tile_shape, shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, lanes=lanes, base_addr=base_addr, - micro_op_name="vadd", + context="TADD", + micro_op=pto.vadd, impl=impl, ) @@ -62,7 +72,11 @@ def tsub( out_view, *, dtype, - shape, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, lanes=None, base_addr=0, impl=VF_IMPL_DEFAULT, @@ -72,10 +86,15 @@ def tsub( rhs_view, out_view, dtype=dtype, + tile_shape=tile_shape, shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, lanes=lanes, base_addr=base_addr, - micro_op_name="vsub", + context="TSUB", + micro_op=pto.vsub, impl=impl, ) @@ -86,7 +105,11 @@ def tmul( out_view, *, dtype, - shape, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, lanes=None, base_addr=0, impl=VF_IMPL_DEFAULT, @@ -96,10 +119,15 @@ def tmul( rhs_view, out_view, dtype=dtype, + tile_shape=tile_shape, shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, lanes=lanes, base_addr=base_addr, - micro_op_name="vmul", + context="TMUL", + micro_op=pto.vmul, impl=impl, ) @@ -110,7 +138,11 @@ def tdiv( out_view, *, dtype, - shape, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, lanes=None, base_addr=0, impl=VF_IMPL_DEFAULT, @@ -120,10 +152,15 @@ def tdiv( rhs_view, out_view, dtype=dtype, + tile_shape=tile_shape, shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, lanes=lanes, base_addr=base_addr, - micro_op_name="vdiv", + context="TDIV", + micro_op=pto.vdiv, impl=impl, ) @@ -134,7 +171,143 @@ def tor_( out_view, *, dtype, - shape, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, + impl=VF_IMPL_DEFAULT, +): + return _binary_tile_vop( + lhs_view, + rhs_view, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + lanes=lanes, + base_addr=base_addr, + context="TOR", + micro_op=pto.vor, + impl=impl, + ) + + +def tmax( + lhs_view, + rhs_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, + impl=VF_IMPL_DEFAULT, +): + return _binary_tile_vop( + lhs_view, + rhs_view, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + lanes=lanes, + base_addr=base_addr, + context="TMAX", + micro_op=pto.vmax, + impl=impl, + ) + + +def tmin( + lhs_view, + rhs_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, + impl=VF_IMPL_DEFAULT, +): + return _binary_tile_vop( + lhs_view, + rhs_view, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + lanes=lanes, + base_addr=base_addr, + context="TMIN", + micro_op=pto.vmin, + impl=impl, + ) + + +def tand( + lhs_view, + rhs_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, + impl=VF_IMPL_DEFAULT, +): + return _binary_tile_vop( + lhs_view, + rhs_view, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + lanes=lanes, + base_addr=base_addr, + context="TAND", + micro_op=pto.vand, + impl=impl, + ) + + +def txor( + lhs_view, + rhs_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, lanes=None, base_addr=0, impl=VF_IMPL_DEFAULT, @@ -144,21 +317,112 @@ def tor_( rhs_view, out_view, dtype=dtype, + tile_shape=tile_shape, shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, lanes=lanes, base_addr=base_addr, - micro_op_name="vor", + context="TXOR", + micro_op=pto.vxor, impl=impl, ) -def tmov(src_view, out_view, *, dtype, shape, lanes=None, base_addr=0): +def tshl( + lhs_view, + rhs_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, + impl=VF_IMPL_DEFAULT, +): + return _binary_tile_vop( + lhs_view, + rhs_view, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + lanes=lanes, + base_addr=base_addr, + context="TSHL", + micro_op=pto.vshl, + impl=impl, + ) + + +def tshr( + lhs_view, + rhs_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, + impl=VF_IMPL_DEFAULT, +): + return _binary_tile_vop( + lhs_view, + rhs_view, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + lanes=lanes, + base_addr=base_addr, + context="TSHR", + micro_op=pto.vshr, + impl=impl, + ) + + +def tmov( + src_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, +): + rows, cols, valid_row, valid_col, type_valid_shape = resolve_tile_spec( + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + context="TMOV", + ) rows, cols = check_tbinop_operands( src_view, src_view, out_view, dtype=dtype, - shape=shape, + shape=[rows, cols], context="TMOV", ) lanes = resolve_lanes(dtype, lanes) @@ -167,8 +431,24 @@ def tmov(src_view, out_view, *, dtype, shape, lanes=None, base_addr=0): src_addr = const_i64(base_addr) out_addr = const_i64(base_addr + buf_bytes) - src_tile = alloc_tile_buffer(dtype, shape, space="VEC", addr=src_addr) - out_tile = alloc_tile_buffer(dtype, shape, space="VEC", addr=out_addr) + src_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + addr=src_addr, + valid_shape=type_valid_shape, + valid_row=valid_row, + valid_col=valid_col, + ) + out_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + addr=out_addr, + valid_shape=type_valid_shape, + valid_row=valid_row, + valid_col=valid_col, + ) load_view(src_view, src_tile) ptr_type = ptr(dtype, space="VEC") @@ -177,7 +457,7 @@ def tmov(src_view, out_view, *, dtype, shape, lanes=None, base_addr=0): vector_type = vreg_type(lanes, dtype) for offset in range_constexpr(0, element_count, lanes): - active = builtins.min(lanes, element_count - offset) + active = flat_active_lanes(valid_row, valid_col, offset, lanes) mask = mask_for_chunk(dtype, active) index = s.const(offset) src_vec = pto.vlds(vector_type, src_ptr, raw(index)) @@ -193,19 +473,32 @@ def _binary_tile_vop( out_view, *, dtype, + tile_shape, shape, + valid_row, + valid_col, + valid_shape, lanes, base_addr, - micro_op_name, + context, + micro_op, impl, ): + rows, cols, valid_row, valid_col, type_valid_shape = resolve_tile_spec( + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + context=context, + ) rows, cols = check_tbinop_operands( lhs_view, rhs_view, out_view, dtype=dtype, - shape=shape, - context=micro_op_name.upper().replace("V", "T", 1), + shape=[rows, cols], + context=context, ) lanes = resolve_lanes(dtype, lanes) element_count = rows * cols @@ -214,9 +507,33 @@ def _binary_tile_vop( rhs_addr = const_i64(base_addr + buf_bytes) out_addr = const_i64(base_addr + buf_bytes * 2) - lhs_tile = alloc_tile_buffer(dtype, shape, space="VEC", addr=lhs_addr) - rhs_tile = alloc_tile_buffer(dtype, shape, space="VEC", addr=rhs_addr) - out_tile = alloc_tile_buffer(dtype, shape, space="VEC", addr=out_addr) + lhs_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + addr=lhs_addr, + valid_shape=type_valid_shape, + valid_row=valid_row, + valid_col=valid_col, + ) + rhs_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + addr=rhs_addr, + valid_shape=type_valid_shape, + valid_row=valid_row, + valid_col=valid_col, + ) + out_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + addr=out_addr, + valid_shape=type_valid_shape, + valid_row=valid_row, + valid_col=valid_col, + ) load_view(lhs_view, lhs_tile) load_view(rhs_view, rhs_tile) @@ -225,10 +542,9 @@ def _binary_tile_vop( lhs_ptr = pto.castptr(ptr_type, lhs_addr) rhs_ptr = pto.castptr(ptr_type, rhs_addr) out_ptr = pto.castptr(ptr_type, out_addr) - micro_op = getattr(pto, micro_op_name) impl_kind = normalize_vf_impl_kind(impl) - is_contiguous = rows == 1 or cols == element_count + is_contiguous = rows == 1 or type_valid_shape[1] == cols if const_expr(impl_kind == VF_IMPL_DEFAULT): impl_kind = ( VF_IMPL_1D_POST_UPDATE if is_contiguous else VF_IMPL_2D_NO_POST_UPDATE @@ -242,6 +558,8 @@ def _binary_tile_vop( dtype=dtype, lanes=lanes, element_count=element_count, + valid_row=valid_row, + valid_col=valid_col, vector_type=vector_type, micro_op=micro_op, ) @@ -254,6 +572,8 @@ def _binary_tile_vop( dtype=dtype, lanes=lanes, element_count=element_count, + valid_row=valid_row, + valid_col=valid_col, vector_type=vector_type, micro_op=micro_op, ) @@ -266,6 +586,8 @@ def _binary_tile_vop( rows=rows, cols=cols, lanes=lanes, + valid_row=valid_row, + valid_col=valid_col, vector_type=vector_type, micro_op=micro_op, ) @@ -279,6 +601,8 @@ def _binary_tile_vop( rows=rows, cols=cols, lanes=lanes, + valid_row=valid_row, + valid_col=valid_col, vector_type=vector_type, micro_op=micro_op, ) @@ -290,10 +614,20 @@ def _binary_tile_vop( def _binary_1d_no_post_update( - lhs_ptr, rhs_ptr, out_ptr, *, dtype, lanes, element_count, vector_type, micro_op + lhs_ptr, + rhs_ptr, + out_ptr, + *, + dtype, + lanes, + element_count, + valid_row, + valid_col, + vector_type, + micro_op, ): for offset in range_constexpr(0, element_count, lanes): - active = builtins.min(lanes, element_count - offset) + active = flat_active_lanes(valid_row, valid_col, offset, lanes) mask = mask_for_chunk(dtype, active) index = s.const(offset) lhs_vec = pto.vlds(vector_type, lhs_ptr, raw(index)) @@ -311,6 +645,8 @@ def _binary_1d_post_update( dtype, lanes, element_count, + valid_row, + valid_col, vector_type, micro_op, ): @@ -319,7 +655,7 @@ def _binary_1d_post_update( out_cursor = out_ptr lane_step = s.const(lanes) for offset in range_constexpr(0, element_count, lanes): - active = builtins.min(lanes, element_count - offset) + active = flat_active_lanes(valid_row, valid_col, offset, lanes) mask = mask_for_chunk(dtype, active) lhs_vec, lhs_cursor = pto.vlds_post( vector_type, ptr_type, lhs_cursor, raw(lane_step) @@ -332,11 +668,22 @@ def _binary_1d_post_update( def _binary_2d_no_post_update( - lhs_ptr, rhs_ptr, out_ptr, *, dtype, rows, cols, lanes, vector_type, micro_op + lhs_ptr, + rhs_ptr, + out_ptr, + *, + dtype, + rows, + cols, + lanes, + valid_row, + valid_col, + vector_type, + micro_op, ): for row in range_constexpr(rows): for col in range_constexpr(0, cols, lanes): - active = builtins.min(lanes, cols - col) + active = matrix_active_lanes(valid_row, valid_col, row, col, lanes) mask = mask_for_chunk(dtype, active) index = s.const(row * cols + col) lhs_vec = pto.vlds(vector_type, lhs_ptr, raw(index)) @@ -355,6 +702,8 @@ def _binary_2d_post_update( rows, cols, lanes, + valid_row, + valid_col, vector_type, micro_op, ): @@ -368,7 +717,7 @@ def _binary_2d_post_update( rhs_cursor = rhs_row out_cursor = out_row for col in range_constexpr(0, cols, lanes): - active = builtins.min(lanes, cols - col) + active = matrix_active_lanes(valid_row, valid_col, row, col, lanes) mask = mask_for_chunk(dtype, active) lhs_vec, lhs_cursor = pto.vlds_post( vector_type, ptr_type, lhs_cursor, raw(lane_step) @@ -389,9 +738,15 @@ def _binary_2d_post_update( "VF_IMPL_2D_NO_POST_UPDATE", "VF_IMPL_2D_POST_UPDATE", "tadd", + "tand", "tdiv", + "tmax", + "tmin", "tmov", "tmul", "tor_", + "tshl", + "tshr", "tsub", + "txor", ] diff --git a/ptodsl/lib/a5/texpand.py b/ptodsl/lib/a5/texpand.py index f2c841dd..bcaad618 100644 --- a/ptodsl/lib/a5/texpand.py +++ b/ptodsl/lib/a5/texpand.py @@ -11,10 +11,12 @@ const_i64, dtype_byte_width, mask_for_chunk, + matrix_active_lanes, micro_lane_count, ptr, raw, range_constexpr, + resolve_tile_spec, s, store_view, load_view, @@ -24,9 +26,28 @@ ) -def tcol_expand(src_view, out_view, *, dtype, shape, base_addr=0): +def tcol_expand( + src_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + base_addr=0, +): + rows, cols, valid_row, valid_col, type_valid_shape = resolve_tile_spec( + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + context="TCOLEXPAND", + ) rows, cols = check_col_expand_operands( - src_view, out_view, dtype=dtype, shape=shape, context="TCOLEXPAND" + src_view, out_view, dtype=dtype, shape=[rows, cols], context="TCOLEXPAND" ) lanes = micro_lane_count(dtype) vector_type = vreg_type(lanes, dtype) @@ -35,16 +56,30 @@ def tcol_expand(src_view, out_view, *, dtype, shape, base_addr=0): src_addr = const_i64(base_addr) out_addr = const_i64(base_addr + buf_bytes) src_tile = alloc_tile_buffer( - dtype, shape, space="VEC", valid_shape=[1, cols], addr=src_addr + dtype, + [rows, cols], + space="VEC", + valid_shape=[1, type_valid_shape[1]], + addr=src_addr, + valid_row=1, + valid_col=valid_col, + ) + out_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + valid_shape=type_valid_shape, + addr=out_addr, + valid_row=valid_row, + valid_col=valid_col, ) - out_tile = alloc_tile_buffer(dtype, shape, space="VEC", addr=out_addr) load_view(src_view, src_tile) src_ptr = pto.castptr(ptr(dtype, space="VEC"), src_addr) out_ptr = pto.castptr(ptr(dtype, space="VEC"), out_addr) for col in range_constexpr(0, cols, lanes): - active = builtins.min(lanes, cols - col) + active = matrix_active_lanes(1, valid_col, 0, col, lanes) mask = mask_for_chunk(dtype, active) col_offset = s.const(col) vec = pto.vlds(vector_type, src_ptr, raw(col_offset)) @@ -56,9 +91,28 @@ def tcol_expand(src_view, out_view, *, dtype, shape, base_addr=0): return out_view -def trow_expand(src_view, out_view, *, dtype, shape, base_addr=0): +def trow_expand( + src_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + base_addr=0, +): + rows, cols, valid_row, valid_col, type_valid_shape = resolve_tile_spec( + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + context="TROWEXPAND", + ) rows, cols = check_row_expand_operands( - src_view, out_view, dtype=dtype, shape=shape, context="TROWEXPAND" + src_view, out_view, dtype=dtype, shape=[rows, cols], context="TROWEXPAND" ) lanes = micro_lane_count(dtype) vector_type = vreg_type(lanes, dtype) @@ -67,9 +121,23 @@ def trow_expand(src_view, out_view, *, dtype, shape, base_addr=0): src_addr = const_i64(base_addr) out_addr = const_i64(base_addr + buf_bytes) src_tile = alloc_tile_buffer( - dtype, shape, space="VEC", valid_shape=[rows, 1], addr=src_addr + dtype, + [rows, cols], + space="VEC", + valid_shape=[type_valid_shape[0], 1], + addr=src_addr, + valid_row=valid_row, + valid_col=1, + ) + out_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + valid_shape=type_valid_shape, + addr=out_addr, + valid_row=valid_row, + valid_col=valid_col, ) - out_tile = alloc_tile_buffer(dtype, shape, space="VEC", addr=out_addr) load_view(src_view, src_tile) src_ptr = pto.castptr(ptr(dtype, space="VEC"), src_addr) @@ -88,7 +156,7 @@ def trow_expand(src_view, out_view, *, dtype, shape, base_addr=0): ) broadcast = pto.vdup(vector_type, scalar_vec, position="POS_LOWEST") for col in range_constexpr(0, cols, lanes): - active = builtins.min(lanes, cols - col) + active = matrix_active_lanes(valid_row, valid_col, row, col, lanes) mask = mask_for_chunk(dtype, active) dst_offset = s.const(row * cols + col) pto.vsts(broadcast, out_ptr, raw(dst_offset), mask) @@ -97,39 +165,351 @@ def trow_expand(src_view, out_view, *, dtype, shape, base_addr=0): return out_view -def trow_expand_sub(base_view, expand_view, out_view, *, dtype, shape, base_addr=0): +def trow_expand_add( + base_view, + expand_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + base_addr=0, +): + return _trow_expand_binary( + base_view, + expand_view, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + base_addr=base_addr, + context="TROWEXPANDADD", + micro_op=pto.vadd, + ) + + +def trow_expand_sub( + base_view, + expand_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + base_addr=0, +): return _trow_expand_binary( base_view, expand_view, out_view, dtype=dtype, + tile_shape=tile_shape, shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, base_addr=base_addr, - micro_op_name="vsub", + context="TROWEXPANDSUB", + micro_op=pto.vsub, ) -def trow_expand_mul(base_view, expand_view, out_view, *, dtype, shape, base_addr=0): +def trow_expand_mul( + base_view, + expand_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + base_addr=0, +): return _trow_expand_binary( base_view, expand_view, out_view, dtype=dtype, + tile_shape=tile_shape, shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, base_addr=base_addr, - micro_op_name="vmul", + context="TROWEXPANDMUL", + micro_op=pto.vmul, ) -def trow_expand_div(base_view, expand_view, out_view, *, dtype, shape, base_addr=0): +def trow_expand_div( + base_view, + expand_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + base_addr=0, +): return _trow_expand_binary( base_view, expand_view, out_view, dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + base_addr=base_addr, + context="TROWEXPANDDIV", + micro_op=pto.vdiv, + ) + + +def trow_expand_max( + base_view, + expand_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + base_addr=0, +): + return _trow_expand_binary( + base_view, + expand_view, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + base_addr=base_addr, + context="TROWEXPANDMAX", + micro_op=pto.vmax, + ) + + +def trow_expand_min( + base_view, + expand_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + base_addr=0, +): + return _trow_expand_binary( + base_view, + expand_view, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + base_addr=base_addr, + context="TROWEXPANDMIN", + micro_op=pto.vmin, + ) + + +def tcol_expand_add( + base_view, + expand_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + base_addr=0, +): + return _tcol_expand_binary( + base_view, + expand_view, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + base_addr=base_addr, + context="TCOLEXPANDADD", + micro_op=pto.vadd, + ) + + +def tcol_expand_sub( + base_view, + expand_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + base_addr=0, +): + return _tcol_expand_binary( + base_view, + expand_view, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + base_addr=base_addr, + context="TCOLEXPANDSUB", + micro_op=pto.vsub, + ) + + +def tcol_expand_div( + base_view, + expand_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + base_addr=0, +): + return _tcol_expand_binary( + base_view, + expand_view, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + base_addr=base_addr, + context="TCOLEXPANDDIV", + micro_op=pto.vdiv, + ) + + +def tcol_expand_mul( + base_view, + expand_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + base_addr=0, +): + return _tcol_expand_binary( + base_view, + expand_view, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + base_addr=base_addr, + context="TCOLEXPANDMUL", + micro_op=pto.vmul, + ) + + +def tcol_expand_max( + base_view, + expand_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + base_addr=0, +): + return _tcol_expand_binary( + base_view, + expand_view, + out_view, + dtype=dtype, + tile_shape=tile_shape, shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, base_addr=base_addr, - micro_op_name="vdiv", + context="TCOLEXPANDMAX", + micro_op=pto.vmax, + ) + + +def tcol_expand_min( + base_view, + expand_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + base_addr=0, +): + return _tcol_expand_binary( + base_view, + expand_view, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + base_addr=base_addr, + context="TCOLEXPANDMIN", + micro_op=pto.vmin, ) @@ -139,26 +519,39 @@ def _trow_expand_binary( out_view, *, dtype, + tile_shape, shape, + valid_row, + valid_col, + valid_shape, base_addr, - micro_op_name, + context, + micro_op, ): + rows, cols, valid_row, valid_col, type_valid_shape = resolve_tile_spec( + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + context=context, + ) rows, cols = check_row_expand_operands( expand_view, out_view, dtype=dtype, - shape=shape, - context=f"TROWEXPAND_{micro_op_name[1:].upper()}", + shape=[rows, cols], + context=context, ) require_view_shape( base_view, [rows, cols], - message=f"Fix: TROWEXPAND_{micro_op_name[1:].upper()} base input valid shape mismatch with output tile dst shape.", + message=f"Fix: {context} base input valid shape mismatch with output tile dst shape.", ) require_view_dtype( base_view, dtype, - message=f"Fix: TROWEXPAND_{micro_op_name[1:].upper()} input data type must be consistent with the output data type.", + message=f"Fix: {context} input data type must be consistent with the output data type.", ) lanes = micro_lane_count(dtype) @@ -168,18 +561,39 @@ def _trow_expand_binary( expand_addr_value = const_i64(base_addr + buf_bytes) out_addr_value = const_i64(base_addr + buf_bytes * 2) - base_tile = alloc_tile_buffer(dtype, shape, space="VEC", addr=base_addr_value) + base_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + addr=base_addr_value, + valid_shape=type_valid_shape, + valid_row=valid_row, + valid_col=valid_col, + ) expand_tile = alloc_tile_buffer( - dtype, shape, space="VEC", valid_shape=[rows, 1], addr=expand_addr_value + dtype, + [rows, cols], + space="VEC", + valid_shape=[type_valid_shape[0], 1], + addr=expand_addr_value, + valid_row=valid_row, + valid_col=1, + ) + out_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + addr=out_addr_value, + valid_shape=type_valid_shape, + valid_row=valid_row, + valid_col=valid_col, ) - out_tile = alloc_tile_buffer(dtype, shape, space="VEC", addr=out_addr_value) load_view(base_view, base_tile) load_view(expand_view, expand_tile) base_ptr = pto.castptr(ptr(dtype, space="VEC"), base_addr_value) expand_ptr = pto.castptr(ptr(dtype, space="VEC"), expand_addr_value) out_ptr = pto.castptr(ptr(dtype, space="VEC"), out_addr_value) - micro_op = getattr(pto, micro_op_name) for row in range_constexpr(rows): scalar_offset = s.const(row * cols) @@ -194,7 +608,7 @@ def _trow_expand_binary( ) broadcast = pto.vdup(vector_type, scalar_vec, position="POS_LOWEST") for col in range_constexpr(0, cols, lanes): - active = builtins.min(lanes, cols - col) + active = matrix_active_lanes(valid_row, valid_col, row, col, lanes) mask = mask_for_chunk(dtype, active) row_offset = s.const(row * cols + col) base_vec = pto.vlds(vector_type, base_ptr, raw(row_offset)) @@ -205,10 +619,116 @@ def _trow_expand_binary( return out_view +def _tcol_expand_binary( + base_view, + expand_view, + out_view, + *, + dtype, + tile_shape, + shape, + valid_row, + valid_col, + valid_shape, + base_addr, + context, + micro_op, +): + rows, cols, valid_row, valid_col, type_valid_shape = resolve_tile_spec( + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + context=context, + ) + rows, cols = check_col_expand_operands( + expand_view, + out_view, + dtype=dtype, + shape=[rows, cols], + context=context, + ) + require_view_shape( + base_view, + [rows, cols], + message=f"Fix: {context} base input valid shape mismatch with output tile dst shape.", + ) + require_view_dtype( + base_view, + dtype, + message=f"Fix: {context} input data type must be consistent with the output data type.", + ) + + lanes = micro_lane_count(dtype) + vector_type = vreg_type(lanes, dtype) + buf_bytes = rows * cols * dtype_byte_width(dtype) + base_addr_value = const_i64(base_addr) + expand_addr_value = const_i64(base_addr + buf_bytes) + out_addr_value = const_i64(base_addr + buf_bytes * 2) + + base_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + addr=base_addr_value, + valid_shape=type_valid_shape, + valid_row=valid_row, + valid_col=valid_col, + ) + expand_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + valid_shape=[1, type_valid_shape[1]], + addr=expand_addr_value, + valid_row=1, + valid_col=valid_col, + ) + out_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + addr=out_addr_value, + valid_shape=type_valid_shape, + valid_row=valid_row, + valid_col=valid_col, + ) + load_view(base_view, base_tile) + load_view(expand_view, expand_tile) + + base_ptr = pto.castptr(ptr(dtype, space="VEC"), base_addr_value) + expand_ptr = pto.castptr(ptr(dtype, space="VEC"), expand_addr_value) + out_ptr = pto.castptr(ptr(dtype, space="VEC"), out_addr_value) + + for col in range_constexpr(0, cols, lanes): + active = matrix_active_lanes(1, valid_col, 0, col, lanes) + mask = mask_for_chunk(dtype, active) + col_offset = s.const(col) + expand_vec = pto.vlds(vector_type, expand_ptr, raw(col_offset)) + for row in range_constexpr(rows): + offset = s.const(row * cols + col) + base_vec = pto.vlds(vector_type, base_ptr, raw(offset)) + out_vec = micro_op(vector_type, base_vec, expand_vec, mask) + pto.vsts(out_vec, out_ptr, raw(offset), mask) + + store_view(out_tile, out_view) + return out_view + + __all__ = [ + "tcol_expand_add", + "tcol_expand_div", "tcol_expand", + "tcol_expand_max", + "tcol_expand_min", + "tcol_expand_mul", + "tcol_expand_sub", + "trow_expand_add", "trow_expand", "trow_expand_div", + "trow_expand_max", + "trow_expand_min", "trow_expand_mul", "trow_expand_sub", ] diff --git a/ptodsl/lib/a5/treduce.py b/ptodsl/lib/a5/treduce.py index 3419e0d6..b800a71c 100644 --- a/ptodsl/lib/a5/treduce.py +++ b/ptodsl/lib/a5/treduce.py @@ -1,6 +1,10 @@ -"""Implement tile reduce ops with PTO vector micro instructions.""" +"""Implement tile reduce ops with PTO vector micro instructions. -import builtins +This file demonstrates how to write tile-style helpers such as +`pto.trow_sum` and `pto.tcol_sum` directly in terms of PTO vector +micro instructions while keeping compile-time tile capacity separate +from runtime valid bounds. +""" from mlir.dialects import pto @@ -19,91 +23,194 @@ dtype_byte_width, full_mask, mask_for_chunk, + matrix_active_lanes, micro_lane_count, normalize_vf_impl_kind, onept_dist, ptr, raw, range_constexpr, + resolve_tile_spec, s, - store_view, - tail_mask, load_view, + store_view, vreg_type, ) -def trow_sum(src_view, out_view, *, dtype, shape, base_addr=0): +def trow_sum( + src_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + base_addr=0, +): return _trow_reduce( src_view, out_view, dtype=dtype, + tile_shape=tile_shape, shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, base_addr=base_addr, - reduce_op_name="vcadd", - combine_op_name="vadd", + context="TROWSUM", + reduce_op=pto.vcadd, + combine_op=pto.vadd, init_value=0.0, ) -def trow_max(src_view, out_view, *, dtype, shape, base_addr=0): +def trow_max( + src_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + base_addr=0, +): return _trow_reduce( src_view, out_view, dtype=dtype, + tile_shape=tile_shape, shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, base_addr=base_addr, - reduce_op_name="vcmax", - combine_op_name="vmax", + context="TROWMAX", + reduce_op=pto.vcmax, + combine_op=pto.vmax, init_value=float("-inf"), ) -def trow_min(src_view, out_view, *, dtype, shape, base_addr=0): +def trow_min( + src_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + base_addr=0, +): return _trow_reduce( src_view, out_view, dtype=dtype, + tile_shape=tile_shape, shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, base_addr=base_addr, - reduce_op_name="vcmin", - combine_op_name="vmin", + context="TROWMIN", + reduce_op=pto.vcmin, + combine_op=pto.vmin, init_value=float("inf"), ) -def tcol_sum(src_view, out_view, *, dtype, shape, base_addr=0, impl=VF_IMPL_DEFAULT): +def tcol_sum( + src_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + base_addr=0, + impl=VF_IMPL_DEFAULT, +): return _tcol_reduce( src_view, out_view, dtype=dtype, + tile_shape=tile_shape, shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, base_addr=base_addr, - micro_op_name="vadd", + context="TCOLSUM", + reduce_op=pto.vadd, + init_value=0.0, impl=impl, ) -def tcol_max(src_view, out_view, *, dtype, shape, base_addr=0, impl=VF_IMPL_DEFAULT): +def tcol_max( + src_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + base_addr=0, + impl=VF_IMPL_DEFAULT, +): return _tcol_reduce( src_view, out_view, dtype=dtype, + tile_shape=tile_shape, shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, base_addr=base_addr, - micro_op_name="vmax", + context="TCOLMAX", + reduce_op=pto.vmax, + init_value=float("-inf"), impl=impl, ) -def tcol_min(src_view, out_view, *, dtype, shape, base_addr=0, impl=VF_IMPL_DEFAULT): +def tcol_min( + src_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + base_addr=0, + impl=VF_IMPL_DEFAULT, +): return _tcol_reduce( src_view, out_view, dtype=dtype, + tile_shape=tile_shape, shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, base_addr=base_addr, - micro_op_name="vmin", + context="TCOLMIN", + reduce_op=pto.vmin, + init_value=float("inf"), impl=impl, ) @@ -113,18 +220,34 @@ def _trow_reduce( out_view, *, dtype, + tile_shape, shape, + valid_row, + valid_col, + valid_shape, base_addr, - reduce_op_name, - combine_op_name, + context, + reduce_op, + combine_op, init_value, ): + validation_context = "TROWREDUCE" + rows, cols, valid_row, valid_col, type_valid_shape = resolve_tile_spec( + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + context=validation_context, + ) rows, cols = check_row_reduce_operands( - src_view, out_view, dtype=dtype, shape=shape, context="TROWREDUCE" + src_view, out_view, dtype=dtype, shape=[rows, cols], context=validation_context ) width = dtype_byte_width(dtype) if width not in {2, 4}: - raise ValueError(f"{reduce_op_name} currently supports only float16/float32.") + raise ValueError( + f"{validation_context} currently supports only float16/float32." + ) lanes = micro_lane_count(dtype) vector_type = vreg_type(lanes, dtype) @@ -132,31 +255,45 @@ def _trow_reduce( src_addr = const_i64(base_addr) out_addr = const_i64(base_addr + buf_bytes) - src_tile = alloc_tile_buffer(dtype, shape, space="VEC", addr=src_addr) + src_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + addr=src_addr, + valid_shape=type_valid_shape, + valid_row=valid_row, + valid_col=valid_col, + ) out_tile = alloc_tile_buffer( - dtype, shape, space="VEC", valid_shape=[rows, 1], addr=out_addr + dtype, + [rows, cols], + space="VEC", + valid_shape=[type_valid_shape[0], 1], + addr=out_addr, + valid_row=valid_row, + valid_col=1, ) load_view(src_view, src_tile) src_ptr = pto.castptr(ptr(dtype, space="VEC"), src_addr) out_ptr = pto.castptr(ptr(dtype, space="VEC"), out_addr) - reduce_op = getattr(pto, reduce_op_name) - combine_op = getattr(pto, combine_op_name) - row_mask = full_mask(dtype) - point_mask = tail_mask(dtype, 1) + vector_mask = full_mask(dtype) init_scalar = const_float(dtype, init_value) + neutral_vec = pto.vbr(vector_type, init_scalar) for row in range_constexpr(rows): - accum = pto.vbr(vector_type, init_scalar) + accum = neutral_vec for col in range_constexpr(0, cols, lanes): - active = builtins.min(lanes, cols - col) + active = matrix_active_lanes(valid_row, valid_col, row, col, lanes) mask = mask_for_chunk(dtype, active) offset = s.const(row * cols + col) vec = pto.vlds(vector_type, src_ptr, raw(offset)) - reduced = reduce_op(vector_type, vec, mask) - accum = combine_op(vector_type, accum, reduced, row_mask) + masked_vec = pto.vsel(vector_type, vec, neutral_vec, mask) + reduced = reduce_op(vector_type, masked_vec, vector_mask) + accum = combine_op(vector_type, accum, reduced, vector_mask) out_offset = s.const(row * cols) - pto.vsts(accum, out_ptr, raw(out_offset), point_mask, dist=onept_dist(dtype)) + store_mask = mask_for_chunk(dtype, matrix_active_lanes(valid_row, 1, row, 0, 1)) + pto.vsts(accum, out_ptr, raw(out_offset), store_mask, dist=onept_dist(dtype)) store_view(out_tile, out_view) return out_view @@ -167,22 +304,51 @@ def _tcol_reduce( out_view, *, dtype, + tile_shape, shape, + valid_row, + valid_col, + valid_shape, base_addr, - micro_op_name, + context, + reduce_op, + init_value, impl, ): + validation_context = "TCOLREDUCE" + rows, cols, valid_row, valid_col, type_valid_shape = resolve_tile_spec( + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + context=validation_context, + ) rows, cols = check_col_reduce_operands( - src_view, out_view, dtype=dtype, shape=shape, context="TCOLREDUCE" + src_view, out_view, dtype=dtype, shape=[rows, cols], context=validation_context ) lanes = micro_lane_count(dtype) buf_bytes = rows * cols * dtype_byte_width(dtype) src_addr = const_i64(base_addr) out_addr = const_i64(base_addr + buf_bytes) - src_tile = alloc_tile_buffer(dtype, shape, space="VEC", addr=src_addr) + src_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + addr=src_addr, + valid_shape=type_valid_shape, + valid_row=valid_row, + valid_col=valid_col, + ) out_tile = alloc_tile_buffer( - dtype, [1, cols], space="VEC", valid_shape=[1, cols], addr=out_addr + dtype, + [rows, cols], + space="VEC", + valid_shape=[1, type_valid_shape[1]], + addr=out_addr, + valid_row=1, + valid_col=valid_col, ) load_view(src_view, src_tile) @@ -190,8 +356,10 @@ def _tcol_reduce( vector_type = vreg_type(lanes, dtype) src_ptr = pto.castptr(ptr_type, src_addr) out_ptr = pto.castptr(ptr_type, out_addr) - reduce_op = getattr(pto, micro_op_name) impl_kind = normalize_vf_impl_kind(impl) + init_scalar = const_float(dtype, init_value) + neutral_vec = pto.vbr(vector_type, init_scalar) + vector_mask = full_mask(dtype) if const_expr(impl_kind == VF_IMPL_DEFAULT): impl_kind = VF_IMPL_1D_POST_UPDATE @@ -203,8 +371,12 @@ def _tcol_reduce( rows=rows, cols=cols, lanes=lanes, + valid_row=valid_row, + valid_col=valid_col, vector_type=vector_type, reduce_op=reduce_op, + neutral_vec=neutral_vec, + vector_mask=vector_mask, ) elif const_expr(impl_kind in {VF_IMPL_1D_POST_UPDATE, VF_IMPL_2D_POST_UPDATE}): _tcol_reduce_post_update( @@ -215,8 +387,12 @@ def _tcol_reduce( rows=rows, cols=cols, lanes=lanes, + valid_row=valid_row, + valid_col=valid_col, vector_type=vector_type, reduce_op=reduce_op, + neutral_vec=neutral_vec, + vector_mask=vector_mask, ) else: raise ValueError(f"Unexpected normalized VF impl kind '{impl_kind}'.") @@ -226,46 +402,72 @@ def _tcol_reduce( def _tcol_reduce_no_post_update( - src_ptr, out_ptr, *, dtype, rows, cols, lanes, vector_type, reduce_op + src_ptr, + out_ptr, + *, + dtype, + rows, + cols, + lanes, + valid_row, + valid_col, + vector_type, + reduce_op, + neutral_vec, + vector_mask, ): - loop_pairs = (rows - 1) // 2 - remain = (rows - 1) % 2 for col in range_constexpr(0, cols, lanes): - active = builtins.min(lanes, cols - col) - mask = mask_for_chunk(dtype, active) - accum = pto.vlds(vector_type, src_ptr, raw(s.const(col))) - for pair in range_constexpr(loop_pairs): - row0 = 2 * pair + 1 - row1 = 2 * pair + 2 - src0 = pto.vlds(vector_type, src_ptr, raw(s.const(col + row0 * cols))) - src1 = pto.vlds(vector_type, src_ptr, raw(s.const(col + row1 * cols))) - pair_sum = reduce_op(vector_type, src0, src1, mask) - accum = reduce_op(vector_type, accum, pair_sum, mask) - if const_expr(remain): - src_tail = pto.vlds( - vector_type, src_ptr, raw(s.const(col + (rows - 1) * cols)) + col_mask = mask_for_chunk( + dtype, matrix_active_lanes(1, valid_col, 0, col, lanes) + ) + accum = neutral_vec + for row in range_constexpr(rows): + row_mask = mask_for_chunk( + dtype, matrix_active_lanes(valid_row, lanes, row, 0, lanes) ) - accum = reduce_op(vector_type, accum, src_tail, mask) - pto.vsts(accum, out_ptr, raw(s.const(col)), mask) + src = pto.vlds(vector_type, src_ptr, raw(s.const(col + row * cols))) + row_filtered = pto.vsel(vector_type, src, neutral_vec, row_mask) + contrib = pto.vsel(vector_type, row_filtered, neutral_vec, col_mask) + accum = reduce_op(vector_type, accum, contrib, vector_mask) + pto.vsts(accum, out_ptr, raw(s.const(col)), col_mask) def _tcol_reduce_post_update( - src_ptr, out_ptr, *, ptr_type, dtype, rows, cols, lanes, vector_type, reduce_op + src_ptr, + out_ptr, + *, + ptr_type, + dtype, + rows, + cols, + lanes, + valid_row, + valid_col, + vector_type, + reduce_op, + neutral_vec, + vector_mask, ): lane_step = s.const(lanes) for col in range_constexpr(0, cols, lanes): - active = builtins.min(lanes, cols - col) - mask = mask_for_chunk(dtype, active) + col_mask = mask_for_chunk( + dtype, matrix_active_lanes(1, valid_col, 0, col, lanes) + ) row0_ptr = pto.addptr(src_ptr, raw(s.const(col))) - accum, _ = pto.vlds_post(vector_type, ptr_type, row0_ptr, raw(lane_step)) - row_ptr = pto.addptr(row0_ptr, raw(s.const(cols))) - for _ in range_constexpr(rows - 1): - src_tail, row_ptr = pto.vlds_post( - vector_type, ptr_type, row_ptr, raw(lane_step) + row_ptr = row0_ptr + accum = neutral_vec + for row in range_constexpr(rows): + src_row, row_ptr = pto.vlds_post( + vector_type, ptr_type, row_ptr, raw(s.const(cols)) + ) + row_mask = mask_for_chunk( + dtype, matrix_active_lanes(valid_row, lanes, row, 0, lanes) ) - accum = reduce_op(vector_type, accum, src_tail, mask) + row_filtered = pto.vsel(vector_type, src_row, neutral_vec, row_mask) + contrib = pto.vsel(vector_type, row_filtered, neutral_vec, col_mask) + accum = reduce_op(vector_type, accum, contrib, vector_mask) out_cursor = pto.addptr(out_ptr, raw(s.const(col))) - pto.vsts_post(ptr_type, accum, out_cursor, raw(lane_step), mask) + pto.vsts_post(ptr_type, accum, out_cursor, raw(lane_step), col_mask) __all__ = [ diff --git a/ptodsl/lib/a5/tscalar.py b/ptodsl/lib/a5/tscalar.py new file mode 100644 index 00000000..b62122b5 --- /dev/null +++ b/ptodsl/lib/a5/tscalar.py @@ -0,0 +1,798 @@ +"""Implement tile scalar ops with PTO vector micro instructions. + +This file demonstrates how to write tile-style helpers such as `pto.tadds` +directly in terms of `pto.vlds`, scalar/vector micro opcodes, and `pto.vsts`. +""" + +from mlir.dialects import pto + +from ._common import ( + VF_IMPL_1D_NO_POST_UPDATE, + VF_IMPL_1D_POST_UPDATE, + VF_IMPL_2D_NO_POST_UPDATE, + VF_IMPL_2D_POST_UPDATE, + VF_IMPL_DEFAULT, + alloc_tile_buffer, + check_tscalar_operands, + const_expr, + const_i64, + const_scalar, + dtype_byte_width, + flat_active_lanes, + mask_for_chunk, + matrix_active_lanes, + normalize_vf_impl_kind, + ptr, + raw, + range_constexpr, + resolve_tile_spec, + resolve_lanes, + s, + store_view, + load_view, + vreg_type, +) + + +def tadds( + src_view, + scalar, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, + impl=VF_IMPL_DEFAULT, +): + return _scalar_tile_vop( + src_view, + scalar, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + lanes=lanes, + base_addr=base_addr, + context="TADDS", + impl=impl, + emit_op=_emit_tadds, + ) + + +def tsubs( + src_view, + scalar, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, + impl=VF_IMPL_DEFAULT, +): + return _scalar_tile_vop( + src_view, + scalar, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + lanes=lanes, + base_addr=base_addr, + context="TSUBS", + impl=impl, + emit_op=_emit_tsubs, + ) + + +def tmuls( + src_view, + scalar, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, + impl=VF_IMPL_DEFAULT, +): + return _scalar_tile_vop( + src_view, + scalar, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + lanes=lanes, + base_addr=base_addr, + context="TMULS", + impl=impl, + emit_op=_emit_tmuls, + ) + + +def tdivs( + src_view, + scalar, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, + impl=VF_IMPL_DEFAULT, +): + return _scalar_tile_vop( + src_view, + scalar, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + lanes=lanes, + base_addr=base_addr, + context="TDIVS", + impl=impl, + emit_op=_emit_tdivs, + ) + + +def tmaxs( + src_view, + scalar, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, + impl=VF_IMPL_DEFAULT, +): + return _scalar_tile_vop( + src_view, + scalar, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + lanes=lanes, + base_addr=base_addr, + context="TMAXS", + impl=impl, + emit_op=_emit_tmaxs, + ) + + +def tmins( + src_view, + scalar, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, + impl=VF_IMPL_DEFAULT, +): + return _scalar_tile_vop( + src_view, + scalar, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + lanes=lanes, + base_addr=base_addr, + context="TMINS", + impl=impl, + emit_op=_emit_tmins, + ) + + +def tands( + src_view, + scalar, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, + impl=VF_IMPL_DEFAULT, +): + return _scalar_tile_vop( + src_view, + scalar, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + lanes=lanes, + base_addr=base_addr, + context="TANDS", + impl=impl, + emit_op=_emit_tands, + ) + + +def tors( + src_view, + scalar, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, + impl=VF_IMPL_DEFAULT, +): + return _scalar_tile_vop( + src_view, + scalar, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + lanes=lanes, + base_addr=base_addr, + context="TORS", + impl=impl, + emit_op=_emit_tors, + ) + + +def txors( + src_view, + scalar, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, + impl=VF_IMPL_DEFAULT, +): + return _scalar_tile_vop( + src_view, + scalar, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + lanes=lanes, + base_addr=base_addr, + context="TXORS", + impl=impl, + emit_op=_emit_txors, + ) + + +def tshls( + src_view, + scalar, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, + impl=VF_IMPL_DEFAULT, +): + return _scalar_tile_vop( + src_view, + scalar, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + lanes=lanes, + base_addr=base_addr, + context="TSHLS", + impl=impl, + emit_op=_emit_tshls, + ) + + +def tshrs( + src_view, + scalar, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, + impl=VF_IMPL_DEFAULT, +): + return _scalar_tile_vop( + src_view, + scalar, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + lanes=lanes, + base_addr=base_addr, + context="TSHRS", + impl=impl, + emit_op=_emit_tshrs, + ) + + +def tlrelu( + src_view, + scalar, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, + impl=VF_IMPL_DEFAULT, +): + return _scalar_tile_vop( + src_view, + scalar, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + lanes=lanes, + base_addr=base_addr, + context="TLRELU", + impl=impl, + emit_op=_emit_tlrelu, + allowed_dtypes={"f32", "f16"}, + ) + + +def taxpy( + src_view, + scalar, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, +): + return _scalar_tile_vop( + src_view, + scalar, + out_view, + dtype=dtype, + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + lanes=lanes, + base_addr=base_addr, + context="TAXPY", + impl=VF_IMPL_2D_NO_POST_UPDATE, + emit_op=_emit_taxpy, + use_out_as_acc=True, + allowed_dtypes={"f32", "f16", "bf16"}, + ) + + +def _scalar_tile_vop( + src_view, + scalar, + out_view, + *, + dtype, + tile_shape, + shape, + valid_row, + valid_col, + valid_shape, + lanes, + base_addr, + context, + impl, + emit_op, + use_out_as_acc=False, + allowed_dtypes=None, +): + rows, cols, valid_row, valid_col, type_valid_shape = resolve_tile_spec( + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + context=context, + ) + rows, cols = check_tscalar_operands( + src_view, + out_view, + dtype=dtype, + shape=[rows, cols], + context=context, + allowed=allowed_dtypes, + ) + lanes = resolve_lanes(dtype, lanes) + element_count = rows * cols + buf_bytes = element_count * dtype_byte_width(dtype) + src_addr = const_i64(base_addr) + out_addr = const_i64(base_addr + buf_bytes) + + src_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + addr=src_addr, + valid_shape=type_valid_shape, + valid_row=valid_row, + valid_col=valid_col, + ) + out_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + addr=out_addr, + valid_shape=type_valid_shape, + valid_row=valid_row, + valid_col=valid_col, + ) + load_view(src_view, src_tile) + if use_out_as_acc: + load_view(out_view, out_tile) + + ptr_type = ptr(dtype, space="VEC") + vector_type = vreg_type(lanes, dtype) + src_ptr = pto.castptr(ptr_type, src_addr) + out_ptr = pto.castptr(ptr_type, out_addr) + scalar_value = raw(scalar) + if not hasattr(scalar_value, "type"): + scalar_value = const_scalar(dtype, scalar) + scalar_vector = pto.vbr(vector_type, scalar_value) + + impl_kind = normalize_vf_impl_kind(impl) + is_contiguous = rows == 1 or type_valid_shape[1] == cols + if const_expr(impl_kind == VF_IMPL_DEFAULT): + impl_kind = VF_IMPL_1D_POST_UPDATE if is_contiguous else VF_IMPL_2D_POST_UPDATE + + if const_expr(impl_kind == VF_IMPL_1D_NO_POST_UPDATE): + _scalar_1d_no_post_update( + src_ptr, + out_ptr, + dtype=dtype, + lanes=lanes, + element_count=element_count, + valid_row=valid_row, + valid_col=valid_col, + vector_type=vector_type, + scalar_value=scalar_value, + scalar_vector=scalar_vector, + emit_op=emit_op, + use_out_as_acc=use_out_as_acc, + ) + elif const_expr(impl_kind == VF_IMPL_1D_POST_UPDATE and not use_out_as_acc): + _scalar_1d_post_update( + src_ptr, + out_ptr, + ptr_type=ptr_type, + dtype=dtype, + lanes=lanes, + element_count=element_count, + valid_row=valid_row, + valid_col=valid_col, + vector_type=vector_type, + scalar_value=scalar_value, + scalar_vector=scalar_vector, + emit_op=emit_op, + ) + elif const_expr( + impl_kind + in { + VF_IMPL_1D_POST_UPDATE, + VF_IMPL_2D_NO_POST_UPDATE, + VF_IMPL_2D_POST_UPDATE, + } + ): + _scalar_2d_no_post_update( + src_ptr, + out_ptr, + dtype=dtype, + rows=rows, + cols=cols, + lanes=lanes, + valid_row=valid_row, + valid_col=valid_col, + vector_type=vector_type, + scalar_value=scalar_value, + scalar_vector=scalar_vector, + emit_op=emit_op, + use_out_as_acc=use_out_as_acc, + ) + else: + raise ValueError(f"Unexpected normalized VF impl kind '{impl_kind}'.") + + store_view(out_tile, out_view) + return out_view + + +def _emit_tadds(*, vector_type, src_vec, acc_vec, mask, scalar_value, scalar_vector): + del acc_vec, mask, scalar_vector + return pto.vadds(vector_type, src_vec, scalar_value) + + +def _emit_tsubs(*, vector_type, src_vec, acc_vec, mask, scalar_value, scalar_vector): + del acc_vec, scalar_value + return pto.vsub(vector_type, src_vec, scalar_vector, mask) + + +def _emit_tmuls(*, vector_type, src_vec, acc_vec, mask, scalar_value, scalar_vector): + del acc_vec, mask, scalar_vector + return pto.vmuls(vector_type, src_vec, scalar_value) + + +def _emit_tdivs(*, vector_type, src_vec, acc_vec, mask, scalar_value, scalar_vector): + del acc_vec, scalar_value + return pto.vdiv(vector_type, src_vec, scalar_vector, mask) + + +def _emit_tmaxs(*, vector_type, src_vec, acc_vec, mask, scalar_value, scalar_vector): + del acc_vec, mask, scalar_vector + return pto.vmaxs(vector_type, src_vec, scalar_value) + + +def _emit_tmins(*, vector_type, src_vec, acc_vec, mask, scalar_value, scalar_vector): + del acc_vec, mask, scalar_vector + return pto.vmins(vector_type, src_vec, scalar_value) + + +def _emit_tands(*, vector_type, src_vec, acc_vec, mask, scalar_value, scalar_vector): + del acc_vec, scalar_value + return pto.vand(vector_type, src_vec, scalar_vector, mask) + + +def _emit_tors(*, vector_type, src_vec, acc_vec, mask, scalar_value, scalar_vector): + del acc_vec, scalar_value + return pto.vor(vector_type, src_vec, scalar_vector, mask) + + +def _emit_txors(*, vector_type, src_vec, acc_vec, mask, scalar_value, scalar_vector): + del acc_vec, scalar_value + return pto.vxor(vector_type, src_vec, scalar_vector, mask) + + +def _emit_tshls(*, vector_type, src_vec, acc_vec, mask, scalar_value, scalar_vector): + del acc_vec, mask, scalar_vector + return pto.vshls(vector_type, src_vec, scalar_value) + + +def _emit_tshrs(*, vector_type, src_vec, acc_vec, mask, scalar_value, scalar_vector): + del acc_vec, mask, scalar_vector + return pto.vshrs(vector_type, src_vec, scalar_value) + + +def _emit_tlrelu(*, vector_type, src_vec, acc_vec, mask, scalar_value, scalar_vector): + del acc_vec, mask, scalar_vector + return pto.vlrelu(vector_type, src_vec, scalar_value) + + +def _emit_taxpy(*, vector_type, src_vec, acc_vec, mask, scalar_value, scalar_vector): + del scalar_value + return pto.vmula(vector_type, acc_vec, src_vec, scalar_vector, mask) + + +def _scalar_1d_no_post_update( + src_ptr, + out_ptr, + *, + dtype, + lanes, + element_count, + valid_row, + valid_col, + vector_type, + scalar_value, + scalar_vector, + emit_op, + use_out_as_acc, +): + for offset in range_constexpr(0, element_count, lanes): + active = flat_active_lanes(valid_row, valid_col, offset, lanes) + mask = mask_for_chunk(dtype, active) + index = s.const(offset) + src_vec = pto.vlds(vector_type, src_ptr, raw(index)) + acc_vec = ( + pto.vlds(vector_type, out_ptr, raw(index)) if use_out_as_acc else src_vec + ) + out_vec = emit_op( + vector_type=vector_type, + src_vec=src_vec, + acc_vec=acc_vec, + mask=mask, + scalar_value=scalar_value, + scalar_vector=scalar_vector, + ) + pto.vsts(out_vec, out_ptr, raw(index), mask) + + +def _scalar_1d_post_update( + src_ptr, + out_ptr, + *, + ptr_type, + dtype, + lanes, + element_count, + valid_row, + valid_col, + vector_type, + scalar_value, + scalar_vector, + emit_op, +): + cursor_in = src_ptr + cursor_out = out_ptr + lane_step = s.const(lanes) + + for offset in range_constexpr(0, element_count, lanes): + active = flat_active_lanes(valid_row, valid_col, offset, lanes) + mask = mask_for_chunk(dtype, active) + src_vec, cursor_in = pto.vlds_post( + vector_type, ptr_type, cursor_in, raw(lane_step) + ) + out_vec = emit_op( + vector_type=vector_type, + src_vec=src_vec, + acc_vec=src_vec, + mask=mask, + scalar_value=scalar_value, + scalar_vector=scalar_vector, + ) + cursor_out = pto.vsts_post(ptr_type, out_vec, cursor_out, raw(lane_step), mask) + + +def _scalar_2d_no_post_update( + src_ptr, + out_ptr, + *, + dtype, + rows, + cols, + lanes, + valid_row, + valid_col, + vector_type, + scalar_value, + scalar_vector, + emit_op, + use_out_as_acc, +): + for row in range_constexpr(rows): + row_offset = row * cols + for col in range_constexpr(0, cols, lanes): + active = matrix_active_lanes(valid_row, valid_col, row, col, lanes) + mask = mask_for_chunk(dtype, active) + offset = s.const(row_offset + col) + src_vec = pto.vlds(vector_type, src_ptr, raw(offset)) + acc_vec = ( + pto.vlds(vector_type, out_ptr, raw(offset)) + if use_out_as_acc + else src_vec + ) + out_vec = emit_op( + vector_type=vector_type, + src_vec=src_vec, + acc_vec=acc_vec, + mask=mask, + scalar_value=scalar_value, + scalar_vector=scalar_vector, + ) + pto.vsts(out_vec, out_ptr, raw(offset), mask) + + +__all__ = [ + "VF_IMPL_DEFAULT", + "VF_IMPL_1D_NO_POST_UPDATE", + "VF_IMPL_1D_POST_UPDATE", + "VF_IMPL_2D_NO_POST_UPDATE", + "VF_IMPL_2D_POST_UPDATE", + "taxpy", + "tadds", + "tands", + "tdivs", + "tlrelu", + "tmaxs", + "tmins", + "tmuls", + "tors", + "tshls", + "tshrs", + "tsubs", + "txors", +] diff --git a/ptodsl/lib/a5/tsort.py b/ptodsl/lib/a5/tsort.py index 72884ded..bc972bd8 100644 --- a/ptodsl/lib/a5/tsort.py +++ b/ptodsl/lib/a5/tsort.py @@ -1,8 +1,7 @@ """Implement gather/sort tile ops with PTO vector micro instructions.""" -import builtins - -from mlir.dialects import pto +from mlir.dialects import arith, pto +from mlir.ir import IndexType from ._common import ( alloc_tile_buffer, @@ -12,14 +11,16 @@ const_i64, dtype_byte_width, mask_for_chunk, + matrix_active_lanes, micro_lane_count, ptr, raw, range_constexpr, + resolve_tile_spec, s, store_view, - uint32_type, load_view, + uint32_type, vreg_type, ) @@ -31,17 +32,29 @@ def tgather( *, dtype, index_dtype=None, - shape, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, base_addr=0, ): index_dtype = uint32_type() if index_dtype is None else index_dtype + rows, cols, valid_row, valid_col, type_valid_shape = resolve_tile_spec( + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + context="TGATHER", + ) rows, cols = check_gather_operands( src_view, indices_view, out_view, dtype=dtype, index_dtype=index_dtype, - shape=shape, + shape=[rows, cols], ) src_bytes = rows * cols * dtype_byte_width(dtype) idx_bytes = rows * cols * dtype_byte_width(index_dtype) @@ -50,9 +63,33 @@ def tgather( idx_addr = const_i64(base_addr + src_bytes) out_addr = const_i64(base_addr + src_bytes + idx_bytes) - src_tile = alloc_tile_buffer(dtype, shape, space="VEC", addr=src_addr) - idx_tile = alloc_tile_buffer(index_dtype, shape, space="VEC", addr=idx_addr) - out_tile = alloc_tile_buffer(dtype, shape, space="VEC", addr=out_addr) + src_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + addr=src_addr, + valid_shape=type_valid_shape, + valid_row=valid_row, + valid_col=valid_col, + ) + idx_tile = alloc_tile_buffer( + index_dtype, + [rows, cols], + space="VEC", + addr=idx_addr, + valid_shape=type_valid_shape, + valid_row=valid_row, + valid_col=valid_col, + ) + out_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + addr=out_addr, + valid_shape=type_valid_shape, + valid_row=valid_row, + valid_col=valid_col, + ) load_view(src_view, src_tile) load_view(indices_view, idx_tile) @@ -66,26 +103,60 @@ def tgather( for row in range_constexpr(rows): row_base = row * cols for col in range_constexpr(0, cols, lanes): - active = builtins.min(lanes, cols - col) + active = matrix_active_lanes(valid_row, valid_col, row, col, lanes) offset = s.const(row_base + col) mask = mask_for_chunk(dtype, active) idx_vec = pto.vlds(index_vector_type, idx_ptr, raw(offset)) - out_vec = pto.vgather2(vector_type, src_ptr, idx_vec, raw(s.const(active))) + if isinstance(active, int): + active_lanes = raw(s.const(active)) + else: + active_lanes = arith.IndexCastOp(IndexType.get(), raw(active)).result + out_vec = pto.vgather2(vector_type, src_ptr, idx_vec, active_lanes) pto.vsts(out_vec, out_ptr, raw(offset), mask) store_view(out_tile, out_view) return out_view -def tmrgsort(src_view, out_view, *, dtype, shape, block_len, base_addr=0): +def tmrgsort( + src_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + block_len, + base_addr=0, +): + rows, cols, valid_row, valid_col, type_valid_shape = resolve_tile_spec( + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + context="TMRGSORT", + ) _, cols = check_mrgsort_operands( - src_view, out_view, dtype=dtype, shape=shape, block_len=block_len + src_view, out_view, dtype=dtype, shape=[rows, cols], block_len=block_len ) + if ( + not isinstance(valid_row, int) + or not isinstance(valid_col, int) + or valid_row != rows + or valid_col != cols + or type_valid_shape != [rows, cols] + ): + raise ValueError( + "TMRGSORT micro lowering currently requires a fully valid single-row tile." + ) src_addr = const_i64(base_addr) out_addr = const_i64(base_addr + cols * dtype_byte_width(dtype)) - src_tile = alloc_tile_buffer(dtype, shape, space="VEC", addr=src_addr) - out_tile = alloc_tile_buffer(dtype, shape, space="VEC", addr=out_addr) + src_tile = alloc_tile_buffer(dtype, [rows, cols], space="VEC", addr=src_addr) + out_tile = alloc_tile_buffer(dtype, [rows, cols], space="VEC", addr=out_addr) load_view(src_view, src_tile) ptr_type = ptr(dtype, space="VEC") @@ -118,10 +189,40 @@ def tmrgsort(src_view, out_view, *, dtype, shape, block_len, base_addr=0): return out_view -def tsort32(src_view, idx_view, out_view, *, dtype, shape, base_addr=0): +def tsort32( + src_view, + idx_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + base_addr=0, +): + rows, cols, valid_row, valid_col, type_valid_shape = resolve_tile_spec( + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + context="TSORT32", + ) rows, cols, out_cols = check_sort32_operands( - src_view, idx_view, out_view, dtype=dtype, shape=shape + src_view, idx_view, out_view, dtype=dtype, shape=[rows, cols] ) + if ( + not isinstance(valid_row, int) + or not isinstance(valid_col, int) + or valid_row != rows + or valid_col != cols + or type_valid_shape != [rows, cols] + ): + raise ValueError( + "TSORT32 micro lowering currently requires a fully valid input tile." + ) src_bytes = rows * cols * dtype_byte_width(dtype) idx_bytes = rows * cols * 4 diff --git a/ptodsl/lib/a5/tunary.py b/ptodsl/lib/a5/tunary.py index 3eee318a..771ca4af 100644 --- a/ptodsl/lib/a5/tunary.py +++ b/ptodsl/lib/a5/tunary.py @@ -4,8 +4,6 @@ directly in terms of `pto.vlds`, a unary vector opcode, and `pto.vsts`. """ -import builtins - from mlir.dialects import pto from ._common import ( @@ -13,10 +11,12 @@ check_tbinop_operands, const_i64, dtype_byte_width, + flat_active_lanes, mask_for_chunk, ptr, raw, range_constexpr, + resolve_tile_spec, resolve_lanes, s, store_view, @@ -25,91 +25,207 @@ ) -def texp(src_view, out_view, *, dtype, shape, lanes=None, base_addr=0): +def texp( + src_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, +): return _unary_tile_vop( src_view, out_view, dtype=dtype, + tile_shape=tile_shape, shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, lanes=lanes, base_addr=base_addr, - micro_op_name="vexp", context="TEXP", + micro_op=pto.vexp, ) -def tlog(src_view, out_view, *, dtype, shape, lanes=None, base_addr=0): +def tlog( + src_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, +): return _unary_tile_vop( src_view, out_view, dtype=dtype, + tile_shape=tile_shape, shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, lanes=lanes, base_addr=base_addr, - micro_op_name="vln", context="TLOG", + micro_op=pto.vln, ) -def trelu(src_view, out_view, *, dtype, shape, lanes=None, base_addr=0): +def trelu( + src_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, +): return _unary_tile_vop( src_view, out_view, dtype=dtype, + tile_shape=tile_shape, shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, lanes=lanes, base_addr=base_addr, - micro_op_name="vrelu", context="TRELU", + micro_op=pto.vrelu, ) -def tabs(src_view, out_view, *, dtype, shape, lanes=None, base_addr=0): +def tabs( + src_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, +): return _unary_tile_vop( src_view, out_view, dtype=dtype, + tile_shape=tile_shape, shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, lanes=lanes, base_addr=base_addr, - micro_op_name="vabs", context="TABS", + micro_op=pto.vabs, ) -def tsqrt(src_view, out_view, *, dtype, shape, lanes=None, base_addr=0): +def tsqrt( + src_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, +): return _unary_tile_vop( src_view, out_view, dtype=dtype, + tile_shape=tile_shape, shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, lanes=lanes, base_addr=base_addr, - micro_op_name="vsqrt", context="TSQRT", + micro_op=pto.vsqrt, ) -def trecip(src_view, out_view, *, dtype, shape, lanes=None, base_addr=0): +def trecip( + src_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, +): return _unary_tile_vop( src_view, out_view, dtype=dtype, + tile_shape=tile_shape, shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, lanes=lanes, base_addr=base_addr, - micro_op_name="vrec", context="TRECIP", + micro_op=pto.vrec, ) -def trsqrt(src_view, out_view, *, dtype, shape, lanes=None, base_addr=0): +def trsqrt( + src_view, + out_view, + *, + dtype, + tile_shape=None, + shape=None, + valid_row=None, + valid_col=None, + valid_shape=None, + lanes=None, + base_addr=0, +): + rows, cols, valid_row, valid_col, type_valid_shape = resolve_tile_spec( + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + context="TRSQRT", + ) rows, cols = check_tbinop_operands( src_view, src_view, out_view, dtype=dtype, - shape=shape, + shape=[rows, cols], context="TRSQRT", ) lanes = resolve_lanes(dtype, lanes) @@ -118,8 +234,24 @@ def trsqrt(src_view, out_view, *, dtype, shape, lanes=None, base_addr=0): src_addr = const_i64(base_addr) out_addr = const_i64(base_addr + buf_bytes) - src_tile = alloc_tile_buffer(dtype, shape, space="VEC", addr=src_addr) - out_tile = alloc_tile_buffer(dtype, shape, space="VEC", addr=out_addr) + src_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + addr=src_addr, + valid_shape=type_valid_shape, + valid_row=valid_row, + valid_col=valid_col, + ) + out_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + addr=out_addr, + valid_shape=type_valid_shape, + valid_row=valid_row, + valid_col=valid_col, + ) load_view(src_view, src_tile) ptr_type = ptr(dtype, space="VEC") @@ -128,7 +260,7 @@ def trsqrt(src_view, out_view, *, dtype, shape, lanes=None, base_addr=0): out_ptr = pto.castptr(ptr_type, out_addr) for offset in range_constexpr(0, element_count, lanes): - active = builtins.min(lanes, element_count - offset) + active = flat_active_lanes(valid_row, valid_col, offset, lanes) mask = mask_for_chunk(dtype, active) index = s.const(offset) src_vec = pto.vlds(vector_type, src_ptr, raw(index)) @@ -145,18 +277,30 @@ def _unary_tile_vop( out_view, *, dtype, + tile_shape, shape, + valid_row, + valid_col, + valid_shape, lanes, base_addr, - micro_op_name, context, + micro_op, ): + rows, cols, valid_row, valid_col, type_valid_shape = resolve_tile_spec( + tile_shape=tile_shape, + shape=shape, + valid_row=valid_row, + valid_col=valid_col, + valid_shape=valid_shape, + context=context, + ) rows, cols = check_tbinop_operands( src_view, src_view, out_view, dtype=dtype, - shape=shape, + shape=[rows, cols], context=context, ) lanes = resolve_lanes(dtype, lanes) @@ -165,18 +309,33 @@ def _unary_tile_vop( src_addr = const_i64(base_addr) out_addr = const_i64(base_addr + buf_bytes) - src_tile = alloc_tile_buffer(dtype, shape, space="VEC", addr=src_addr) - out_tile = alloc_tile_buffer(dtype, shape, space="VEC", addr=out_addr) + src_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + addr=src_addr, + valid_shape=type_valid_shape, + valid_row=valid_row, + valid_col=valid_col, + ) + out_tile = alloc_tile_buffer( + dtype, + [rows, cols], + space="VEC", + addr=out_addr, + valid_shape=type_valid_shape, + valid_row=valid_row, + valid_col=valid_col, + ) load_view(src_view, src_tile) ptr_type = ptr(dtype, space="VEC") vector_type = vreg_type(lanes, dtype) src_ptr = pto.castptr(ptr_type, src_addr) out_ptr = pto.castptr(ptr_type, out_addr) - micro_op = getattr(pto, micro_op_name) for offset in range_constexpr(0, element_count, lanes): - active = builtins.min(lanes, element_count - offset) + active = flat_active_lanes(valid_row, valid_col, offset, lanes) mask = mask_for_chunk(dtype, active) index = s.const(offset) src_vec = pto.vlds(vector_type, src_ptr, raw(index)) diff --git a/scripts/generate_a5_pto.py b/scripts/generate_a5_pto.py index a7d97d6c..6e7fcdcb 100644 --- a/scripts/generate_a5_pto.py +++ b/scripts/generate_a5_pto.py @@ -17,30 +17,72 @@ _DEFAULT_PTOAS = _ROOT.parent / "PTOAS" / "build-src312" / "tools" / "ptoas" / "ptoas" -def emit_kernels(*, output_dir, ptoas_bin=None, emit_cpp=False): +def _run_ptoas(ptoas_bin, args, *, output_dir, warning): + try: + subprocess.run( + [str(ptoas_bin), *args], + check=True, + cwd=str(output_dir), + ) + except subprocess.CalledProcessError as exc: + print(f"warning: {warning}: {exc}", file=sys.stderr) + + +def emit_kernels( + *, + output_dir, + ptoas_bin=None, + emit_cpp=False, + emit_hivm_llvm=False, + kernel_names=None, +): output_dir.mkdir(parents=True, exist_ok=True) generated = [] + selected = set(kernel_names) if kernel_names is not None else None + hivm_selected = ( + set(a5.HIVM_LLVM_KERNELS) + if selected is None + else set(a5.HIVM_LLVM_KERNELS).intersection(selected) + ) for kernel_name, builder in a5.KERNEL_BUILDERS.items(): + if selected is not None and kernel_name not in selected: + continue module = builder() pto_path = output_dir / f"{kernel_name}.pto" pto_path.write_text(f"{module}\n", encoding="utf-8") generated.append(pto_path) - if emit_cpp: + if emit_cpp or emit_hivm_llvm: if ptoas_bin is None: - raise ValueError("`emit_cpp=True` requires `ptoas_bin`.") - cpp_path = output_dir / f"{kernel_name}.cpp" - try: - subprocess.run( - [str(ptoas_bin), str(pto_path), "-o", str(cpp_path)], - check=True, - cwd=str(output_dir), - ) - except subprocess.CalledProcessError as exc: - print( - f"warning: failed to lower {pto_path.name} to C++ with ptoas: {exc}", - file=sys.stderr, + raise ValueError( + "`emit_cpp=True` and `emit_hivm_llvm=True` require `ptoas_bin`." ) + + if emit_cpp: + cpp_path = output_dir / f"{kernel_name}.cpp" + _run_ptoas( + ptoas_bin, + [str(pto_path), "-o", str(cpp_path)], + output_dir=output_dir, + warning=f"failed to lower {pto_path.name} to C++ with ptoas", + ) + + if emit_hivm_llvm and kernel_name in hivm_selected: + llvm_path = output_dir / f"{kernel_name}.ll" + _run_ptoas( + ptoas_bin, + [ + "--pto-arch=a5", + "--pto-level=level3", + "--pto-backend=vpto", + "--vpto-emit-hivm-llvm", + str(pto_path), + "-o", + str(llvm_path), + ], + output_dir=output_dir, + warning=f"failed to lower {pto_path.name} to HIVM LLVM with ptoas", + ) return generated @@ -65,16 +107,22 @@ def _parse_args(): action="store_true", help="Also run ptoas and write `.cpp` files next to the generated `.pto` files.", ) + parser.add_argument( + "--emit-hivm-llvm", + action="store_true", + help="Also run PTOAS VPTO lowering and write `.ll` files for kernels that are marked HIVM-lowerable.", + ) return parser.parse_args() def main(): args = _parse_args() - ptoas_bin = args.ptoas if args.emit_cpp else None + ptoas_bin = args.ptoas if args.emit_cpp or args.emit_hivm_llvm else None generated = emit_kernels( output_dir=args.output_dir, ptoas_bin=ptoas_bin, emit_cpp=args.emit_cpp, + emit_hivm_llvm=args.emit_hivm_llvm, ) for path in generated: print(path) diff --git a/scripts/update_a5_header_coverage.py b/scripts/update_a5_header_coverage.py new file mode 100644 index 00000000..f9b7964c --- /dev/null +++ b/scripts/update_a5_header_coverage.py @@ -0,0 +1,21 @@ +#!/usr/bin/env python3 + +from pathlib import Path +import sys + + +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from ptodsl.lib.a5.a5_header_coverage import a5_header_coverage_markdown + + +def main(): + target = _REPO_ROOT / "ptodsl" / "lib" / "a5" / "A5_HEADER_COVERAGE.md" + target.write_text(a5_header_coverage_markdown(), encoding="utf-8") + print(target) + + +if __name__ == "__main__": + main() diff --git a/tests/regression/test_a5_header_coverage.py b/tests/regression/test_a5_header_coverage.py new file mode 100644 index 00000000..baf6eb82 --- /dev/null +++ b/tests/regression/test_a5_header_coverage.py @@ -0,0 +1,48 @@ +from pathlib import Path + +from ptodsl.lib import a5 +from ptodsl.lib.a5.a5_header_coverage import ( + A5_HEADER_COVERAGE, + A5_HEADER_INVENTORY, + a5_header_coverage_markdown, + a5_header_coverage_summary, +) + + +def test_a5_header_inventory_is_tracked_completely(): + assert len(A5_HEADER_INVENTORY) == 116 + assert set(A5_HEADER_COVERAGE) == set(A5_HEADER_INVENTORY) + + +def test_implemented_and_partial_a5_helpers_exist(): + for entry in A5_HEADER_COVERAGE.values(): + if entry["status"] not in {"implemented", "partial"}: + continue + helper = entry["helper"] + if helper is None or "." in helper: + continue + assert getattr(a5, helper) is not None + + +def test_a5_header_coverage_markdown_mentions_all_headers(): + text = a5_header_coverage_markdown() + for name in A5_HEADER_INVENTORY: + assert f"`{name}`" in text + + +def test_a5_header_coverage_summary_matches_inventory(): + counts = a5_header_coverage_summary() + assert sum(counts.values()) == len(A5_HEADER_INVENTORY) + assert counts["implemented"] > 0 + assert counts["pending"] > 0 + + +def test_checked_in_a5_header_checklist_is_in_sync(): + checklist = ( + Path(__file__).resolve().parents[2] + / "ptodsl" + / "lib" + / "a5" + / "A5_HEADER_COVERAGE.md" + ) + assert checklist.read_text(encoding="utf-8") == a5_header_coverage_markdown() diff --git a/tests/regression/test_a5_lib_regression.py b/tests/regression/test_a5_lib_regression.py index 981c5eae..3a65a79e 100644 --- a/tests/regression/test_a5_lib_regression.py +++ b/tests/regression/test_a5_lib_regression.py @@ -1,11 +1,24 @@ +from pathlib import Path + import pytest +from mlir.dialects import arith from mlir.ir import IndexType, IntegerType import ptodsl.language as pto +import ptodsl.pto as public_pto from ptodsl import to_ir_module from ptodsl.lib import a5 from scripts.generate_a5_pto import emit_kernels +_PTOAS_BIN = ( + Path(__file__).resolve().parents[3] + / "PTOAS" + / "build-src312" + / "tools" + / "ptoas" + / "ptoas" +) + def _index(value): return pto.const(value) if isinstance(value, int) else value @@ -41,12 +54,24 @@ def _slice_tensor(source, *, offsets, sizes, dtype): def test_a5_split_modules_are_publicly_exposed(): assert a5.tbinary.tadd is a5.tadd + assert a5.tscalar.tadds is a5.tadds assert a5.tunary.trsqrt is a5.trsqrt assert a5.texpand.trow_expand is a5.trow_expand assert a5.treduce.trow_sum is a5.trow_sum assert a5.tsort.tgather is a5.tgather +def test_public_pto_ptr_supports_explicit_memory_spaces(): + def meta_data(): + return {"vec_ptr_t": public_pto.ptr(pto.float32, space="VEC")} + + @to_ir_module(meta_data=meta_data) + def ptr_surface_demo(arg: "vec_ptr_t") -> None: + return None + + assert "!pto.ptr" in str(ptr_surface_demo) + + def test_a5_elementwise_add_kernel_emits_tile_flow(): text = str(a5.build_elementwise_add()) @@ -89,6 +114,80 @@ def test_a5_vector_copy_emits_vector_opcodes(): assert "pto.vsts" in text +def test_a5_tadd_separates_tile_shape_from_dynamic_valid_bounds(): + def meta_data(): + return { + "ptr_t": pto.PtrType(pto.float32), + "index_t": IndexType.get(), + } + + @to_ir_module(meta_data=meta_data) + def a5_tadd_dynamic_valid( + src0: "ptr_t", + src1: "ptr_t", + dst: "ptr_t", + valid_row: "index_t", + valid_col: "index_t", + ) -> None: + lhs = _make_tensor(src0, shape=[8, 64], dtype=pto.float32) + rhs = _make_tensor(src1, shape=[8, 64], dtype=pto.float32) + out = _make_tensor(dst, shape=[8, 64], dtype=pto.float32) + with pto.vector_section(): + a5.tadd( + _slice_tensor(lhs, offsets=[0, 0], sizes=[8, 64], dtype=pto.float32), + _slice_tensor(rhs, offsets=[0, 0], sizes=[8, 64], dtype=pto.float32), + _slice_tensor(out, offsets=[0, 0], sizes=[8, 64], dtype=pto.float32), + dtype=pto.float32, + tile_shape=[8, 64], + valid_row=valid_row, + valid_col=valid_col, + ) + + text = str(a5_tadd_dynamic_valid) + + assert "!pto.tile_buf" in text + assert "valid_row = %arg3 valid_col = %arg4" in text + assert "pto.vadd" in text + assert "pto.tadd" not in text + + +def test_a5_trow_sum_uses_dynamic_valid_bounds_and_masked_reduction(): + def meta_data(): + return { + "ptr_src": pto.PtrType(pto.float32), + "ptr_dst": pto.PtrType(pto.float32), + "index_t": IndexType.get(), + } + + @to_ir_module(meta_data=meta_data) + def a5_trow_sum_dynamic_valid( + src: "ptr_src", dst: "ptr_dst", valid_row: "index_t", valid_col: "index_t" + ) -> None: + src_view = _make_tensor(src, shape=[8, 64], dtype=pto.float32) + dst_view = _make_tensor(dst, shape=[8, 1], dtype=pto.float32) + with pto.vector_section(): + a5.trow_sum( + _slice_tensor( + src_view, offsets=[0, 0], sizes=[8, 64], dtype=pto.float32 + ), + _slice_tensor( + dst_view, offsets=[0, 0], sizes=[8, 1], dtype=pto.float32 + ), + dtype=pto.float32, + tile_shape=[8, 64], + valid_row=valid_row, + valid_col=valid_col, + ) + + text = str(a5_trow_sum_dynamic_valid) + + assert "!pto.tile_buf" in text + assert "valid_row = %arg2" in text + assert "pto.vsel" in text + assert "pto.vcadd" in text + assert 'dist = "ONEPT_B32"' in text + + def test_a5_tcol_expand_emits_broadcast_micro_ops(): def meta_data(): return { @@ -160,6 +259,52 @@ def a5_tgather(src: "ptr_src", idx: "ptr_idx", dst: "ptr_src") -> None: assert "pto.tgather" not in text +def test_a5_tgather_supports_dynamic_valid_bounds(): + def uint32(): + return IntegerType.get_unsigned(32) + + def meta_data(): + return { + "ptr_src": pto.PtrType(pto.float32), + "ptr_idx": pto.PtrType(uint32()), + "index_t": IndexType.get(), + } + + @to_ir_module(meta_data=meta_data) + def a5_tgather_dynamic_valid( + src: "ptr_src", + idx: "ptr_idx", + dst: "ptr_src", + valid_row: "index_t", + valid_col: "index_t", + ) -> None: + src_view = _make_tensor(src, shape=[8, 64], dtype=pto.float32) + idx_view = _make_tensor(idx, shape=[8, 64], dtype=uint32()) + dst_view = _make_tensor(dst, shape=[8, 64], dtype=pto.float32) + with pto.vector_section(): + a5.tgather( + _slice_tensor( + src_view, offsets=[0, 0], sizes=[8, 64], dtype=pto.float32 + ), + _slice_tensor(idx_view, offsets=[0, 0], sizes=[8, 64], dtype=uint32()), + _slice_tensor( + dst_view, offsets=[0, 0], sizes=[8, 64], dtype=pto.float32 + ), + dtype=pto.float32, + index_dtype=uint32(), + tile_shape=[8, 64], + valid_row=valid_row, + valid_col=valid_col, + ) + + text = str(a5_tgather_dynamic_valid) + + assert "!pto.tile_buf" in text + assert "valid_row = %arg3 valid_col = %arg4" in text + assert "pto.vgather2" in text + assert "arith.index_cast" in text + + def test_a5_trow_expand_emits_broadcast_micro_ops(): def meta_data(): return { @@ -264,6 +409,179 @@ def a5_trsqrt(src: "ptr_t", dst: "ptr_t") -> None: assert "pto.trsqrt" not in text +@pytest.mark.parametrize( + ("helper_name", "dtype", "micro_op", "tile_op"), + [ + ("tmax", "float32", "pto.vmax", "pto.tmax"), + ("tmin", "float32", "pto.vmin", "pto.tmin"), + ("tand", "int32", "pto.vand", "pto.tand"), + ("txor", "int32", "pto.vxor", "pto.txor"), + ("tshl", "int32", "pto.vshl", "pto.tshl"), + ("tshr", "int32", "pto.vshr", "pto.tshr"), + ], +) +def test_a5_binary_header_helpers_emit_micro_opcodes( + helper_name, dtype, micro_op, tile_op +): + def meta_data(): + mlir_dtype = getattr(pto, dtype) + return {"ptr_t": pto.PtrType(mlir_dtype)} + + helper = getattr(a5, helper_name) + + @to_ir_module(meta_data=meta_data) + def a5_binary_helper(src0: "ptr_t", src1: "ptr_t", dst: "ptr_t") -> None: + mlir_dtype = getattr(pto, dtype) + lhs = _make_tensor(src0, shape=[1, 64], dtype=mlir_dtype) + rhs = _make_tensor(src1, shape=[1, 64], dtype=mlir_dtype) + out = _make_tensor(dst, shape=[1, 64], dtype=mlir_dtype) + with pto.vector_section(): + helper( + _slice_tensor(lhs, offsets=[0, 0], sizes=[1, 64], dtype=mlir_dtype), + _slice_tensor(rhs, offsets=[0, 0], sizes=[1, 64], dtype=mlir_dtype), + _slice_tensor(out, offsets=[0, 0], sizes=[1, 64], dtype=mlir_dtype), + dtype=mlir_dtype, + shape=[1, 64], + ) + + text = str(a5_binary_helper) + + assert micro_op in text + assert tile_op not in text + + +@pytest.mark.parametrize( + ("helper_name", "dtype", "scalar_value", "micro_op", "tile_op"), + [ + ("tadds", "float32", 2.0, "pto.vadds", "pto.tadds"), + ("tsubs", "float32", 2.0, "pto.vsub", "pto.tsubs"), + ("tmaxs", "float32", 2.0, "pto.vmaxs", "pto.tmaxs"), + ("tands", "int32", 7, "pto.vand", "pto.tands"), + ("tshls", "int32", 1, "pto.vshls", "pto.tshls"), + ("tlrelu", "float32", 0.1, "pto.vlrelu", "pto.tlrelu"), + ], +) +def test_a5_scalar_header_helpers_emit_micro_opcodes( + helper_name, dtype, scalar_value, micro_op, tile_op +): + def meta_data(): + mlir_dtype = getattr(pto, dtype) + return {"ptr_t": pto.PtrType(mlir_dtype)} + + helper = getattr(a5, helper_name) + + @to_ir_module(meta_data=meta_data) + def a5_scalar_helper(src: "ptr_t", dst: "ptr_t") -> None: + mlir_dtype = getattr(pto, dtype) + src_view = _make_tensor(src, shape=[1, 64], dtype=mlir_dtype) + dst_view = _make_tensor(dst, shape=[1, 64], dtype=mlir_dtype) + scalar = arith.ConstantOp(mlir_dtype, scalar_value).result + with pto.vector_section(): + helper( + _slice_tensor( + src_view, offsets=[0, 0], sizes=[1, 64], dtype=mlir_dtype + ), + scalar, + _slice_tensor( + dst_view, offsets=[0, 0], sizes=[1, 64], dtype=mlir_dtype + ), + dtype=mlir_dtype, + shape=[1, 64], + ) + + text = str(a5_scalar_helper) + + assert micro_op in text + assert tile_op not in text + + +def test_a5_taxpy_emits_vmula(): + def meta_data(): + return {"ptr_t": pto.PtrType(pto.float32)} + + @to_ir_module(meta_data=meta_data) + def a5_taxpy(src: "ptr_t", dst: "ptr_t") -> None: + src_view = _make_tensor(src, shape=[8, 64], dtype=pto.float32) + dst_view = _make_tensor(dst, shape=[8, 64], dtype=pto.float32) + scalar = arith.ConstantOp(pto.float32, 0.5).result + with pto.vector_section(): + a5.taxpy( + _slice_tensor( + src_view, offsets=[0, 0], sizes=[8, 64], dtype=pto.float32 + ), + scalar, + _slice_tensor( + dst_view, offsets=[0, 0], sizes=[8, 64], dtype=pto.float32 + ), + dtype=pto.float32, + shape=[8, 64], + ) + + text = str(a5_taxpy) + + assert "pto.vmula" in text + assert "pto.taxpy" not in text + + +@pytest.mark.parametrize( + ("helper_name", "micro_op", "tile_op"), + [ + ("trow_expand_add", "pto.vadd", "pto.trowexpandadd"), + ("trow_expand_max", "pto.vmax", "pto.trowexpandmax"), + ("tcol_expand_add", "pto.vadd", "pto.tcolexpandadd"), + ("tcol_expand_max", "pto.vmax", "pto.tcolexpandmax"), + ], +) +def test_a5_expand_header_helpers_emit_micro_opcodes(helper_name, micro_op, tile_op): + def meta_data(): + return {"ptr_t": pto.PtrType(pto.float32)} + + helper = getattr(a5, helper_name) + + @to_ir_module(meta_data=meta_data) + def a5_expand_helper(src0: "ptr_t", src1: "ptr_t", dst: "ptr_t") -> None: + with pto.vector_section(): + if helper_name.startswith("trow_"): + base_view = _make_tensor(src0, shape=[8, 64], dtype=pto.float32) + expand_view = _make_tensor(src1, shape=[8, 1], dtype=pto.float32) + out_view = _make_tensor(dst, shape=[8, 64], dtype=pto.float32) + helper( + _slice_tensor( + base_view, offsets=[0, 0], sizes=[8, 64], dtype=pto.float32 + ), + _slice_tensor( + expand_view, offsets=[0, 0], sizes=[8, 1], dtype=pto.float32 + ), + _slice_tensor( + out_view, offsets=[0, 0], sizes=[8, 64], dtype=pto.float32 + ), + dtype=pto.float32, + shape=[8, 64], + ) + else: + base_view = _make_tensor(src0, shape=[8, 64], dtype=pto.float32) + expand_view = _make_tensor(src1, shape=[1, 64], dtype=pto.float32) + out_view = _make_tensor(dst, shape=[8, 64], dtype=pto.float32) + helper( + _slice_tensor( + base_view, offsets=[0, 0], sizes=[8, 64], dtype=pto.float32 + ), + _slice_tensor( + expand_view, offsets=[0, 0], sizes=[1, 64], dtype=pto.float32 + ), + _slice_tensor( + out_view, offsets=[0, 0], sizes=[8, 64], dtype=pto.float32 + ), + dtype=pto.float32, + shape=[8, 64], + ) + + text = str(a5_expand_helper) + + assert micro_op in text + assert tile_op not in text + + @pytest.mark.parametrize( ("helper_name", "reduce_op", "combine_op", "tile_op"), [ @@ -384,6 +702,51 @@ def a5_tsort32(src: "ptr_src", idx: "ptr_idx", dst: "ptr_src") -> None: assert "pto.tsort32" not in text +def test_a5_tsort32_rejects_dynamic_valid_input_tiles(): + def uint32(): + return IntegerType.get_unsigned(32) + + def meta_data(): + return { + "ptr_src": pto.PtrType(pto.float32), + "ptr_idx": pto.PtrType(uint32()), + "index_t": IndexType.get(), + } + + with pytest.raises( + ValueError, + match="TSORT32 micro lowering currently requires a fully valid input tile", + ): + + @to_ir_module(meta_data=meta_data) + def invalid_tsort32( + src: "ptr_src", + idx: "ptr_idx", + dst: "ptr_src", + valid_row: "index_t", + valid_col: "index_t", + ) -> None: + src_view = _make_tensor(src, shape=[1, 64], dtype=pto.float32) + idx_view = _make_tensor(idx, shape=[1, 64], dtype=uint32()) + dst_view = _make_tensor(dst, shape=[1, 128], dtype=pto.float32) + with pto.vector_section(): + a5.tsort32( + _slice_tensor( + src_view, offsets=[0, 0], sizes=[1, 64], dtype=pto.float32 + ), + _slice_tensor( + idx_view, offsets=[0, 0], sizes=[1, 64], dtype=uint32() + ), + _slice_tensor( + dst_view, offsets=[0, 0], sizes=[1, 128], dtype=pto.float32 + ), + dtype=pto.float32, + tile_shape=[1, 64], + valid_row=valid_row, + valid_col=valid_col, + ) + + def test_a5_tmrgsort_emits_vmrgsort4(): def meta_data(): return {"ptr_t": pto.PtrType(pto.float32)} @@ -419,6 +782,7 @@ def test_a5_generation_script_emits_pto_files(tmp_path): assert generated_names == [ "a5_cube_matmul.pto", "a5_elementwise_add.pto", + "a5_hivm_vadd_demo.pto", "a5_vector_copy.pto", ] @@ -427,6 +791,34 @@ def test_a5_generation_script_emits_pto_files(tmp_path): assert "func.func @" in text +def test_a5_generation_script_can_emit_hivm_llvm_for_micro_kernel(tmp_path): + generated = emit_kernels( + output_dir=tmp_path, + ptoas_bin=_PTOAS_BIN, + emit_hivm_llvm=True, + kernel_names={"a5_hivm_vadd_demo"}, + ) + + assert [path.name for path in generated] == ["a5_hivm_vadd_demo.pto"] + + llvm_path = tmp_path / "a5_hivm_vadd_demo.ll" + assert llvm_path.exists() + + text = llvm_path.read_text(encoding="utf-8") + assert "llvm.hivm.vldsx1.v64f32" in text + assert "llvm.hivm.vadd.v64f32.x" in text + assert "llvm.hivm.vstsx1.v64f32" in text + + +def test_a5_generation_script_only_emits_hivm_sidecars_for_supported_kernels(tmp_path): + emit_kernels(output_dir=tmp_path, ptoas_bin=_PTOAS_BIN, emit_hivm_llvm=True) + + assert (tmp_path / "a5_hivm_vadd_demo.ll").exists() + assert (tmp_path / "a5_vector_copy.ll").exists() + assert not (tmp_path / "a5_elementwise_add.ll").exists() + assert not (tmp_path / "a5_cube_matmul.ll").exists() + + def test_a5_tadd_rejects_view_dtype_mismatch(): def meta_data(): return {"ptr_t": pto.PtrType(pto.float16)}