[PATCH] D73437: [mlir][spirv] Convert linalg.generic for reduction to SPIR-V ops

[Mahesh Ravishankar via Phabricator via llvm-commits]

mravishankar accepted this revision.
mravishankar marked 2 inline comments as done.
mravishankar added inline comments.


================
Comment at: mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRV.cpp:73
+
+  using mlir::matchers::m_Val;
+  auto a = m_Val(block.getArgument(0));
----------------
mravishankar wrote:
> antiagainst wrote:
> > mravishankar wrote:
> > > I feel like this is not required. If the linalg.generic operation says this is a reduction, we should not be checking the region to verify that it is. linalg as a contract is guaranteeing this is a reduction.
> > I'm not sure I get your idea correctly; but this is checking we are doing the expected kind of reduction (`BinaryOp`). 
> That is what I am not sure we need to do. You have already checked the generic op is a 1D loop with reduction iterator type. So the body of the generic op is assumed to be a reduction right, i.e. a "binaryOp". Here you are checking whether its an operation which takes two operands, but not necessarily a binary operation that can be used for reduction. In some sense you dont need to check that since the generic op description already told you to assume that the region of the op represents a binary op that can be used for reduction (this is based on my understanding of "reduction" loops in linalg, but I need to double check).
> 
Ok, looked into this a little bit more. This is indeed required as of now. Ideally we dont need to do this, but that is orthogonal. No issues here.


================
Comment at: mlir/lib/Conversion/LinalgToSPIRV/LinalgToSPIRV.cpp:182
+  // Perform the group reduction operation.
+  Value groupOperation = rewriter.create<spirv::GroupNonUniformIAddOp>(
+      loc, originalInputType.getElementType(), spirv::Scope::Subgroup,
----------------
mravishankar wrote:
> antiagainst wrote:
> > mravishankar wrote:
> > > This is hard-wiring to IAdd. We probably need to structure this differently. We need to have a pattern to lower the linalg.generic with reduction iterator into the kernel generated here, and then lower the operations within the region separately.
> > It was intentional to only support IAdd. I've re-structured this a bit so it's extensible to other binary op kinds.
> The way I see this structured is
> 
> 1) You check the "structure" of the linalg.generic op to see if it is a 1D reduction. You assume that the body of the generic op represents a binary operation that can be used for the reduction.
> 2) You write a separate pattern that converts the operations of the body itself into the spirv::GroupNonUniform*Op.
> 
> The dialect conversion will first convert the generic op, and then it will attempt to convert the body of the op. The way this is strucutred, it can only handle straight-forward binary operations. THere could be more complicated operations in the region of a generic op that implements the reduction, which would be hard to incorporate in this approach.
> 
> With your updated changes, it is easy to extend to other ops, but I think a more general solution is needed where we are not constrained in handling just simple reductions. It might need some more careful design though, which I dont have a full picture of right now. So for now this OK, but please leave a TODO to say something like "handle more general reductions".
Thought about this a little bit more. It will be a more sane strategy to not handle cases where the region is "inlined" into the generic op, but rather the region is specified as a function using the `fun` attribute of the `linalg.generic` op. Then we can split the lowering of the generic op and the lowering of the function that does the reduction separately.
This too not needed for this CL. Should revisit this


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D73437/new/

https://reviews.llvm.org/D73437