mirrors/llvm-project
0
0
Fork 0
mirror of https://github.com/llvm/llvm-project.git synced 2025-04-26 10:06:05 +00:00
Code Issues Projects Releases Wiki Activity
llvm-project/mlir/test/python/dialects/affine.py
Amy Wang d50fbe43c9
[MLIR][Python] Python binding support for AffineIfOp (#108323) 
Fix the AffineIfOp's default builder such that it takes in an
IntegerSetAttr. AffineIfOp has skipDefaultBuilders=1 which effectively
skips the creation of the default AffineIfOp::builder on the C++ side.
(AffineIfOp has two custom OpBuilder defined in the
extraClassDeclaration.) However, on the python side, _affine_ops_gen.py
shows that the default builder is being created, but it does not accept
IntegerSet and thus is useless. This fix at line 411 makes the default
python AffineIfOp builder take in an IntegerSet input and does not
impact the C++ side of things.
2024-11-13 16:27:46 -05:00

336 lines
13 KiB
Python
Raw Blame History

# RUN: %PYTHON %s | FileCheck %s
from mlir.ir import *
from mlir.dialects import func
from mlir.dialects import arith
from mlir.dialects import memref
from mlir.dialects import affine
import mlir.extras.types as T
def constructAndPrintInModule(f):
print("\nTEST:", f.__name__)
with Context(), Location.unknown():
module = Module.create()
with InsertionPoint(module.body):
f()
print(module)
return f
# CHECK-LABEL: TEST: testAffineStoreOp
@constructAndPrintInModule
def testAffineStoreOp():
f32 = F32Type.get()
index_type = IndexType.get()
memref_type_out = MemRefType.get([12, 12], f32)
# CHECK: func.func @affine_store_test(%[[ARG0:.*]]: index) -> memref<12x12xf32> {
@func.FuncOp.from_py_func(index_type)
def affine_store_test(arg0):
# CHECK: %[[O_VAR:.*]] = memref.alloc() : memref<12x12xf32>
mem = memref.AllocOp(memref_type_out, [], []).result
d0 = AffineDimExpr.get(0)
s0 = AffineSymbolExpr.get(0)
map = AffineMap.get(1, 1, [s0 * 3, d0 + s0 + 1])
# CHECK: %[[A1:.*]] = arith.constant 2.100000e+00 : f32
a1 = arith.ConstantOp(f32, 2.1)
# CHECK: affine.store %[[A1]], %alloc[symbol(%[[ARG0]]) * 3, %[[ARG0]] + symbol(%[[ARG0]]) + 1] : memref<12x12xf32>
affine.AffineStoreOp(a1, mem, indices=[arg0, arg0], map=map)
return mem
# CHECK-LABEL: TEST: testAffineDelinearizeInfer
@constructAndPrintInModule
def testAffineDelinearizeInfer():
# CHECK: %[[C1:.*]] = arith.constant 1 : index
c1 = arith.ConstantOp(T.index(), 1)
# CHECK: %{{.*}}:2 = affine.delinearize_index %[[C1:.*]] into (2, 3) : index, index
two_indices = affine.AffineDelinearizeIndexOp(c1, [], [2, 3])
# CHECK-LABEL: TEST: testAffineLoadOp
@constructAndPrintInModule
def testAffineLoadOp():
f32 = F32Type.get()
index_type = IndexType.get()
memref_type_in = MemRefType.get([10, 10], f32)
# CHECK: func.func @affine_load_test(%[[I_VAR:.*]]: memref<10x10xf32>, %[[ARG0:.*]]: index) -> f32 {
@func.FuncOp.from_py_func(memref_type_in, index_type)
def affine_load_test(I, arg0):
d0 = AffineDimExpr.get(0)
s0 = AffineSymbolExpr.get(0)
map = AffineMap.get(1, 1, [s0 * 3, d0 + s0 + 1])
# CHECK: {{.*}} = affine.load %[[I_VAR]][symbol(%[[ARG0]]) * 3, %[[ARG0]] + symbol(%[[ARG0]]) + 1] : memref<10x10xf32>
a1 = affine.AffineLoadOp(f32, I, indices=[arg0, arg0], map=map)
return a1
# CHECK-LABEL: TEST: testAffineForOp
@constructAndPrintInModule
def testAffineForOp():
f32 = F32Type.get()
index_type = IndexType.get()
memref_type = MemRefType.get([1024], f32)
# CHECK: #[[MAP0:.*]] = affine_map<(d0)[s0] -> (0, d0 + s0)>
# CHECK: #[[MAP1:.*]] = affine_map<(d0, d1) -> (d0 - 2, d1 * 32)>
# CHECK: func.func @affine_for_op_test(%[[BUFFER:.*]]: memref<1024xf32>) {
@func.FuncOp.from_py_func(memref_type)
def affine_for_op_test(buffer):
# CHECK: %[[C1:.*]] = arith.constant 1 : index
c1 = arith.ConstantOp(index_type, 1)
# CHECK: %[[C2:.*]] = arith.constant 2 : index
c2 = arith.ConstantOp(index_type, 2)
# CHECK: %[[C3:.*]] = arith.constant 3 : index
c3 = arith.ConstantOp(index_type, 3)
# CHECK: %[[C9:.*]] = arith.constant 9 : index
c9 = arith.ConstantOp(index_type, 9)
# CHECK: %[[AC0:.*]] = arith.constant 0.000000e+00 : f32
ac0 = AffineConstantExpr.get(0)
d0 = AffineDimExpr.get(0)
d1 = AffineDimExpr.get(1)
s0 = AffineSymbolExpr.get(0)
lb = AffineMap.get(1, 1, [ac0, d0 + s0])
ub = AffineMap.get(2, 0, [d0 - 2, 32 * d1])
sum_0 = arith.ConstantOp(f32, 0.0)
# CHECK: %0 = affine.for %[[INDVAR:.*]] = max #[[MAP0]](%[[C2]])[%[[C3]]] to min #[[MAP1]](%[[C9]], %[[C1]]) step 2 iter_args(%[[SUM0:.*]] = %[[AC0]]) -> (f32) {
sum = affine.AffineForOp(
lb,
ub,
2,
iter_args=[sum_0],
lower_bound_operands=[c2, c3],
upper_bound_operands=[c9, c1],
)
with InsertionPoint(sum.body):
# CHECK: %[[TMP:.*]] = memref.load %[[BUFFER]][%[[INDVAR]]] : memref<1024xf32>
tmp = memref.LoadOp(buffer, [sum.induction_variable])
sum_next = arith.AddFOp(sum.inner_iter_args[0], tmp)
affine.AffineYieldOp([sum_next])
# CHECK-LABEL: TEST: testAffineForOpErrors
@constructAndPrintInModule
def testAffineForOpErrors():
c1 = arith.ConstantOp(T.index(), 1)
c2 = arith.ConstantOp(T.index(), 2)
c3 = arith.ConstantOp(T.index(), 3)
d0 = AffineDimExpr.get(0)
try:
affine.AffineForOp(
c1,
c2,
1,
lower_bound_operands=[c3],
upper_bound_operands=[],
)
except ValueError as e:
assert (
e.args[0]
== "Either a concrete lower bound or an AffineMap in combination with lower bound operands, but not both, is supported."
)
try:
affine.AffineForOp(
AffineMap.get_constant(1),
c2,
1,
lower_bound_operands=[c3, c3],
upper_bound_operands=[],
)
except ValueError as e:
assert (
e.args[0]
== "Wrong number of lower bound operands passed to AffineForOp; Expected 0, got 2."
)
try:
two_indices = affine.AffineDelinearizeIndexOp(c1, [], [1, 1])
affine.AffineForOp(
two_indices,
c2,
1,
lower_bound_operands=[],
upper_bound_operands=[],
)
except ValueError as e:
assert e.args[0] == "Only a single concrete value is supported for lower bound."
try:
affine.AffineForOp(
1.0,
c2,
1,
lower_bound_operands=[],
upper_bound_operands=[],
)
except ValueError as e:
assert e.args[0] == "lower bound must be int | ResultValueT | AffineMap."
@constructAndPrintInModule
def testForSugar():
memref_t = T.memref(10, T.index())
range = affine.for_
# CHECK: #[[$ATTR_2:.+]] = affine_map<(d0) -> (d0)>
# CHECK-LABEL: func.func @range_loop_1(
# CHECK-SAME: %[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: memref<10xindex>) {
# CHECK: affine.for %[[VAL_3:.*]] = #[[$ATTR_2]](%[[VAL_0]]) to #[[$ATTR_2]](%[[VAL_1]]) {
# CHECK: %[[VAL_4:.*]] = arith.addi %[[VAL_3]], %[[VAL_3]] : index
# CHECK: memref.store %[[VAL_4]], %[[VAL_2]]{{\[}}%[[VAL_3]]] : memref<10xindex>
# CHECK: }
# CHECK: return
# CHECK: }
@func.FuncOp.from_py_func(T.index(), T.index(), memref_t)
def range_loop_1(lb, ub, memref_v):
for i in range(lb, ub, step=1):
add = arith.addi(i, i)
memref.store(add, memref_v, [i])
affine.yield_([])
# CHECK-LABEL: func.func @range_loop_2(
# CHECK-SAME: %[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: memref<10xindex>) {
# CHECK: affine.for %[[VAL_3:.*]] = #[[$ATTR_2]](%[[VAL_0]]) to 10 {
# CHECK: %[[VAL_4:.*]] = arith.addi %[[VAL_3]], %[[VAL_3]] : index
# CHECK: memref.store %[[VAL_4]], %[[VAL_2]]{{\[}}%[[VAL_3]]] : memref<10xindex>
# CHECK: }
# CHECK: return
# CHECK: }
@func.FuncOp.from_py_func(T.index(), T.index(), memref_t)
def range_loop_2(lb, ub, memref_v):
for i in range(lb, 10, step=1):
add = arith.addi(i, i)
memref.store(add, memref_v, [i])
affine.yield_([])
# CHECK-LABEL: func.func @range_loop_3(
# CHECK-SAME: %[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: memref<10xindex>) {
# CHECK: affine.for %[[VAL_3:.*]] = 0 to #[[$ATTR_2]](%[[VAL_1]]) {
# CHECK: %[[VAL_4:.*]] = arith.addi %[[VAL_3]], %[[VAL_3]] : index
# CHECK: memref.store %[[VAL_4]], %[[VAL_2]]{{\[}}%[[VAL_3]]] : memref<10xindex>
# CHECK: }
# CHECK: return
# CHECK: }
@func.FuncOp.from_py_func(T.index(), T.index(), memref_t)
def range_loop_3(lb, ub, memref_v):
for i in range(0, ub, step=1):
add = arith.addi(i, i)
memref.store(add, memref_v, [i])
affine.yield_([])
# CHECK-LABEL: func.func @range_loop_4(
# CHECK-SAME: %[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: memref<10xindex>) {
# CHECK: affine.for %[[VAL_3:.*]] = 0 to 10 {
# CHECK: %[[VAL_4:.*]] = arith.addi %[[VAL_3]], %[[VAL_3]] : index
# CHECK: memref.store %[[VAL_4]], %[[VAL_2]]{{\[}}%[[VAL_3]]] : memref<10xindex>
# CHECK: }
# CHECK: return
# CHECK: }
@func.FuncOp.from_py_func(T.index(), T.index(), memref_t)
def range_loop_4(lb, ub, memref_v):
for i in range(0, 10, step=1):
add = arith.addi(i, i)
memref.store(add, memref_v, [i])
affine.yield_([])
# CHECK-LABEL: func.func @range_loop_8(
# CHECK-SAME: %[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: memref<10xindex>) {
# CHECK: %[[VAL_3:.*]] = affine.for %[[VAL_4:.*]] = 0 to 10 iter_args(%[[VAL_5:.*]] = %[[VAL_2]]) -> (memref<10xindex>) {
# CHECK: %[[VAL_6:.*]] = arith.addi %[[VAL_4]], %[[VAL_4]] : index
# CHECK: memref.store %[[VAL_6]], %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<10xindex>
# CHECK: affine.yield %[[VAL_5]] : memref<10xindex>
# CHECK: }
# CHECK: return
# CHECK: }
@func.FuncOp.from_py_func(T.index(), T.index(), memref_t)
def range_loop_8(lb, ub, memref_v):
for i, it in range(0, 10, iter_args=[memref_v]):
add = arith.addi(i, i)
memref.store(add, it, [i])
affine.yield_([it])
# CHECK-LABEL: TEST: testAffineIfWithoutElse
@constructAndPrintInModule
def testAffineIfWithoutElse():
index = IndexType.get()
i32 = IntegerType.get_signless(32)
d0 = AffineDimExpr.get(0)
# CHECK: #[[$SET0:.*]] = affine_set<(d0) : (d0 - 5 >= 0)>
cond = IntegerSet.get(1, 0, [d0 - 5], [False])
# CHECK-LABEL: func.func @simple_affine_if(
# CHECK-SAME: %[[VAL_0:.*]]: index) {
# CHECK: affine.if #[[$SET0]](%[[VAL_0]]) {
# CHECK: %[[VAL_1:.*]] = arith.constant 1 : i32
# CHECK: %[[VAL_2:.*]] = arith.addi %[[VAL_1]], %[[VAL_1]] : i32
# CHECK: }
# CHECK: return
# CHECK: }
@func.FuncOp.from_py_func(index)
def simple_affine_if(cond_operands):
if_op = affine.AffineIfOp(cond, cond_operands=[cond_operands])
with InsertionPoint(if_op.then_block):
one = arith.ConstantOp(i32, 1)
add = arith.AddIOp(one, one)
affine.AffineYieldOp([])
return
# CHECK-LABEL: TEST: testAffineIfWithElse
@constructAndPrintInModule
def testAffineIfWithElse():
index = IndexType.get()
i32 = IntegerType.get_signless(32)
d0 = AffineDimExpr.get(0)
# CHECK: #[[$SET0:.*]] = affine_set<(d0) : (d0 - 5 >= 0)>
cond = IntegerSet.get(1, 0, [d0 - 5], [False])
# CHECK-LABEL: func.func @simple_affine_if_else(
# CHECK-SAME: %[[VAL_0:.*]]: index) {
# CHECK: %[[VAL_IF:.*]]:2 = affine.if #[[$SET0]](%[[VAL_0]]) -> (i32, i32) {
# CHECK: %[[VAL_XT:.*]] = arith.constant 0 : i32
# CHECK: %[[VAL_YT:.*]] = arith.constant 1 : i32
# CHECK: affine.yield %[[VAL_XT]], %[[VAL_YT]] : i32, i32
# CHECK: } else {
# CHECK: %[[VAL_XF:.*]] = arith.constant 2 : i32
# CHECK: %[[VAL_YF:.*]] = arith.constant 3 : i32
# CHECK: affine.yield %[[VAL_XF]], %[[VAL_YF]] : i32, i32
# CHECK: }
# CHECK: %[[VAL_ADD:.*]] = arith.addi %[[VAL_IF]]#0, %[[VAL_IF]]#1 : i32
# CHECK: return
# CHECK: }
@func.FuncOp.from_py_func(index)
def simple_affine_if_else(cond_operands):
if_op = affine.AffineIfOp(
cond, [i32, i32], cond_operands=[cond_operands], has_else=True
)
with InsertionPoint(if_op.then_block):
x_true = arith.ConstantOp(i32, 0)
y_true = arith.ConstantOp(i32, 1)
affine.AffineYieldOp([x_true, y_true])
with InsertionPoint(if_op.else_block):
x_false = arith.ConstantOp(i32, 2)
y_false = arith.ConstantOp(i32, 3)
affine.AffineYieldOp([x_false, y_false])
add = arith.AddIOp(if_op.results[0], if_op.results[1])
return
Reference in New Issue View Git Blame Copy Permalink