[Mlir-commits] [mlir] [mlir][python] python binding wrapper for the affine.AffineForOp (PR #74408)
Amy Wang
llvmlistbot at llvm.org
Mon Dec 4 19:44:02 PST 2023
https://github.com/kaitingwang created https://github.com/llvm/llvm-project/pull/74408
This PR creates the wrapper class AffineForOp and adds a testcase for it. A testcase for the AffineLoadOp is also added.
>From baf42cdea2f243a1c9d2f89c456589e351c50521 Mon Sep 17 00:00:00 2001
From: Amy Wang <kai.ting.wang at huawei.com>
Date: Mon, 4 Dec 2023 22:31:23 -0500
Subject: [PATCH] [mlir][python] python binding for the affine.for op
This PR creates the wrapper class AffineForOp and adds a testcase
for it. A testcase for AffineLoadOp is also added
---
mlir/python/mlir/dialects/affine.py | 91 +++++++++++++++++++++++++++++
mlir/test/python/dialects/affine.py | 87 +++++++++++++++++++++++++--
2 files changed, 174 insertions(+), 4 deletions(-)
diff --git a/mlir/python/mlir/dialects/affine.py b/mlir/python/mlir/dialects/affine.py
index 80d3873e19a05..71b44e5492716 100644
--- a/mlir/python/mlir/dialects/affine.py
+++ b/mlir/python/mlir/dialects/affine.py
@@ -3,3 +3,94 @@
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
from ._affine_ops_gen import *
+from ._affine_ops_gen import _Dialect, AffineForOp
+
+try:
+ from ..ir import *
+ from ._ods_common import (
+ get_op_result_or_value as _get_op_result_or_value,
+ get_op_results_or_values as _get_op_results_or_values,
+ _cext as _ods_cext,
+ )
+except ImportError as e:
+ raise RuntimeError("Error loading imports from extension module") from e
+
+from typing import Optional, Sequence, Union
+
+
+ at _ods_cext.register_operation(_Dialect, replace=True)
+class AffineForOp(AffineForOp):
+ """Specialization for the Affine for op class"""
+
+ def __init__(
+ self,
+ lower_bound,
+ upper_bound,
+ step,
+ iter_args: Optional[Union[Operation, OpView, Sequence[Value]]] = None,
+ *,
+ lower_bound_operands=[],
+ upper_bound_operands=[],
+ loc=None,
+ ip=None
+ ):
+ """Creates an Affine `for` operation.
+
+ - `lower_bound` is the affine map to use as lower bound of the loop.
+ - `upper_bound` is the affine map to use as upper bound of the loop.
+ - `step` is the value to use as loop step.
+ - `iter_args` is a list of additional loop-carried arguments or an operation
+ producing them as results.
+ - `lower_bound_operands` is the list of arguments to substitute the dimensions,
+ then symbols in the `lower_bound` affine map, in an increasing order
+ - `upper_bound_operands` is the list of arguments to substitute the dimensions,
+ then symbols in the `upper_bound` affine map, in an increasing order
+ """
+
+ if iter_args is None:
+ iter_args = []
+ iter_args = _get_op_results_or_values(iter_args)
+
+ if len(lower_bound_operands) != lower_bound.n_inputs:
+ raise ValueError(
+ f"Wrong number of lower bound operands passed to AffineForOp. " +
+ "Expected {lower_bound.n_symbols}, got {len(lower_bound_operands)}.")
+
+ if len(upper_bound_operands) != upper_bound.n_inputs:
+ raise ValueError(
+ f"Wrong number of upper bound operands passed to AffineForOp. " +
+ "Expected {upper_bound.n_symbols}, got {len(upper_bound_operands)}.")
+
+ results = [arg.type for arg in iter_args]
+ super().__init__(
+ results_=results,
+ lowerBoundOperands=[_get_op_result_or_value(
+ o) for o in lower_bound_operands],
+ upperBoundOperands=[_get_op_result_or_value(
+ o) for o in upper_bound_operands],
+ inits=list(iter_args),
+ lowerBoundMap=AffineMapAttr.get(lower_bound),
+ upperBoundMap=AffineMapAttr.get(upper_bound),
+ step=IntegerAttr.get(IndexType.get(), step),
+ loc=loc,
+ ip=ip
+ )
+ self.regions[0].blocks.append(IndexType.get(), *results)
+
+ @property
+ def body(self):
+ """Returns the body (block) of the loop."""
+ return self.regions[0].blocks[0]
+
+ @property
+ def induction_variable(self):
+ """Returns the induction variable of the loop."""
+ return self.body.arguments[0]
+
+ @property
+ def inner_iter_args(self):
+ """Returns the loop-carried arguments usable within the loop.
+
+ To obtain the loop-carried operands, use `iter_args`.
+ """
+ return self.body.arguments[1:]
diff --git a/mlir/test/python/dialects/affine.py b/mlir/test/python/dialects/affine.py
index c5ec85457493b..c780bf01beac2 100644
--- a/mlir/test/python/dialects/affine.py
+++ b/mlir/test/python/dialects/affine.py
@@ -1,10 +1,10 @@
# RUN: %PYTHON %s | FileCheck %s
from mlir.ir import *
-import mlir.dialects.func as func
-import mlir.dialects.arith as arith
-import mlir.dialects.affine as affine
-import mlir.dialects.memref as memref
+from mlir.dialects import func
+from mlir.dialects import arith
+from mlir.dialects import memref
+from mlir.dialects import affine
def run(f):
@@ -42,3 +42,82 @@ def affine_store_test(arg0):
return mem
print(module)
+
+
+# CHECK-LABEL: TEST: testAffineLoadOp
+ at run
+def testAffineLoadOp():
+ with Context() as ctx, Location.unknown():
+ module = Module.create()
+ with InsertionPoint(module.body):
+ 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
+
+ print(module)
+
+
+# CHECK-LABEL: TEST: testAffineForOp
+ at run
+def testAffineForOp():
+ with Context() as ctx, Location.unknown():
+ module = Module.create()
+ with InsertionPoint(module.body):
+ 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])
+
+ return
+ print(module)
More information about the Mlir-commits
mailing list