9a2769db80
The patch extends the yaml code generation to support the following new OpDSL constructs: - captures - constants - iteration index accesses - predefined types These changes have been introduced by revision https://reviews.llvm.org/D101364. Differential Revision: https://reviews.llvm.org/D102075
188 lines
8.4 KiB
Python
188 lines
8.4 KiB
Python
# RUN: %PYTHON %s | FileCheck %s
|
|
|
|
from mlir.ir import *
|
|
from mlir.dialects import builtin
|
|
from mlir.dialects import linalg
|
|
from mlir.dialects import std
|
|
|
|
from mlir.dialects.linalg.opdsl.lang import *
|
|
|
|
|
|
@linalg_structured_op
|
|
def matmul_mono(
|
|
A=TensorDef(T, S.M, S.K),
|
|
B=TensorDef(T, S.K, S.N),
|
|
C=TensorDef(T, S.M, S.N, output=True)):
|
|
C[D.m, D.n] += A[D.m, D.k] * B[D.k, D.n]
|
|
|
|
|
|
@linalg_structured_op
|
|
def matmul_poly(
|
|
A=TensorDef(TV.T1, S.M, S.K),
|
|
B=TensorDef(TV.T2, S.K, S.N),
|
|
C=TensorDef(U, S.M, S.N, output=True)):
|
|
C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])
|
|
|
|
|
|
@linalg_structured_op
|
|
def fill_rng(
|
|
O=TensorDef(T, S.M, S.N, output=True),
|
|
min=CaptureDef(F64),
|
|
max=CaptureDef(F64),
|
|
seed=CaptureDef(I32)):
|
|
multiplier = cast(I32, const(1103515245))
|
|
increment = cast(I32, const(12345))
|
|
rand1 = (cast(I32, index(D.m)) + seed) * multiplier + increment
|
|
rand2 = (cast(I32, index(D.n)) + rand1) * multiplier + increment
|
|
inv_range = cast(F64, const(2.3283064e-10))
|
|
offset = cast(F64, const(2147483647))
|
|
scaling = (max - min) * inv_range
|
|
O[D.m, D.n] = cast(T, (offset + cast(F64, rand2)) * scaling + min)
|
|
|
|
|
|
with Context() as ctx, Location.unknown():
|
|
module = Module.create()
|
|
f16 = F16Type.get()
|
|
f32 = F32Type.get()
|
|
f64 = F64Type.get()
|
|
i8 = IntegerType.get_signless(8)
|
|
i16 = IntegerType.get_signless(16)
|
|
i32 = IntegerType.get_signless(32)
|
|
with InsertionPoint(module.body):
|
|
|
|
# Note that these all have the same indexing maps. We verify the first and
|
|
# then do more permutation tests on casting and body generation
|
|
# behavior.
|
|
# CHECK: #[[$MAPA:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>
|
|
# CHECK: #[[$MAPB:.+]] = affine_map<(d0, d1, d2) -> (d2, d1)>
|
|
# CHECK: #[[$MAPC:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
|
|
|
|
# CHECK-LABEL: func @test_matmul_mono
|
|
# CHECK-SAME: %[[A:.+]]: tensor<4x16xf32>
|
|
# CHECK-SAME: %[[B:.+]]: tensor<16x8xf32>
|
|
|
|
# CHECK: %[[INITC:.+]] = linalg.init_tensor [4, 8] : tensor<4x8xf32>
|
|
# CHECK: linalg.generic
|
|
# CHECK-SAME: indexing_maps = [#[[$MAPA]], #[[$MAPB]], #[[$MAPC]]]
|
|
# CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
|
|
# CHECK-SAME: ins(%[[A]], %[[B]]
|
|
# CHECK-SAME: outs(%[[INITC]]
|
|
|
|
@builtin.FuncOp.from_py_func(
|
|
RankedTensorType.get((4, 16), f32), RankedTensorType.get((16, 8), f32))
|
|
def test_matmul_mono(lhs, rhs):
|
|
init_result = linalg.InitTensorOp([4, 8], f32)
|
|
return matmul_mono(lhs, rhs, outs=[init_result.result])
|
|
|
|
# CHECK-LABEL: @test_i8i8i32_matmul
|
|
# CHECK: ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i8, %[[C_ARG:.+]]: i32)
|
|
# CHECK-NEXT: %[[A_CAST:.+]] = sexti %[[A_ARG]] : i8 to i32
|
|
# CHECK-NEXT: %[[B_CAST:.+]] = sexti %[[B_ARG]] : i8 to i32
|
|
# CHECK-NEXT: %[[MUL:.+]] = muli %[[A_CAST]], %[[B_CAST]] : i32
|
|
# CHECK-NEXT: %[[ADD:.+]] = addi %[[C_ARG]], %[[MUL]] : i32
|
|
# CHECK-NEXT: linalg.yield %[[ADD]] : i32
|
|
# CHECK-NEXT: -> tensor<4x8xi32>
|
|
@builtin.FuncOp.from_py_func(
|
|
RankedTensorType.get((4, 16), i8), RankedTensorType.get((16, 8), i8),
|
|
RankedTensorType.get((4, 8), i32))
|
|
def test_i8i8i32_matmul(lhs, rhs, init_result):
|
|
return matmul_poly(lhs, rhs, outs=[init_result])
|
|
|
|
# CHECK-LABEL: @test_i8i16i32_matmul
|
|
# CHECK: ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i16, %[[C_ARG:.+]]: i32)
|
|
# CHECK-NEXT: %[[A_CAST:.+]] = sexti %[[A_ARG]] : i8 to i32
|
|
# CHECK-NEXT: %[[B_CAST:.+]] = sexti %[[B_ARG]] : i16 to i32
|
|
# CHECK-NEXT: %[[MUL:.+]] = muli %[[A_CAST]], %[[B_CAST]] : i32
|
|
# CHECK-NEXT: %[[ADD:.+]] = addi %[[C_ARG]], %[[MUL]] : i32
|
|
# CHECK-NEXT: linalg.yield %[[ADD]] : i32
|
|
# CHECK-NEXT: -> tensor<4x8xi32>
|
|
@builtin.FuncOp.from_py_func(
|
|
RankedTensorType.get((4, 16), i8), RankedTensorType.get((16, 8), i16),
|
|
RankedTensorType.get((4, 8), i32))
|
|
def test_i8i16i32_matmul(lhs, rhs, init_result):
|
|
return matmul_poly(lhs, rhs, outs=[init_result])
|
|
|
|
# CHECK-LABEL: @test_i32i32i16_matmul
|
|
# CHECK: ^{{.*}}(%[[A_ARG:.+]]: i32, %[[B_ARG:.+]]: i32, %[[C_ARG:.+]]: i16)
|
|
# CHECK-NEXT: %[[A_CAST:.+]] = trunci %[[A_ARG]] : i32 to i16
|
|
# CHECK-NEXT: %[[B_CAST:.+]] = trunci %[[B_ARG]] : i32 to i16
|
|
# CHECK-NEXT: %[[MUL:.+]] = muli %[[A_CAST]], %[[B_CAST]] : i16
|
|
# CHECK-NEXT: %[[ADD:.+]] = addi %[[C_ARG]], %[[MUL]] : i16
|
|
# CHECK-NEXT: linalg.yield %[[ADD]] : i16
|
|
# CHECK-NEXT: -> tensor<4x8xi16>
|
|
@builtin.FuncOp.from_py_func(
|
|
RankedTensorType.get((4, 16), i32), RankedTensorType.get((16, 8), i32),
|
|
RankedTensorType.get((4, 8), i16))
|
|
def test_i32i32i16_matmul(lhs, rhs, init_result):
|
|
return matmul_poly(lhs, rhs, outs=[init_result])
|
|
|
|
# CHECK-LABEL: @test_i8i8f32_matmul
|
|
# CHECK: ^{{.*}}(%[[A_ARG:.+]]: i8, %[[B_ARG:.+]]: i8, %[[C_ARG:.+]]: f32)
|
|
# CHECK-NEXT: %[[A_CAST:.+]] = sitofp %[[A_ARG]] : i8 to f32
|
|
# CHECK-NEXT: %[[B_CAST:.+]] = sitofp %[[B_ARG]] : i8 to f32
|
|
# CHECK-NEXT: %[[MUL:.+]] = mulf %[[A_CAST]], %[[B_CAST]] : f32
|
|
# CHECK-NEXT: %[[ADD:.+]] = addf %[[C_ARG]], %[[MUL]] : f32
|
|
# CHECK-NEXT: linalg.yield %[[ADD]] : f32
|
|
# CHECK-NEXT: -> tensor<4x8xf32>
|
|
@builtin.FuncOp.from_py_func(
|
|
RankedTensorType.get((4, 16), i8), RankedTensorType.get((16, 8), i8),
|
|
RankedTensorType.get((4, 8), f32))
|
|
def test_i8i8f32_matmul(lhs, rhs, init_result):
|
|
return matmul_poly(lhs, rhs, outs=[init_result])
|
|
|
|
# CHECK-LABEL: @test_f16f16f32_matmul
|
|
# CHECK: ^{{.*}}(%[[A_ARG:.+]]: f16, %[[B_ARG:.+]]: f16, %[[C_ARG:.+]]: f32)
|
|
# CHECK-NEXT: %[[A_CAST:.+]] = fpext %[[A_ARG]] : f16 to f32
|
|
# CHECK-NEXT: %[[B_CAST:.+]] = fpext %[[B_ARG]] : f16 to f32
|
|
# CHECK-NEXT: %[[MUL:.+]] = mulf %[[A_CAST]], %[[B_CAST]] : f32
|
|
# CHECK-NEXT: %[[ADD:.+]] = addf %[[C_ARG]], %[[MUL]] : f32
|
|
# CHECK-NEXT: linalg.yield %[[ADD]] : f32
|
|
# CHECK-NEXT: -> tensor<4x8xf32>
|
|
@builtin.FuncOp.from_py_func(
|
|
RankedTensorType.get((4, 16), f16), RankedTensorType.get((16, 8), f16),
|
|
RankedTensorType.get((4, 8), f32))
|
|
def test_f16f16f32_matmul(lhs, rhs, init_result):
|
|
return matmul_poly(lhs, rhs, outs=[init_result])
|
|
|
|
# CHECK-LABEL: @test_f64f64f32_matmul
|
|
# CHECK: ^{{.*}}(%[[A_ARG:.+]]: f64, %[[B_ARG:.+]]: f64, %[[C_ARG:.+]]: f32)
|
|
# CHECK-NEXT: %[[A_CAST:.+]] = fptrunc %[[A_ARG]] : f64 to f32
|
|
# CHECK-NEXT: %[[B_CAST:.+]] = fptrunc %[[B_ARG]] : f64 to f32
|
|
# CHECK-NEXT: %[[MUL:.+]] = mulf %[[A_CAST]], %[[B_CAST]] : f32
|
|
# CHECK-NEXT: %[[ADD:.+]] = addf %[[C_ARG]], %[[MUL]] : f32
|
|
# CHECK-NEXT: linalg.yield %[[ADD]] : f32
|
|
# CHECK-NEXT: -> tensor<4x8xf32>
|
|
@builtin.FuncOp.from_py_func(
|
|
RankedTensorType.get((4, 16), f64), RankedTensorType.get((16, 8), f64),
|
|
RankedTensorType.get((4, 8), f32))
|
|
def test_f64f64f32_matmul(lhs, rhs, init_result):
|
|
return matmul_poly(lhs, rhs, outs=[init_result])
|
|
|
|
# CHECK-LABEL: @test_fill_rng
|
|
# CHECK-SAME: %{{.*}} tensor<4x16xi32>, %[[MIN:.+]]: f64, %[[MAX:.+]]: f64, %[[SEED:.+]]: i32
|
|
# CHECK-DAG: %[[IDX0:.+]] = linalg.index 0 : index
|
|
# CHECK-DAG: %[[IDX1:.+]] = linalg.index 1 : index
|
|
# CHECK-DAG: %[[IDX0_CAST:.+]] = index_cast %[[IDX0]] : index to i32
|
|
# CHECK-DAG: %[[IDX1_CAST:.+]] = index_cast %[[IDX1]] : index to i32
|
|
# CHECK-DAG: %[[RND0:.+]] = addi %[[IDX0_CAST]], %[[SEED]] : i32
|
|
# CHECK-DAG: %[[CST0:.+]] = constant 1103515245 : i64
|
|
# CHECK-DAG: %[[CST0_CAST:.+]] = trunci %[[CST0]] : i64 to i32
|
|
# CHECK-DAG: %[[CST1:.+]] = constant 12345 : i64
|
|
# CHECK-DAG: %[[CST1_CAST:.+]] = trunci %[[CST1]] : i64 to i32
|
|
# CHECK-DAG: %[[RND1:.+]] = muli %[[RND0]], %[[CST0_CAST]] : i32
|
|
# CHECK-DAG: %[[RND2:.+]] = addi %[[RND1]], %[[CST1_CAST]] : i32
|
|
# Skip random number computation for the second index.
|
|
# CHECK-DAG: %[[DIFF:.+]] = subf %[[MAX]], %[[MIN]] : f64
|
|
# CHECK-DAG: %[[CST3:.+]] = constant 2.3283063999999999E-10 : f64
|
|
# CHECK-DAG: %[[FACT:.+]] = mulf %[[DIFF]], %[[CST3]] : f64
|
|
# CHECK-DAG: %[[RND4:.+]] = mulf %{{.+}}, %[[FACT]] : f64
|
|
# CHECK-DAG: %[[RND5:.+]] = addf %[[RND4]], %[[MIN]] : f64
|
|
# CHECK-DAG: %{{.*}} = fptosi %[[RND5]] : f64 to i32
|
|
@builtin.FuncOp.from_py_func(
|
|
RankedTensorType.get((4, 16), i32), f64, f64, i32)
|
|
def test_fill_rng(init_result, min, max, seed):
|
|
return fill_rng(outs=[init_result], captures=[min, max, seed])
|
|
|
|
|
|
print(module)
|