Repository navigation
Reviewing MLIRBuilder for Impero->MLIR - #5552
SamSJackson wants to merge 22 commits into
Conversation
- Requires parsing GEM and Impero types - Function arguments can be known in advance if we parse Impero AST - Need to determine if Parallel indices are determined at GEM->Impero or Impero->MLIR pass. - GEM->Impero happens in FIAT, no changes from `fiat/connorjward/mlir` branch (but there are some changes, such as __str__ definitions that I have locally)
- Some major bugs to solve: - Evaluate does not work for tensor expressions (literally everything from GEM) - Temporary buffer memory is allocated on alloc/alloca - it should be alloc for runtime & non-read-only vs global for compile-time readonly - gem.Variable improperly assumed to be function arguments. Function arguments should be parsed ahead of time. - Need way to determine between ForOp and ParallelOp. Ideally comes from GEM -> Impero parsing step.
- Evaluate expressions return to temporary variables - Need to determine a way to solve their mapping
Regretfully, the MLIR is invalid Bugs to fix: - Not appropriately handling multiindex/VariableIndex - imp.Return is not aligning the processed values with the expected memref store type (i.e. actual: memref.store val, buf[indices] : f64 expected: memref.store val, buf[indices] : memref<Nxf64> for some value N or dynamic index )
Getting segfaults on execution. Debugging in process. cc.py is still goofy with implementation but just prototype for now.
Not flattening indices correctly
Removed segfault and traversing tree in reverse. Only issue now is that the final scatter is not accumulating
Assemble now works for one-form. Barring that the results are not correct. Implementing ReturnAccumulate would fix this but upstream changes will make that pointless. (ReturnAccumulate being replaced with Assignment)
Getting rid of overlapping functions Pretty much finished now for CPU implementations. Waiting for upstream changes to validate. Possible improvements: - Pass kernel buffers to GemExecutable from GemCodegenContext: - Allows us to create FuncOp from getgo and remove Block/detaching and reattaching ops. - Also avoids sorting of arguments
Added a mapping for buffer->device based on desired location, given generation - A device is placed on CPU if it is used for runtime-variable domain extent Need to make temporary arrays use memref.Global as opposed to memref.Alloc because it enables easier shared memory workgroups on GPU and we can easily initialise to 0 on CPU and GPU
- GlobalOps are used instead of `alloc` because not re-initialised per thread. Memory scope can be assigned (per-thread, per-threadblock, global memory) - Outermost function developed for entering outermost parallel (when we enter cell loop) - Additional cleaning and refactoring
|
|
||
| dt = np.dtype(dtype) | ||
| if dt.kind == "f": | ||
| return f64 if dt.itemsize == 8 else f32 |
There was a problem hiding this comment.
Single precision etc? I think we should be explicit about catching everything and raising errors if we don't know what to do
| self.process(root) | ||
| n = len(self._arguments) | ||
|
|
||
| perm = sorted(range(n), key=lambda i: self._arguments[i].name) |
There was a problem hiding this comment.
I am happy to have an arbitrary argument order in the generated code. In which case, can this go?
There was a problem hiding this comment.
The order in which I was collecting the arguments was not the same as the kernel_names collected during GemCodegenContext. Sorting was aligning them.
I think I would prefer to pass the kernel arguments from GemCodegenContext - it would avoid all of the reordering and attaching, so I shall make that change.
| if len(multiindex) > 1: | ||
| multiindex = multiindex[1:] | ||
|
|
||
| return self._load(aggregate, multiindex, [1]) |
There was a problem hiding this comment.
As your comments suggest, this does look quite ugly and will need cleanup
| _process.register(gem.FloorDiv)(_binary_handler(None, arith.FloorDivSIOp)) | ||
| _process.register(gem.Remainder)(_binary_handler(None, arith.RemSIOp)) | ||
| _process.register(gem.LogicalAnd)(_binary_handler(None, arith.AndIOp)) | ||
| _process.register(gem.LogicalOr)(_binary_handler(None, arith.OrIOp)) |
There was a problem hiding this comment.
_handler_mapping = {
gem.Sum: (AddfOp, AddiOp),
...
}
def _process(op: gem.BinaryOp):
val1, val2 = _handler_mapping[type(op)]
return _binary_handler(val1, val2)(assuming something like BinaryOp exists, it should)
also fix the bad variable names
There was a problem hiding this comment.
For reasons, unknown to me, GEM has some operations as Operators (such as Sum, Product) but others as Scalar (such as MinValue and MaxValue) - Operator is a subclass of Scalar.
Hence, the most suitable choice (for an existing node referring to gem.BinaryOp) is gem.Scalar.
Any other nodes subclassing gem.Scalar are processed with more specificity (such as gem.Index).
Anyway, for now, I have singledispatch registered under gem.Scalar but slightly misleading.
I am just cautious of changing the GEM inheritance system.
| print("Not removing cached libraries") | ||
|
|
||
|
|
||
| @contextmanager |
There was a problem hiding this comment.
I already have this in pyop3.debug.timer!
There was a problem hiding this comment.
Removed it anyway as I was just using this for testing earlier on.
| os.close(descriptor) | ||
|
|
||
| if not compiler.ld: | ||
| if extension == "mlir": |
There was a problem hiding this comment.
I think this needs a refactor, but I can do that because I know the code better
There was a problem hiding this comment.
Yeah, I have avoided properly doing this for a bit. I have not figured out a clean refactor. Happy to let you do this.
Processing IndexSum now. This required pre-walking an imp.For objects to seek any potential iterator arguments. This is because MLIR is an SSA compiler. Summing scalars requires iterator arguments over an scf loop. The iterator arguments must be known in advance of generation
- Adding demo case for non-constant one-form example - Taking function arguments into function generation - This avoids arg sorting and reattaching MLIR operations - Only the dtypes of the function arguments are used - Refactoring the processing of GEM binary operations
MLIRBuilder - Class to build MLIR from an Impero AST
Principal idea is to receive an Impero AST and eagerly generate MLIR.
Most translation happens within
def process. If an object is already processed, the previously-generated SSA value will be used as opposed to re-processing. Additional functions are helper functions for a_processmethod.We use memref.Globals for temporary buffers as opposed to memref.alloc (heap memory) or memref.alloca (stack memory) because it generalises to GPU easier and we assign memory hierarchy on GPU (in per-thread, shared, or global memory). This assumes that the temporary buffers have static sizing.
There is also a constant-cache - to avoid a bloated MLIR, even if it is cleaned in the compiler passes.
Core assumptions being made
To improve
Ongoing bugs