TileOPs-Metax/tileops/manifest/elementwise_binary.yaml

758 lines
27 KiB
YAML

# elementwise_binary.yaml -- manifest entries for elementwise binary ops.
#
# Two-input single-output elementwise ops (arithmetic, comparison, logical,
# bitwise). Source of truth for op interfaces in this family. Loaded and
# merged with other family files by tileops.manifest at runtime. See
# docs/design/manifest.md for the full schema.
PreluFwdOp:
ref_api: "torch.nn.functional.prelu"
family: elementwise
status: implemented
signature:
inputs:
input: {dtype: "float16 | bfloat16 | float32"}
weight: {dtype: "same_as(input)"}
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
# PyTorch prelu: weight is either scalar or per-channel along dim 1
# for inputs with ndim >= 2; 1-D inputs accept scalar weight only.
- "weight.ndim == 0 or (weight.ndim == 1 and (weight.shape[0] == 1 or (input.ndim >= 2 and weight.shape[0] == input.shape[1])))"
- "output.shape == input.shape"
workloads:
# PReLU CNN feature map (per-channel weight)
- {input_shape: [16, 256, 56, 56], weight_shape: [256], dtypes: [float16, bfloat16], label: "cnn-feat-per-channel"}
- {input_shape: [16, 512, 28, 28], weight_shape: [512], dtypes: [float16, bfloat16], label: "cnn-feat-per-channel-deep"}
roofline:
vars:
N: "product(input.shape)"
W: "1 if weight.ndim == 0 or weight.shape[0] == 1 else weight.shape[0]"
# FLOPs: prelu(x) = x if x>0 else weight*x. Per roofline.md §1.3,
# compare-and-select(1) + mul(1) = 2 per elem.
flops: "2 * N"
# Read input (N) + read weight (small, ~C) + write output (N)
bytes: "(2 * N + W) * elem_bytes"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
prelu: PreluFwdKernel
op: tileops/ops/elementwise/prelu.py
test: tests/ops/test_activation.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
MaskedFillFwdOp:
ref_api: "torch.Tensor.masked_fill"
# Primary entry: PyTorch's torch.Tensor.masked_fill(mask, value: Tensor)
# where ``value`` is a 0-dim Tensor. The out-of-place
# ``Tensor.masked_fill`` returns a new tensor whose shape is the
# bidirectional broadcast of ``input`` and ``mask`` — either operand
# may be expanded (verified directly: ``torch.zeros((2,1)).masked_fill(
# mask=torch.zeros((2,3), dtype=torch.bool), value=1.0)`` returns shape
# ``(2, 3)``). The in-place ``masked_fill_`` is unidirectional, but
# the manifest models only the out-of-place form. The scalar (Number)
# value form lands as MaskedFillScalarFwdOp below per the
# "No Optional[Tensor]" manifest rule for variant splits.
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "bool | uint8 | int8 | int16 | int32 | int64 | float16 | bfloat16 | float32"}
mask: {dtype: "bool"}
value: {dtype: "same_as(input)"}
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
# Out-of-place masked_fill returns the bidirectional broadcast of
# input and mask; value is 0-dim.
- "value.shape == ()"
- "output.shape == broadcast_shapes(input.shape, mask.shape)"
workloads:
- {input_shape: [4096, 4096], mask_shape: [4096, 4096], value_shape: [], dtypes: [float16, bfloat16, float32], label: "elementwise-16M"}
- {input_shape: [16384, 16384], mask_shape: [16384, 16384], value_shape: [], dtypes: [float16, bfloat16], label: "elementwise-256M"}
roofline:
func: "tileops.perf.formulas.masked_fill_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
masked_fill_tensor_value: MaskedFillTensorValueFwdKernel
op: tileops/ops/elementwise/masked_fill.py
test: tests/ops/test_special_elementwise_conformance.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
MaskedFillScalarFwdOp:
ref_api: "torch.Tensor.masked_fill"
# Scalar-value variant: PyTorch's torch.Tensor.masked_fill(mask, value: Number).
# Per the "No Optional[Tensor]" manifest rule, this is split from the
# 0-dim-Tensor-value primary entry above.
family: elementwise
status: implemented
torch_compile_fullgraph: true
variant_of: MaskedFillFwdOp
signature:
inputs:
input: {dtype: "bool | uint8 | int8 | int16 | int32 | int64 | float16 | bfloat16 | float32"}
mask: {dtype: "bool"}
outputs:
output: {dtype: "same_as(input)"}
params:
value: {type: Number, default: 0.0}
shape_rules:
# Out-of-place masked_fill returns the bidirectional broadcast of
# input and mask; out shape follows that broadcast (verified against
# ``torch.Tensor.masked_fill`` — input may also be expanded up).
- "output.shape == broadcast_shapes(input.shape, mask.shape)"
workloads:
- {input_shape: [4096, 4096], dtypes: [float16, bfloat16, float32], label: "elementwise-16M"}
- {input_shape: [16384, 16384], dtypes: [float16, bfloat16], label: "elementwise-256M"}
roofline:
# Func mode: shared with MaskedFillFwdOp (Tensor-value primary). See
# MaskedFillFwdOp.roofline for rationale (broadcast_shapes not in
# inline-roofline vars-layer namespace).
func: "tileops.perf.formulas.masked_fill_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
masked_fill: MaskedFillFwdKernel
op: tileops/ops/elementwise/masked_fill.py
test: tests/ops/test_special_elementwise.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
AddFwdOp:
ref_api: "torch.add"
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "bool | uint8 | int8 | int16 | int32 | int64 | float16 | bfloat16 | float32"}
other: {dtype: "same_as(input)"}
outputs:
output: {dtype: "same_as(input)"}
params:
alpha: {type: "int | float", default: 1, kw_only: true}
shape_rules:
# Output follows PyTorch broadcasting; numel uses the broadcast shape.
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
- {input_shape: [16, 256, 56, 56], other_shape: [256, 1, 1], dtypes: [float16, bfloat16, float32], label: cnn-feat-broadcast}
roofline:
func: "tileops.perf.formulas.add_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
add: AddFwdKernel
op: tileops/ops/elementwise/arithmetic.py
test: tests/ops/test_binary_arith.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
SubFwdOp:
ref_api: "torch.sub"
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "uint8 | int8 | int16 | int32 | int64 | float16 | bfloat16 | float32"}
other: {dtype: "same_as(input)"}
outputs:
output: {dtype: "same_as(input)"}
params:
alpha: {type: "int | float", default: 1, kw_only: true}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
- {input_shape: [16, 256, 56, 56], other_shape: [256, 1, 1], dtypes: [float16, bfloat16, float32], label: cnn-feat-broadcast}
roofline:
func: "tileops.perf.formulas.sub_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
sub: SubFwdKernel
op: tileops/ops/elementwise/arithmetic.py
test: tests/ops/test_binary_arith.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
MulFwdOp:
ref_api: "torch.mul"
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "bool | uint8 | int8 | int16 | int32 | int64 | float16 | bfloat16 | float32"}
other: {dtype: "same_as(input)"}
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
- {input_shape: [16, 256, 56, 56], other_shape: [256, 1, 1], dtypes: [float16, bfloat16, float32], label: cnn-feat-broadcast}
roofline:
func: "tileops.perf.formulas.mul_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
mul: MulFwdKernel
op: tileops/ops/elementwise/arithmetic.py
test: tests/ops/test_binary_arith.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
DivFwdOp:
ref_api: "torch.div"
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "float16 | bfloat16 | float32"}
other: {dtype: "same_as(input)"}
outputs:
output: {dtype: "same_as(input)"}
params:
rounding_mode: {type: "str | None", default: null, kw_only: true}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
- "rounding_mode is None or rounding_mode in ('trunc', 'floor')"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
- {input_shape: [16, 256, 56, 56], other_shape: [256, 1, 1], dtypes: [float16, bfloat16, float32], label: cnn-feat-broadcast}
roofline:
func: "tileops.perf.formulas.div_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
div: DivFwdKernel
op: tileops/ops/elementwise/arithmetic.py
test: tests/ops/test_binary_arith.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
RemainderFwdOp:
ref_api: "torch.remainder"
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "float16 | bfloat16 | float32"}
other: {dtype: "same_as(input)"}
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
- {input_shape: [16, 256, 56, 56], other_shape: [256, 1, 1], dtypes: [float16, bfloat16, float32], label: cnn-feat-broadcast}
roofline:
func: "tileops.perf.formulas.remainder_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
remainder: RemainderFwdKernel
op: tileops/ops/elementwise/arithmetic.py
test: tests/ops/test_binary_arith.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
PowFwdOp:
ref_api: "torch.pow"
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "float16 | bfloat16 | float32"}
exponent: {dtype: "same_as(input)"}
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, exponent.shape)"
workloads:
- {input_shape: [2048, 4096], exponent_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
- {input_shape: [16, 256, 56, 56], exponent_shape: [256, 1, 1], dtypes: [float16, bfloat16, float32], label: cnn-feat-broadcast}
roofline:
func: "tileops.perf.formulas.pow_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
pow: PowFwdKernel
op: tileops/ops/elementwise/arithmetic.py
test: tests/ops/test_binary_arith.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
FloorDivideFwdOp:
ref_api: "torch.floor_divide"
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "float16 | bfloat16 | float32"}
other: {dtype: "same_as(input)"}
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
- {input_shape: [16, 256, 56, 56], other_shape: [256, 1, 1], dtypes: [float16, bfloat16, float32], label: cnn-feat-broadcast}
roofline:
func: "tileops.perf.formulas.floor_divide_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
floor_divide: FloorDivideFwdKernel
op: tileops/ops/elementwise/arithmetic.py
test: tests/ops/test_binary_arith.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
LerpFwdOp:
ref_api: "torch.lerp"
# Scalar-weight variant: PyTorch's torch.lerp also accepts a Tensor
# weight (see LerpTensorFwdOp); this entry tracks the scalar (Number)
# slice that the existing kernel implements.
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "float16 | bfloat16 | float32"}
end: {dtype: "same_as(input)"}
outputs:
output: {dtype: "same_as(input)"}
params:
weight: {type: float, default: 0.5}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, end.shape)"
workloads:
- {input_shape: [2048, 4096], end_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
- {input_shape: [16, 256, 56, 56], end_shape: [256, 1, 1], dtypes: [float16, bfloat16, float32], label: cnn-feat-broadcast}
roofline:
func: "tileops.perf.formulas.lerp_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
lerp: LerpFwdKernel
op: tileops/ops/elementwise/arithmetic.py
test: tests/ops/test_binary_arith.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
MaximumFwdOp:
ref_api: "torch.maximum"
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "bool | uint8 | int8 | int16 | int32 | int64 | float16 | bfloat16 | float32"}
other: {dtype: "same_as(input)"}
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
- {input_shape: [16, 256, 56, 56], other_shape: [256, 1, 1], dtypes: [float16, bfloat16, float32], label: cnn-feat-broadcast}
roofline:
func: "tileops.perf.formulas.maximum_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
maximum: MaximumFwdKernel
op: tileops/ops/elementwise/arithmetic.py
test: tests/ops/test_binary_arith.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
MinimumFwdOp:
ref_api: "torch.minimum"
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "bool | uint8 | int8 | int16 | int32 | int64 | float16 | bfloat16 | float32"}
other: {dtype: "same_as(input)"}
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
- {input_shape: [16, 256, 56, 56], other_shape: [256, 1, 1], dtypes: [float16, bfloat16, float32], label: cnn-feat-broadcast}
roofline:
func: "tileops.perf.formulas.minimum_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
minimum: MinimumFwdKernel
op: tileops/ops/elementwise/arithmetic.py
test: tests/ops/test_binary_arith.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
EqFwdOp:
ref_api: "torch.eq"
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "bool | uint8 | int8 | int16 | int32 | int64 | float16 | bfloat16 | float32"}
other: {dtype: "same_as(input)"}
outputs:
output: {dtype: "bool"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
- {input_shape: [16, 256, 56, 56], other_shape: [256, 1, 1], dtypes: [float16, bfloat16, float32], label: cnn-feat-broadcast}
roofline:
func: "tileops.perf.formulas.eq_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
eq: EqFwdKernel
op: tileops/ops/elementwise/comparison.py
test: tests/ops/test_comparison.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
NeFwdOp:
ref_api: "torch.ne"
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "bool | uint8 | int8 | int16 | int32 | int64 | float16 | bfloat16 | float32"}
other: {dtype: "same_as(input)"}
outputs:
output: {dtype: "bool"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
- {input_shape: [16, 256, 56, 56], other_shape: [256, 1, 1], dtypes: [float16, bfloat16, float32], label: cnn-feat-broadcast}
roofline:
func: "tileops.perf.formulas.ne_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
ne: NeFwdKernel
op: tileops/ops/elementwise/comparison.py
test: tests/ops/test_comparison.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
GtFwdOp:
ref_api: "torch.gt"
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "bool | uint8 | int8 | int16 | int32 | int64 | float16 | bfloat16 | float32"}
other: {dtype: "same_as(input)"}
outputs:
output: {dtype: "bool"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
- {input_shape: [16, 256, 56, 56], other_shape: [256, 1, 1], dtypes: [float16, bfloat16, float32], label: cnn-feat-broadcast}
roofline:
func: "tileops.perf.formulas.gt_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
gt: GtFwdKernel
op: tileops/ops/elementwise/comparison.py
test: tests/ops/test_comparison.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
LtFwdOp:
ref_api: "torch.lt"
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "bool | uint8 | int8 | int16 | int32 | int64 | float16 | bfloat16 | float32"}
other: {dtype: "same_as(input)"}
outputs:
output: {dtype: "bool"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
- {input_shape: [16, 256, 56, 56], other_shape: [256, 1, 1], dtypes: [float16, bfloat16, float32], label: cnn-feat-broadcast}
roofline:
func: "tileops.perf.formulas.lt_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
lt: LtFwdKernel
op: tileops/ops/elementwise/comparison.py
test: tests/ops/test_comparison.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
GeFwdOp:
ref_api: "torch.ge"
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "bool | uint8 | int8 | int16 | int32 | int64 | float16 | bfloat16 | float32"}
other: {dtype: "same_as(input)"}
outputs:
output: {dtype: "bool"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
- {input_shape: [16, 256, 56, 56], other_shape: [256, 1, 1], dtypes: [float16, bfloat16, float32], label: cnn-feat-broadcast}
roofline:
func: "tileops.perf.formulas.ge_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
ge: GeFwdKernel
op: tileops/ops/elementwise/comparison.py
test: tests/ops/test_comparison.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
LeFwdOp:
ref_api: "torch.le"
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "bool | uint8 | int8 | int16 | int32 | int64 | float16 | bfloat16 | float32"}
other: {dtype: "same_as(input)"}
outputs:
output: {dtype: "bool"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [float16, bfloat16, float32], label: hidden-state-prefill}
- {input_shape: [16, 256, 56, 56], other_shape: [256, 1, 1], dtypes: [float16, bfloat16, float32], label: cnn-feat-broadcast}
roofline:
func: "tileops.perf.formulas.le_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
le: LeFwdKernel
op: tileops/ops/elementwise/comparison.py
test: tests/ops/test_comparison.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
LogicalAndFwdOp:
ref_api: "torch.logical_and"
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "bool | uint8 | int8 | int16 | int32 | int64 | float16 | bfloat16 | float32"}
other: {dtype: "same_as(input)"}
outputs:
output: {dtype: "bool"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [bool, float16, bfloat16, float32], label: hidden-state-prefill}
- {input_shape: [16, 256, 56, 56], other_shape: [256, 1, 1], dtypes: [bool, float16, bfloat16, float32], label: cnn-feat-broadcast}
roofline:
func: "tileops.perf.formulas.logical_and_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
logical_and: LogicalAndFwdKernel
op: tileops/ops/elementwise/logical.py
test: tests/ops/test_logical.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
LogicalOrFwdOp:
ref_api: "torch.logical_or"
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "bool | uint8 | int8 | int16 | int32 | int64 | float16 | bfloat16 | float32"}
other: {dtype: "same_as(input)"}
outputs:
output: {dtype: "bool"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [bool, float16, bfloat16, float32], label: hidden-state-prefill}
- {input_shape: [16, 256, 56, 56], other_shape: [256, 1, 1], dtypes: [bool, float16, bfloat16, float32], label: cnn-feat-broadcast}
roofline:
func: "tileops.perf.formulas.logical_or_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
logical_or: LogicalOrFwdKernel
op: tileops/ops/elementwise/logical.py
test: tests/ops/test_logical.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
BitwiseAndFwdOp:
ref_api: "torch.bitwise_and"
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "bool | uint8 | int8 | int16 | int32 | int64"}
other: {dtype: "same_as(input)"}
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [bool, int32, int64], label: hidden-state-prefill}
- {input_shape: [16, 256, 56, 56], other_shape: [256, 1, 1], dtypes: [bool, int32, int64], label: cnn-feat-broadcast}
roofline:
func: "tileops.perf.formulas.bitwise_and_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
bitwise_and: BitwiseAndFwdKernel
op: tileops/ops/elementwise/bitwise.py
test: tests/ops/test_bitwise.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
BitwiseOrFwdOp:
ref_api: "torch.bitwise_or"
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "bool | uint8 | int8 | int16 | int32 | int64"}
other: {dtype: "same_as(input)"}
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [bool, int32, int64], label: hidden-state-prefill}
- {input_shape: [16, 256, 56, 56], other_shape: [256, 1, 1], dtypes: [bool, int32, int64], label: cnn-feat-broadcast}
roofline:
func: "tileops.perf.formulas.bitwise_or_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
bitwise_or: BitwiseOrFwdKernel
op: tileops/ops/elementwise/bitwise.py
test: tests/ops/test_bitwise.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true
BitwiseXorFwdOp:
ref_api: "torch.bitwise_xor"
family: elementwise
status: implemented
torch_compile_fullgraph: true
signature:
inputs:
input: {dtype: "bool | uint8 | int8 | int16 | int32 | int64"}
other: {dtype: "same_as(input)"}
outputs:
output: {dtype: "same_as(input)"}
shape_rules:
- "output.shape == broadcast_shapes(input.shape, other.shape)"
workloads:
- {input_shape: [2048, 4096], other_shape: [2048, 4096], dtypes: [bool, int32, int64], label: hidden-state-prefill}
- {input_shape: [16, 256, 56, 56], other_shape: [256, 1, 1], dtypes: [bool, int32, int64], label: cnn-feat-broadcast}
roofline:
func: "tileops.perf.formulas.bitwise_xor_fwd_roofline"
source:
kernel: tileops/kernels/elementwise.py
kernel_map:
bitwise_xor: BitwiseXorFwdKernel
op: tileops/ops/elementwise/bitwise.py
test: tests/ops/test_bitwise.py
bench: benchmarks/ops/bench_elementwise_manifest.py
bench_manifest_driven: true